diff --git a/.clang-format b/.clang-format
index a4de8e7be8..4bb0ebed18 100644
--- a/.clang-format
+++ b/.clang-format
@@ -26,4 +26,5 @@ BinPackParameters: false
 BinPackArguments: false
 IncludeBlocks: Preserve
 IncludeIsMainSourceRegex: (\.cu)$
+SortIncludes: false
 ...
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index eeb50d4201..1f5791e398 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,6 +1,7 @@
 exclude: |
   (?x)^(
-    dockerfiles/.+
+    dockerfiles/.+|
+    custom_ops/third_party/.+
   )$
 default_install_hook_types:
   - pre-commit
diff --git a/custom_ops/gpu_ops/append_attn/decode_attention_func.cuh b/custom_ops/gpu_ops/append_attn/decode_attention_func.cuh
index 4ff4a02298..8f7b096e6b 100644
--- a/custom_ops/gpu_ops/append_attn/decode_attention_func.cuh
+++ b/custom_ops/gpu_ops/append_attn/decode_attention_func.cuh
@@ -42,13 +42,10 @@ struct softmax_state_t {
     }
   }
 
-  __device__ __forceinline__ softmax_state_t() {
-    init();
-  }
+  __device__ __forceinline__ softmax_state_t() { init(); }
 
-  __device__ __forceinline__ void merge(const AlignedVector<T, vec_size>& other_o,
-                                        T other_m,
-                                        T other_d) {
+  __device__ __forceinline__ void merge(
+      const AlignedVector<T, vec_size>& other_o, T other_m, T other_d) {
     // using kType = typename cascade_attn_nv_type2_traits<T>::type;
     T m_prev = m, d_prev = d;
     m = m_prev > other_m ? m_prev : other_m;
@@ -63,13 +60,11 @@ struct softmax_state_t {
   }
 
   __device__ __forceinline__ void normalize() {
-
 #pragma unroll
     for (size_t i = 0; i < vec_size; ++i) {
       o[i] /= d;
     }
   }
-
 };
 
 template <size_t vec_size, typename T, uint32_t num_tiles = 0>
@@ -102,65 +97,79 @@ struct softmax_state_ts {
     }
   }
 
-  __device__ __forceinline__ softmax_state_ts() {
-    init();
-  }
+  __device__ __forceinline__ softmax_state_ts() { init(); }
 
   __device__ __forceinline__ void normalize(const uint32_t tile_id) {
-
 #pragma unroll
     for (size_t i = 0; i < vec_size; i++) {
       o[tile_id][i] /= d;
     }
   }
-
 };
 
-template <SharedMemFillMode fill_mode, uint32_t HEAD_DIM_QK, uint32_t vec_size, uint32_t NUM_VEC_PER_HEAD, uint32_t bdx, uint32_t BLOCK_SIZE, uint32_t CACHE_VEC_SIZE, typename CacheT>
-__device__ __forceinline__ void produce_kv(CacheT *smem,
-                                          CacheT *kv_base_gptr,
-                                          const int * block_table_smem,
-                                          const uint32_t seq_offset_gmem,
-                                          const uint32_t seq_offset_smem,
-                                          const uint32_t kv_head_idx,
-                                          const uint32_t kv_num_heads,
-                                          const uint32_t tidx,
-                                          const uint32_t chunk_start,
-                                          const uint32_t chunk_end) {
+template <SharedMemFillMode fill_mode,
+          uint32_t HEAD_DIM_QK,
+          uint32_t vec_size,
+          uint32_t NUM_VEC_PER_HEAD,
+          uint32_t bdx,
+          uint32_t BLOCK_SIZE,
+          uint32_t CACHE_VEC_SIZE,
+          typename CacheT>
+__device__ __forceinline__ void produce_kv(CacheT* smem,
+                                           CacheT* kv_base_gptr,
+                                           const int* block_table_smem,
+                                           const uint32_t seq_offset_gmem,
+                                           const uint32_t seq_offset_smem,
+                                           const uint32_t kv_head_idx,
+                                           const uint32_t kv_num_heads,
+                                           const uint32_t tidx,
+                                           const uint32_t chunk_start,
+                                           const uint32_t chunk_end) {
   int block_id = __ldg(&block_table_smem[seq_offset_gmem / BLOCK_SIZE]);
   if (block_id < 0) {
     block_id = 0;
   }
   const uint32_t block_offset = seq_offset_gmem % BLOCK_SIZE;
   // 8/16 T/int8 each time
-  const uint32_t k_offset_base = ((block_id * kv_num_heads + kv_head_idx) * BLOCK_SIZE + block_offset) * HEAD_DIM_QK;
+  const uint32_t k_offset_base =
+      ((block_id * kv_num_heads + kv_head_idx) * BLOCK_SIZE + block_offset) *
+      HEAD_DIM_QK;
   const uint32_t smem_offset_base = seq_offset_smem * HEAD_DIM_QK;
-  for(uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
+  for (uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
     pred_load<128, PrefetchMode::kPrefetch, fill_mode, CacheT>(
-      smem + smem_offset_base + vid * CACHE_VEC_SIZE,
-      kv_base_gptr + k_offset_base + vid * CACHE_VEC_SIZE,
-      seq_offset_gmem < chunk_end
-    );
+        smem + smem_offset_base + vid * CACHE_VEC_SIZE,
+        kv_base_gptr + k_offset_base + vid * CACHE_VEC_SIZE,
+        seq_offset_gmem < chunk_end);
   }
 }
 
-template <uint32_t vec_size, uint32_t NUM_VEC_PER_HEAD, uint32_t bdx, uint32_t bdy, uint32_t HEAD_DIM, uint32_t DEAL_EACH_TIME, uint32_t num_tile_v, typename T, typename CacheT>
-__device__ __forceinline__ void compute_qk(const T* cu_q_smem,
-                                           const CacheT* k_smem,
-                                           const uint32_t kv_idx_base,
-                                           const uint32_t stage_idx,
-                                           const uint32_t iter_base,
-                                           const uint32_t iter_bound,
-                                           const uint32_t tidx,
-                                           const uint32_t gid,
-                                           const float scale,
-                                           float *s,
-                                           softmax_state_ts<vec_size, T, num_tile_v>& st) {
+template <uint32_t vec_size,
+          uint32_t NUM_VEC_PER_HEAD,
+          uint32_t bdx,
+          uint32_t bdy,
+          uint32_t HEAD_DIM,
+          uint32_t DEAL_EACH_TIME,
+          uint32_t num_tile_v,
+          typename T,
+          typename CacheT>
+__device__ __forceinline__ void compute_qk(
+    const T* cu_q_smem,
+    const CacheT* k_smem,
+    const uint32_t kv_idx_base,
+    const uint32_t stage_idx,
+    const uint32_t iter_base,
+    const uint32_t iter_bound,
+    const uint32_t tidx,
+    const uint32_t gid,
+    const float scale,
+    float* s,
+    softmax_state_ts<vec_size, T, num_tile_v>& st) {
   const CacheT* smem;
   AlignedVector<T, vec_size> q_vec;
   AlignedVector<T, vec_size> k_vec;
   float m_prev = st.m;
-  // smem = base_smem + (stage_idx * DEAL_EACH_TIME + zid * tile_size) * HEAD_DIM;
+  // smem = base_smem + (stage_idx * DEAL_EACH_TIME + zid * tile_size) *
+  // HEAD_DIM;
   smem = k_smem + stage_idx * DEAL_EACH_TIME * HEAD_DIM;
 #pragma unroll
   for (uint32_t j = 0; j < DEAL_EACH_TIME; ++j) {
@@ -171,7 +180,7 @@ __device__ __forceinline__ void compute_qk(const T* cu_q_smem,
         s[j] = 0.f;
       }
 #pragma unroll
-      for(uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
+      for (uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
         Load<T, vec_size>(cu_q_smem + vid * vec_size, &q_vec);
         Load<CacheT, vec_size>(smem + j * HEAD_DIM + vid * vec_size, &k_vec);
         for (uint32_t i = 0; i < vec_size; ++i) {
@@ -211,20 +220,29 @@ __device__ __forceinline__ void compute_qk(const T* cu_q_smem,
   }
 }
 
-template<uint32_t vec_size, uint32_t NUM_VEC_PER_HEAD, uint32_t bdx, uint32_t DEAL_EACH_TIME, uint32_t HEAD_DIM_QK, uint32_t num_tile, typename T, typename CacheT>
-__device__ __forceinline__ void compute_sv(const float *s,
-                                           const CacheT *base_v_smem,
-                                           const uint32_t stage_idx,
-                                           const uint32_t iter_base,
-                                           const uint32_t iter_bound,
-                                           const uint32_t tidx,
-                                           softmax_state_ts<vec_size, T, num_tile>& st) {
+template <uint32_t vec_size,
+          uint32_t NUM_VEC_PER_HEAD,
+          uint32_t bdx,
+          uint32_t DEAL_EACH_TIME,
+          uint32_t HEAD_DIM_QK,
+          uint32_t num_tile,
+          typename T,
+          typename CacheT>
+__device__ __forceinline__ void compute_sv(
+    const float* s,
+    const CacheT* base_v_smem,
+    const uint32_t stage_idx,
+    const uint32_t iter_base,
+    const uint32_t iter_bound,
+    const uint32_t tidx,
+    softmax_state_ts<vec_size, T, num_tile>& st) {
   const CacheT* v_smem;
   AlignedVector<T, vec_size> v_vec;
 #pragma unroll
   for (int j = 0; (j < DEAL_EACH_TIME) && (iter_base + j < iter_bound); ++j) {
-    v_smem = base_v_smem + stage_idx * DEAL_EACH_TIME * HEAD_DIM_QK + j * HEAD_DIM_QK;
-    for(uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
+    v_smem = base_v_smem + stage_idx * DEAL_EACH_TIME * HEAD_DIM_QK +
+             j * HEAD_DIM_QK;
+    for (uint32_t vid = tidx; vid < NUM_VEC_PER_HEAD; vid += bdx) {
       Load<T, vec_size>(v_smem + vid * vec_size, &v_vec);
       uint32_t tile_id = vid / bdx;
 #pragma unroll
diff --git a/custom_ops/gpu_ops/append_attn/decoder_write_cache_with_rope_kernel.h b/custom_ops/gpu_ops/append_attn/decoder_write_cache_with_rope_kernel.h
index 8872761956..5095ef1062 100644
--- a/custom_ops/gpu_ops/append_attn/decoder_write_cache_with_rope_kernel.h
+++ b/custom_ops/gpu_ops/append_attn/decoder_write_cache_with_rope_kernel.h
@@ -41,4 +41,5 @@ void DecoderWriteCacheWithRoPEKernel(
     paddle::Tensor* key_cache_out,
     paddle::Tensor* value_cache_out,
     const paddle::optional<paddle::Tensor>& q_norm_weight,
-    const paddle::optional<paddle::Tensor>& k_norm_weight,    const float rms_norm_eps);
+    const paddle::optional<paddle::Tensor>& k_norm_weight,
+    const float rms_norm_eps);
diff --git a/custom_ops/gpu_ops/append_attn/encoder_write_cache_with_rope_kernel.h b/custom_ops/gpu_ops/append_attn/encoder_write_cache_with_rope_kernel.h
index 5ea9ce213f..2349a0e5e1 100644
--- a/custom_ops/gpu_ops/append_attn/encoder_write_cache_with_rope_kernel.h
+++ b/custom_ops/gpu_ops/append_attn/encoder_write_cache_with_rope_kernel.h
@@ -56,46 +56,53 @@ void EncoderWriteCacheWithRopeKernel(
   auto head_dim = meta_data.head_dims;
   bool is_scale_channel_wise = false;
   int rotary_dim = head_dim;
-  if (cache_k_scale && cache_k_scale.get().dims()[0] == head_dim * kv_num_heads) {
+  if (cache_k_scale &&
+      cache_k_scale.get().dims()[0] == head_dim * kv_num_heads) {
     is_scale_channel_wise = true;
   }
-  if (rotary_embs){
-    rotary_dim = rotary_embs.get().dims()[rotary_embs.get().dims().size()-1] * 2;
-    if(rotary_dim < head_dim){
-      if (!use_neox_style || q_norm_weight || k_norm_weight || num_heads == kv_num_heads || is_scale_channel_wise){
+  if (rotary_embs) {
+    rotary_dim =
+        rotary_embs.get().dims()[rotary_embs.get().dims().size() - 1] * 2;
+    if (rotary_dim < head_dim) {
+      if (!use_neox_style || q_norm_weight || k_norm_weight ||
+          num_heads == kv_num_heads || is_scale_channel_wise) {
         PADDLE_THROW(phi::errors::Fatal(
-          "partial_rotary_factor < 1.0 only supports use_neox_rotary_style=True, q_norm_weight/k_norm_weight) is None, GQA and is_scale_channel_wise=false."));
+            "partial_rotary_factor < 1.0 only supports "
+            "use_neox_rotary_style=True, q_norm_weight/k_norm_weight) is None, "
+            "GQA and is_scale_channel_wise=false."));
       }
     }
   }
 
   if (q_norm_weight && k_norm_weight) {
-    if (num_heads != kv_num_heads && !is_scale_channel_wise && !use_neox_style) {
+    if (num_heads != kv_num_heads && !is_scale_channel_wise &&
+        !use_neox_style) {
       gqa_rotary_qk_norm_variable(
-        qkv_out->data<T>(),
-        qkv.data<QKV_TYPE>(),
-        qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-        qkv_biases ? qkv_biases.get().data<T>() : nullptr,
-        rotary_embs.get().data<float>(),
-        batch_id_per_token.data<int>(),
-        cu_seqlens_q.data<int>(),
-        seq_lens_encoder.data<int>(),
-        seq_lens_decoder.data<int>(),
-        token_num,
-        num_heads,
-        kv_num_heads,
-        max_seq_len,
-        rope_3d ? rotary_embs.get().dims()[3] : rotary_embs.get().dims()[2],
-        head_dim,
-        stream,
-        use_neox_style,
-        rope_3d,
-        q_norm_weight ? q_norm_weight.get().data<float>() : nullptr,
-        k_norm_weight ? k_norm_weight.get().data<float>() : nullptr,
-        rms_norm_eps);
+          qkv_out->data<T>(),
+          qkv.data<QKV_TYPE>(),
+          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+          qkv_biases ? qkv_biases.get().data<T>() : nullptr,
+          rotary_embs.get().data<float>(),
+          batch_id_per_token.data<int>(),
+          cu_seqlens_q.data<int>(),
+          seq_lens_encoder.data<int>(),
+          seq_lens_decoder.data<int>(),
+          token_num,
+          num_heads,
+          kv_num_heads,
+          max_seq_len,
+          rope_3d ? rotary_embs.get().dims()[3] : rotary_embs.get().dims()[2],
+          head_dim,
+          stream,
+          use_neox_style,
+          rope_3d,
+          q_norm_weight ? q_norm_weight.get().data<float>() : nullptr,
+          k_norm_weight ? k_norm_weight.get().data<float>() : nullptr,
+          rms_norm_eps);
     } else {
       PD_THROW(
-          "gqa_rotary_qk_norm_variable only support gqa mode. channel wise scale and neox style are not supported");
+          "gqa_rotary_qk_norm_variable only support gqa mode. channel wise "
+          "scale and neox style are not supported");
     }
   } else {
     if (num_heads == kv_num_heads) {
@@ -120,49 +127,48 @@ void EncoderWriteCacheWithRopeKernel(
     } else {
       if (!is_scale_channel_wise) {
         gqa_rotary_qk_variable(
-          qkv_out->data<T>(),
-          qkv.data<QKV_TYPE>(),
-          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-          qkv_biases ? qkv_biases.get().data<T>() : nullptr,
-          rotary_embs.get().data<float>(),
-          batch_id_per_token.data<int>(),
-          cu_seqlens_q.data<int>(),
-          seq_lens_encoder.data<int>(),
-          seq_lens_decoder.data<int>(),
-          token_num,
-          num_heads,
-          kv_num_heads,
-          max_seq_len,
-          rope_3d ? rotary_embs.get().dims()[3] : rotary_embs.get().dims()[2],
-          head_dim,
-          rotary_dim,
-          stream,
-          use_neox_style,
-          rope_3d);
+            qkv_out->data<T>(),
+            qkv.data<QKV_TYPE>(),
+            qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+            qkv_biases ? qkv_biases.get().data<T>() : nullptr,
+            rotary_embs.get().data<float>(),
+            batch_id_per_token.data<int>(),
+            cu_seqlens_q.data<int>(),
+            seq_lens_encoder.data<int>(),
+            seq_lens_decoder.data<int>(),
+            token_num,
+            num_heads,
+            kv_num_heads,
+            max_seq_len,
+            rope_3d ? rotary_embs.get().dims()[3] : rotary_embs.get().dims()[2],
+            head_dim,
+            rotary_dim,
+            stream,
+            use_neox_style,
+            rope_3d);
       } else {
         gqa_rotary_qk_quant_variable(
-          qkv_out->data<T>(),
-          qkv.data<QKV_TYPE>(),
-          qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
-          qkv_biases ? qkv_biases.get().data<T>() : nullptr,
-          cache_k_scale ? cache_k_scale.get().data<T>() : nullptr,
-          cache_v_scale ? cache_v_scale.get().data<T>() : nullptr,
-          rotary_embs.get().data<float>(),
-          batch_id_per_token.data<int>(),
-          cu_seqlens_q.data<int>(),
-          seq_lens_encoder.data<int>(),
-          seq_lens_decoder.data<int>(),
-          token_num,
-          num_heads,
-          kv_num_heads,
-          max_seq_len,
-          rotary_embs.get().dims()[2],
-          head_dim,
-          stream,
-          use_neox_style,
-          rope_3d);
+            qkv_out->data<T>(),
+            qkv.data<QKV_TYPE>(),
+            qkv_out_scales ? qkv_out_scales.get().data<float>() : nullptr,
+            qkv_biases ? qkv_biases.get().data<T>() : nullptr,
+            cache_k_scale ? cache_k_scale.get().data<T>() : nullptr,
+            cache_v_scale ? cache_v_scale.get().data<T>() : nullptr,
+            rotary_embs.get().data<float>(),
+            batch_id_per_token.data<int>(),
+            cu_seqlens_q.data<int>(),
+            seq_lens_encoder.data<int>(),
+            seq_lens_decoder.data<int>(),
+            token_num,
+            num_heads,
+            kv_num_heads,
+            max_seq_len,
+            rotary_embs.get().dims()[2],
+            head_dim,
+            stream,
+            use_neox_style,
+            rope_3d);
       }
-
     }
   }
   const uint32_t block_size = meta_data.block_size;
@@ -178,7 +184,9 @@ void EncoderWriteCacheWithRopeKernel(
                                     stream,
                                     key_cache_out,
                                     value_cache_out);
-  } else if (cache_quant_type_str == "cache_int8" or cache_quant_type_str == "cache_fp8" or cache_quant_type_str == "block_wise_fp8") {
+  } else if (cache_quant_type_str == "cache_int8" or
+             cache_quant_type_str == "cache_fp8" or
+             cache_quant_type_str == "block_wise_fp8") {
     DISPATCH_HEAD_DIM(
         head_dim, HEAD_DIM, {DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {
           CascadeAppendWriteCacheKVC8QKV<T, HEAD_DIM, BLOCK_SIZE>(
@@ -234,23 +242,29 @@ void EncoderWriteCacheWithRopeKernel(
         "cache_int4_zp]");
   }
 
-  const char* fmt_write_cache_completed_signal_str = std::getenv("FLAGS_fmt_write_cache_completed_signal");
-  const char* FLAGS_use_pd_disaggregation_per_chunk = std::getenv("FLAGS_use_pd_disaggregation_per_chunk");
+  const char* fmt_write_cache_completed_signal_str =
+      std::getenv("FLAGS_fmt_write_cache_completed_signal");
+  const char* FLAGS_use_pd_disaggregation_per_chunk =
+      std::getenv("FLAGS_use_pd_disaggregation_per_chunk");
   if (fmt_write_cache_completed_signal_str &&
       (std::strcmp(fmt_write_cache_completed_signal_str, "true") == 0 ||
        std::strcmp(fmt_write_cache_completed_signal_str, "1") == 0)) {
-      if (FLAGS_use_pd_disaggregation_per_chunk &&
-          (std::strcmp(FLAGS_use_pd_disaggregation_per_chunk, "true") == 0 ||
-           std::strcmp(FLAGS_use_pd_disaggregation_per_chunk, "1") == 0)) {
-        cudaLaunchHostFunc(qkv.stream(),
-                           &(RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_per_query),
-                           (void*)nullptr);
-      } else {
-        if (kv_signal_data) {
-          cudaLaunchHostFunc(qkv.stream(),
-                            &RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise,
-                            (void*)(const_cast<int64_t*>(kv_signal_data.get().data<int64_t>())));
-        }
+    if (FLAGS_use_pd_disaggregation_per_chunk &&
+        (std::strcmp(FLAGS_use_pd_disaggregation_per_chunk, "true") == 0 ||
+         std::strcmp(FLAGS_use_pd_disaggregation_per_chunk, "1") == 0)) {
+      cudaLaunchHostFunc(
+          qkv.stream(),
+          &(RemoteCacheKvIpc::
+                save_cache_kv_complete_signal_layerwise_per_query),
+          (void*)nullptr);
+    } else {
+      if (kv_signal_data) {
+        cudaLaunchHostFunc(
+            qkv.stream(),
+            &RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise,
+            (void*)(const_cast<int64_t*>(
+                kv_signal_data.get().data<int64_t>())));
       }
+    }
   }
 }
diff --git a/custom_ops/gpu_ops/append_attn/mem_util.cuh b/custom_ops/gpu_ops/append_attn/mem_util.cuh
index fb735be7ae..25c7a62388 100644
--- a/custom_ops/gpu_ops/append_attn/mem_util.cuh
+++ b/custom_ops/gpu_ops/append_attn/mem_util.cuh
@@ -66,10 +66,10 @@ __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
 #ifdef PADDLE_WITH_CUSTOM_DEVICE_METAX_GPU
   if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
     memset(__cvta_shared_to_generic(smem_int_ptr), 0, 16);
-    memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, 16);
+    memcpy(__cvta_shared_to_generic(smem_int_ptr), (void*)gmem_ptr, 16);
   } else {
     memset(__cvta_shared_to_generic(smem_int_ptr), 0, 16);
-    memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, 16);
+    memcpy(__cvta_shared_to_generic(smem_int_ptr), (void*)gmem_ptr, 16);
   }
 #else
   if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
@@ -100,19 +100,23 @@ __device__ __forceinline__ void pred_load_128b(T* smem_ptr,
     int src_in_bytes = predicate ? 16 : 0;
     if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
       memset(__cvta_shared_to_generic(smem_int_ptr), 0, 16);
-      memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, src_in_bytes);
+      memcpy(__cvta_shared_to_generic(smem_int_ptr),
+             (void*)gmem_ptr,
+             src_in_bytes);
     } else {
       memset(__cvta_shared_to_generic(smem_int_ptr), 0, 16);
-      memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, src_in_bytes);
+      memcpy(__cvta_shared_to_generic(smem_int_ptr),
+             (void*)gmem_ptr,
+             src_in_bytes);
     }
   } else {
     if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
       if (predicate) {
-        memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, 16);
+        memcpy(__cvta_shared_to_generic(smem_int_ptr), (void*)gmem_ptr, 16);
       }
     } else {
       if (predicate) {
-        memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, 16);
+        memcpy(__cvta_shared_to_generic(smem_int_ptr), (void*)gmem_ptr, 16);
       }
     }
   }
@@ -169,10 +173,11 @@ __device__ __forceinline__ void pred_load_64b(T* smem_ptr,
   if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
     int src_in_bytes = predicate ? 8 : 0;
     memset(__cvta_shared_to_generic(smem_int_ptr), 0, 8);
-    memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, src_in_bytes);
+    memcpy(
+        __cvta_shared_to_generic(smem_int_ptr), (void*)gmem_ptr, src_in_bytes);
   } else {
     if (predicate) {
-      memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, 8);
+      memcpy(__cvta_shared_to_generic(smem_int_ptr), (void*)gmem_ptr, 8);
     }
   }
 #else
@@ -207,10 +212,11 @@ __device__ __forceinline__ void pred_load_32b(T* smem_ptr,
   if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
     int src_in_bytes = predicate ? 4 : 0;
     memset(__cvta_shared_to_generic(smem_int_ptr), 0, 4);
-    memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, src_in_bytes);
+    memcpy(
+        __cvta_shared_to_generic(smem_int_ptr), (void*)gmem_ptr, src_in_bytes);
   } else {
     if (predicate) {
-      memcpy(__cvta_shared_to_generic(smem_int_ptr), (void *)gmem_ptr, 4);
+      memcpy(__cvta_shared_to_generic(smem_int_ptr), (void*)gmem_ptr, 4);
     }
   }
 #else
@@ -275,7 +281,6 @@ struct smem_t {
   template <typename T>
   __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {}
 
-
   template <uint32_t stride, uint32_t inv_stride = 0>
   static __device__ __forceinline__ uint32_t get_permuted_offset(uint32_t i,
                                                                  uint32_t j) {
diff --git a/custom_ops/gpu_ops/append_attn/mla_cache_kernel.cuh b/custom_ops/gpu_ops/append_attn/mla_cache_kernel.cuh
index 2efcb7a8c6..ec5b428bda 100644
--- a/custom_ops/gpu_ops/append_attn/mla_cache_kernel.cuh
+++ b/custom_ops/gpu_ops/append_attn/mla_cache_kernel.cuh
@@ -20,10 +20,10 @@
 template <typename T, int VecSize = 1>
 __global__ void decode_absorb_cache_kernel(
     const T* __restrict__ kv_nope,  // [bsz, kv_num_heads, pe_size] 512
-    const T* __restrict__ kv_pe,  // [bsz, kv_num_heads, nope_size] 64
-    T* __restrict__ kv_cache,    // [num_blocks, kv_num_heads, block_size,
-                                  // nope_size]
-    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
+    const T* __restrict__ kv_pe,    // [bsz, kv_num_heads, nope_size] 64
+    T* __restrict__ kv_cache,       // [num_blocks, kv_num_heads, block_size,
+                                    // nope_size]
+    const int* __restrict__ block_tables,  // [bsz, max_blocks_per_seq]
     const int* __restrict__ cu_seqlens_q,
     const int* __restrict__ seq_lens,          // [bsz]
     const int* __restrict__ seq_lens_encoder,  // [bsz]
@@ -62,26 +62,25 @@ __global__ void decode_absorb_cache_kernel(
     const int block_idx = block_table_now[write_seq_id / block_size];
     const int block_offset = write_seq_id % block_size;
 
-    if (bias < nope_hidden_size) { // pe
+    if (bias < nope_hidden_size) {  // pe
       const uint32_t inner_bias = bias;
       const uint32_t hi = inner_bias / nope_size;
       const uint32_t h_bias = inner_bias % nope_size;
-      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
-                             hi * block_size * all_size +
-                             block_offset * all_size + h_bias;
-      const uint32_t ori_idx =
-          start_token_idx * nope_hidden_size + inner_bias;
+      const uint32_t tgt_idx =
+          block_idx * kv_num_heads * block_size * all_size +
+          hi * block_size * all_size + block_offset * all_size + h_bias;
+      const uint32_t ori_idx = start_token_idx * nope_hidden_size + inner_bias;
       Load<T, VecSize>(&kv_nope[ori_idx], &src_vec);
       Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
     } else {
       const uint32_t inner_bias = bias - nope_hidden_size;
       const uint32_t hi = inner_bias / pe_size;
       const uint32_t h_bias = inner_bias % pe_size;
-      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
-                             hi * block_size * all_size +
-                             block_offset * all_size + nope_size + h_bias;
-      const uint32_t ori_idx =
-          start_token_idx * pe_hidden_size + inner_bias;
+      const uint32_t tgt_idx =
+          block_idx * kv_num_heads * block_size * all_size +
+          hi * block_size * all_size + block_offset * all_size + nope_size +
+          h_bias;
+      const uint32_t ori_idx = start_token_idx * pe_hidden_size + inner_bias;
       Load<T, VecSize>(&kv_pe[ori_idx], &src_vec);
       Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
     }
@@ -91,10 +90,10 @@ __global__ void decode_absorb_cache_kernel(
 template <typename T, int VecSize = 1>
 __global__ void speculate_decode_absorb_cache_kernel(
     const T* __restrict__ kv_nope,  // [bsz, kv_num_heads, pe_size] 512
-    const T* __restrict__ kv_pe,  // [bsz, kv_num_heads, nope_size] 64
-    T* __restrict__ kv_cache,    // [num_blocks, kv_num_heads, block_size,
-                                  // nope_size]
-    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
+    const T* __restrict__ kv_pe,    // [bsz, kv_num_heads, nope_size] 64
+    T* __restrict__ kv_cache,       // [num_blocks, kv_num_heads, block_size,
+                                    // nope_size]
+    const int* __restrict__ block_tables,  // [bsz, max_blocks_per_seq]
     const int* __restrict__ batch_id_per_token,
     const int* __restrict__ cu_seqlens_q,
     const int* __restrict__ seq_lens,          // [bsz]
@@ -125,8 +124,7 @@ __global__ void speculate_decode_absorb_cache_kernel(
     if (seq_lens[ori_bi] == 0) continue;
     const int bias = linear_index % hidden_size;
     const int start_token_idx = cu_seqlens_q[ori_bi];
-    const int write_seq_id =
-        seq_lens[ori_bi] + token_id - start_token_idx;
+    const int write_seq_id = seq_lens[ori_bi] + token_id - start_token_idx;
     if (write_seq_id == 0) continue;
 
     const int* block_table_now = nullptr;
@@ -145,26 +143,25 @@ __global__ void speculate_decode_absorb_cache_kernel(
           token_id,
           cu_seqlens_q[ori_bi]);
     }
-    if (bias < nope_hidden_size) { // pe
+    if (bias < nope_hidden_size) {  // pe
       const uint32_t inner_bias = bias;
       const uint32_t hi = inner_bias / nope_size;
       const uint32_t h_bias = inner_bias % nope_size;
-      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
-                             hi * block_size * all_size +
-                             block_offset * all_size + h_bias;
-      const uint32_t ori_idx =
-          token_id * nope_hidden_size + inner_bias;
+      const uint32_t tgt_idx =
+          block_idx * kv_num_heads * block_size * all_size +
+          hi * block_size * all_size + block_offset * all_size + h_bias;
+      const uint32_t ori_idx = token_id * nope_hidden_size + inner_bias;
       Load<T, VecSize>(&kv_nope[ori_idx], &src_vec);
       Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
     } else {
       const uint32_t inner_bias = bias - nope_hidden_size;
       const uint32_t hi = inner_bias / pe_size;
       const uint32_t h_bias = inner_bias % pe_size;
-      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
-                             hi * block_size * all_size +
-                             block_offset * all_size + nope_size + h_bias;
-      const uint32_t ori_idx =
-          token_id * pe_hidden_size + inner_bias;
+      const uint32_t tgt_idx =
+          block_idx * kv_num_heads * block_size * all_size +
+          hi * block_size * all_size + block_offset * all_size + nope_size +
+          h_bias;
+      const uint32_t ori_idx = token_id * pe_hidden_size + inner_bias;
       Load<T, VecSize>(&kv_pe[ori_idx], &src_vec);
       Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
     }
@@ -174,10 +171,10 @@ __global__ void speculate_decode_absorb_cache_kernel(
 template <typename T, int VecSize = 1>
 __global__ void prefill_absorb_cache_kernel(
     const T* __restrict__ kv_nope,  // [bsz, kv_num_heads, pe_size] 512
-    const T* __restrict__ kv_pe,  // [bsz, kv_num_heads, nope_size] 64
-    T* __restrict__ kv_cache,    // [num_blocks, kv_num_heads, block_size,
-                                  // nope_size]
-    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
+    const T* __restrict__ kv_pe,    // [bsz, kv_num_heads, nope_size] 64
+    T* __restrict__ kv_cache,       // [num_blocks, kv_num_heads, block_size,
+                                    // nope_size]
+    const int* __restrict__ block_tables,  // [bsz, max_blocks_per_seq]
     const int* __restrict__ batch_id_per_token,
     const int* __restrict__ cu_seqlens_q,
     const int* __restrict__ seq_lens,          // [bsz]
@@ -206,33 +203,33 @@ __global__ void prefill_absorb_cache_kernel(
     const uint32_t bias = linear_index % hidden_size;
     const uint32_t ori_bi = batch_id_per_token[token_idx];
     if (seq_lens[ori_bi] == 0) continue;
-    const uint32_t ori_seq_id = (token_idx - cu_seqlens_q[ori_bi]) + seq_lens_decoder[ori_bi];
+    const uint32_t ori_seq_id =
+        (token_idx - cu_seqlens_q[ori_bi]) + seq_lens_decoder[ori_bi];
 
     const int* block_table_now = nullptr;
     block_table_now = block_tables + ori_bi * max_blocks_per_seq;
     const uint32_t block_idx = block_table_now[ori_seq_id / block_size];
     const uint32_t block_offset = ori_seq_id % block_size;
 
-    if (bias < nope_hidden_size) { // pe
+    if (bias < nope_hidden_size) {  // pe
       const uint32_t inner_bias = bias;
       const uint32_t hi = inner_bias / nope_size;
       const uint32_t h_bias = inner_bias % nope_size;
-      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
-                             hi * block_size * all_size +
-                             block_offset * all_size + h_bias;
-      const uint32_t ori_idx =
-          token_idx * nope_hidden_size + inner_bias;
+      const uint32_t tgt_idx =
+          block_idx * kv_num_heads * block_size * all_size +
+          hi * block_size * all_size + block_offset * all_size + h_bias;
+      const uint32_t ori_idx = token_idx * nope_hidden_size + inner_bias;
       Load<T, VecSize>(&kv_nope[ori_idx], &src_vec);
       Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
     } else {
       const uint32_t inner_bias = bias - nope_hidden_size;
       const uint32_t hi = inner_bias / pe_size;
       const uint32_t h_bias = inner_bias % pe_size;
-      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
-                             hi * block_size * all_size +
-                             block_offset * all_size + nope_size + h_bias;
-      const uint32_t ori_idx =
-          token_idx * pe_hidden_size + inner_bias;
+      const uint32_t tgt_idx =
+          block_idx * kv_num_heads * block_size * all_size +
+          hi * block_size * all_size + block_offset * all_size + nope_size +
+          h_bias;
+      const uint32_t ori_idx = token_idx * pe_hidden_size + inner_bias;
       Load<T, VecSize>(&kv_pe[ori_idx], &src_vec);
       Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
     }
diff --git a/custom_ops/gpu_ops/append_attn/multiquery_decoder_attention_impl.cuh b/custom_ops/gpu_ops/append_attn/multiquery_decoder_attention_impl.cuh
index 759aa7fbd6..2eb6d6bde5 100644
--- a/custom_ops/gpu_ops/append_attn/multiquery_decoder_attention_impl.cuh
+++ b/custom_ops/gpu_ops/append_attn/multiquery_decoder_attention_impl.cuh
@@ -16,37 +16,40 @@
 #include "decode_attention_func.cuh"
 #include "multiquery_decoder_attention_kernel.h"
 
-#define CHECK(call)                                                             \
-do                                                                              \
-{                                                                               \
-    const cudaError_t error_code = call;                                        \
-    if (error_code != cudaSuccess)                                              \
-    {                                                                           \
-        printf("CUDA Error:\n");                                                \
-        printf("     File:      %s\n", __FILE__);                               \
-        printf("     Line       %d:\n", __LINE__);                              \
-        printf("     Error code:%d\n", error_code);                             \
-        printf("     Error text:%s\n", cudaGetErrorString(error_code));         \
-        exit(1);                                                                \
-    }                                                                           \
-}while(0)
+#define CHECK(call)                                                   \
+  do {                                                                \
+    const cudaError_t error_code = call;                              \
+    if (error_code != cudaSuccess) {                                  \
+      printf("CUDA Error:\n");                                        \
+      printf("     File:      %s\n", __FILE__);                       \
+      printf("     Line       %d:\n", __LINE__);                      \
+      printf("     Error code:%d\n", error_code);                     \
+      printf("     Error text:%s\n", cudaGetErrorString(error_code)); \
+      exit(1);                                                        \
+    }                                                                 \
+  } while (0)
 
-template <typename T, typename OutT, int vec_size, uint32_t bdy, uint32_t HEAD_DIM>
-__global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi_out, // [bsz, num_chunks, num_heads, head_dim]
-                                                    const T * __restrict__ multi_m, // [bsz, num_chunks, num_heads]
-                                                    const T * __restrict__ multi_d, // [bsz, num_chunks, num_heads]
-                                                    const int * __restrict__ seq_lens_q,
-                                                    const int * __restrict__ seq_lens_kv,
-                                                    const int * __restrict__ cu_seqlens_q,
-                                                    const T * __restrict__ shift_bias, // [q_num_heads * HEAD_DIM]
-                                                    const T * __restrict__ smooth_weight, // [q_num_heads * HEAD_DIM]
-                                                    OutT * __restrict__ out, // [token_num, num_heads, head_dim]
-                                                    const float in_scale,
-                                                    const int num_chunks,
-                                                    const int chunk_size,
-                                                    const int max_seq_len,
-                                                    const int num_heads,
-                                                    const int head_dim) {
+template <typename T,
+          typename OutT,
+          int vec_size,
+          uint32_t bdy,
+          uint32_t HEAD_DIM>
+__global__ void merge_varlen_multi_chunks_v2_kernel(
+    const T *__restrict__ multi_out,  // [bsz, num_chunks, num_heads, head_dim]
+    const T *__restrict__ multi_m,    // [bsz, num_chunks, num_heads]
+    const T *__restrict__ multi_d,    // [bsz, num_chunks, num_heads]
+    const int *__restrict__ seq_lens_q,
+    const int *__restrict__ seq_lens_kv,
+    const int *__restrict__ cu_seqlens_q,
+    const T *__restrict__ shift_bias,     // [q_num_heads * HEAD_DIM]
+    const T *__restrict__ smooth_weight,  // [q_num_heads * HEAD_DIM]
+    OutT *__restrict__ out,               // [token_num, num_heads, head_dim]
+    const float in_scale,
+    const int num_chunks,
+    const int chunk_size,
+    const int max_seq_len,
+    const int num_heads,
+    const int head_dim) {
   const int vid = threadIdx.x, ty = threadIdx.y;
   const int qid = blockIdx.x, hid = blockIdx.y;
   const int seq_len_q = seq_lens_q[qid];
@@ -68,12 +71,12 @@ __global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi
   if constexpr (std::is_same<T, half>::value) {
 #pragma unroll
     for (int i = 0; i < vec_size / 2; ++i) {
-      *((half2*)(&res_vec) + i) = make_half2(0, 0);
+      *((half2 *)(&res_vec) + i) = make_half2(0, 0);
     }
   } else if constexpr (std::is_same<T, nv_bfloat16>::value) {
 #pragma unroll
     for (int i = 0; i < vec_size / 2; ++i) {
-      *((nv_bfloat162*)(&res_vec) + i) = make_bfloat162(0, 0);
+      *((nv_bfloat162 *)(&res_vec) + i) = make_bfloat162(0, 0);
     }
   }
   T m;
@@ -92,7 +95,8 @@ __global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi
     const T m_now = multi_m[offset];
     const T d_now = multi_d[offset];
     m = m_prev > m_now ? m_prev : m_now;
-    offset = (qid * num_chunks * num_heads + i * num_heads + hid) * head_dim + vid * vec_size;
+    offset = (qid * num_chunks * num_heads + i * num_heads + hid) * head_dim +
+             vid * vec_size;
     Load<T, vec_size>(&multi_out[offset], &load_vec);
     const T scale1 = hexp(m_prev - m), scale2 = hexp(m_now - m);
     d = d * scale1 + d_now * scale2;
@@ -124,30 +128,47 @@ __global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi
   for (int i = 0; i < vec_size; ++i) {
     out_vec[i] = static_cast<OutT>(st.o[i]);
   }
-  Store<OutT, vec_size>(out_vec, &out[(start_token_ids * num_heads + hid) * head_dim + vid * vec_size]);
+  Store<OutT, vec_size>(
+      out_vec,
+      &out[(start_token_ids * num_heads + hid) * head_dim + vid * vec_size]);
 }
 
-template <bool partition_kv, typename T, typename OutT, typename CacheT, uint32_t NUM_STAGES, uint32_t DEAL_EACH_TIME, uint32_t GROUP_SIZE, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_V,
-          uint32_t BLOCK_SIZE, uint32_t VEC_SIZE, uint32_t CACHE_VEC_SIZE, uint32_t bdx, uint32_t bdy>
-__global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [token_num, num_heads, head_dim]
-                                                    CacheT * __restrict__ cache_k, // [max_block_num, num_heads, block_size, head_dim]
-                                                    CacheT * __restrict__ cache_v,
-                                                    const T * __restrict__ shift_bias, // [q_num_heads * HEAD_DIM]
-                                                    const T * __restrict__ smooth_weight, // [q_num_heads * HEAD_DIM]
-                                                    const int * __restrict__ seq_lens_q,
-                                                    const int * __restrict__ seq_lens_kv,
-                                                    const int * __restrict__ cu_seqlens_q,
-                                                    const int * __restrict__ block_table, // [bsz, block_num_per_seq]
-                                                    const int max_seq_len,
-                                                    const int max_dec_len,
-                                                    const int max_block_num_per_seq,
-                                                    const float scale,
-                                                    const float in_scale,
-                                                    const uint32_t chunk_size,
-                                                    T * __restrict__ tmp_workspace, // [batch_size, num_chunks, num_heads, head_dim]
-                                                    T * __restrict__ tmp_m, // [batch_size, num_chunks, num_heads]
-                                                    T * __restrict__ tmp_d, // [batch_size, num_chunks, num_heads]
-                                                    OutT * __restrict__ out) {
+template <bool partition_kv,
+          typename T,
+          typename OutT,
+          typename CacheT,
+          uint32_t NUM_STAGES,
+          uint32_t DEAL_EACH_TIME,
+          uint32_t GROUP_SIZE,
+          uint32_t HEAD_DIM_QK,
+          uint32_t HEAD_DIM_V,
+          uint32_t BLOCK_SIZE,
+          uint32_t VEC_SIZE,
+          uint32_t CACHE_VEC_SIZE,
+          uint32_t bdx,
+          uint32_t bdy>
+__global__ void multi_query_decode_attention_kernel(
+    T *__restrict__ q,             // [token_num, num_heads, head_dim]
+    CacheT *__restrict__ cache_k,  // [max_block_num, num_heads, block_size,
+                                   // head_dim]
+    CacheT *__restrict__ cache_v,
+    const T *__restrict__ shift_bias,     // [q_num_heads * HEAD_DIM]
+    const T *__restrict__ smooth_weight,  // [q_num_heads * HEAD_DIM]
+    const int *__restrict__ seq_lens_q,
+    const int *__restrict__ seq_lens_kv,
+    const int *__restrict__ cu_seqlens_q,
+    const int *__restrict__ block_table,  // [bsz, block_num_per_seq]
+    const int max_seq_len,
+    const int max_dec_len,
+    const int max_block_num_per_seq,
+    const float scale,
+    const float in_scale,
+    const uint32_t chunk_size,
+    T *__restrict__ tmp_workspace,  // [batch_size, num_chunks, num_heads,
+                                    // head_dim]
+    T *__restrict__ tmp_m,          // [batch_size, num_chunks, num_heads]
+    T *__restrict__ tmp_d,          // [batch_size, num_chunks, num_heads]
+    OutT *__restrict__ out) {
   const uint32_t bidx = blockIdx.x, kv_head_idx = blockIdx.z;
   const uint32_t bid = bidx, gid = threadIdx.y;
   const uint32_t tidx = threadIdx.x;
@@ -167,9 +188,9 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   if (q_len <= 0) {
     return;
   }
-  uint32_t kv_len = seq_lens_kv[bid]; // !!!!!!!!
+  uint32_t kv_len = seq_lens_kv[bid];  // !!!!!!!!
   if (kv_len <= 0) {
-     return;
+    return;
   }
   kv_len += q_len;
   const uint32_t num_chunk_this_seq = div_up(kv_len, chunk_size);
@@ -180,23 +201,24 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   }
 
   const uint32_t chunk_start = partition_kv ? chunk_id * chunk_size : 0;
-  const uint32_t chunk_end = partition_kv ? min(kv_len, chunk_start + chunk_size) : kv_len;
+  const uint32_t chunk_end =
+      partition_kv ? min(kv_len, chunk_start + chunk_size) : kv_len;
   const uint32_t chunk_len = chunk_end - chunk_start;
 
   extern __shared__ uint8_t smem[];
   const T *q_now = q + (q_start_idx * q_num_heads + q_head_idx) * HEAD_DIM_QK;
-  T *q_smem = reinterpret_cast<T*>(smem); // [HEAD_DIM_QK * sizeof(T)]
+  T *q_smem = reinterpret_cast<T *>(smem);  // [HEAD_DIM_QK * sizeof(T)]
   T *cu_q_smem = q_smem + gid * HEAD_DIM_QK;
 #pragma unroll
-  for(uint32_t vid = tidx; vid < num_vec_per_head_qk; vid += bdx) {
-    ((float4*)(&cu_q_smem[vid * VEC_SIZE]))[0] = ((float4*)(&q_now[vid * VEC_SIZE]))[0];
-
+  for (uint32_t vid = tidx; vid < num_vec_per_head_qk; vid += bdx) {
+    ((float4 *)(&cu_q_smem[vid * VEC_SIZE]))[0] =
+        ((float4 *)(&q_now[vid * VEC_SIZE]))[0];
   }
   __syncthreads();
   using VecT = AlignedVector<T, VEC_SIZE>;
   VecT q_vec;
 #pragma unroll
-  for(uint32_t vid = tidx; vid < num_vec_per_head_qk; vid += bdx) {
+  for (uint32_t vid = tidx; vid < num_vec_per_head_qk; vid += bdx) {
     Load<T, VEC_SIZE>(cu_q_smem + vid * VEC_SIZE, &q_vec);
     for (uint32_t i = 0; i < VEC_SIZE; ++i) {
       q_vec[i] *= scale;
@@ -204,8 +226,8 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
     Store<T, VEC_SIZE>(q_vec, cu_q_smem + vid * VEC_SIZE);
   }
 
-
-  CacheT *kv_smem = reinterpret_cast<CacheT*>(smem + GROUP_SIZE * HEAD_DIM_QK * sizeof(CacheT));
+  CacheT *kv_smem = reinterpret_cast<CacheT *>(smem + GROUP_SIZE * HEAD_DIM_QK *
+                                                          sizeof(CacheT));
   uint32_t stage_idx = 0;
   constexpr int loop_times = DEAL_EACH_TIME / bdy;
 #pragma unroll
@@ -214,24 +236,27 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
     for (int j = 0; j < loop_times; ++j) {
       const uint32_t k_seq_offset = i * DEAL_EACH_TIME + j * bdy + gid;
       const uint32_t k_seq_id = chunk_start + k_seq_offset;
-      produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
-        kv_smem,
-        cache_k,
-        block_table_now,
-        k_seq_id,
-        k_seq_offset,
-        kv_head_idx,
-        kv_num_heads,
-        tidx,
-        chunk_start,
-        chunk_end
-      );
+      produce_kv<SharedMemFillMode::kNoFill,
+                 HEAD_DIM_QK,
+                 VEC_SIZE,
+                 num_vec_per_head_qk,
+                 bdx,
+                 BLOCK_SIZE,
+                 CACHE_VEC_SIZE>(kv_smem,
+                                 cache_k,
+                                 block_table_now,
+                                 k_seq_id,
+                                 k_seq_offset,
+                                 kv_head_idx,
+                                 kv_num_heads,
+                                 tidx,
+                                 chunk_start,
+                                 chunk_end);
     }
     commit_group();
     stage_idx = (stage_idx + 1) % NUM_STAGES;
   }
 
-
   softmax_state_ts<VEC_SIZE, T, num_tile_v> st;
   float s[DEAL_EACH_TIME];
 
@@ -240,48 +265,55 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
     wait_group<NUM_STAGES - 1>();
     __syncthreads();
     // compute qk
-    compute_qk<VEC_SIZE, num_vec_per_head_qk, bdx, bdy, HEAD_DIM_QK, DEAL_EACH_TIME, num_tile_v>(
-      cu_q_smem,
-      kv_smem,
-      chunk_start + iter * DEAL_EACH_TIME,
-      stage_idx,
-      iter * DEAL_EACH_TIME,
-      chunk_len,
-      tidx,
-      gid,
-      scale,
-      s,
-      st
-    );
+    compute_qk<VEC_SIZE,
+               num_vec_per_head_qk,
+               bdx,
+               bdy,
+               HEAD_DIM_QK,
+               DEAL_EACH_TIME,
+               num_tile_v>(cu_q_smem,
+                           kv_smem,
+                           chunk_start + iter * DEAL_EACH_TIME,
+                           stage_idx,
+                           iter * DEAL_EACH_TIME,
+                           chunk_len,
+                           tidx,
+                           gid,
+                           scale,
+                           s,
+                           st);
     __syncthreads();
 
     // compute sv
-    compute_sv<VEC_SIZE, num_vec_per_head_v, bdx, DEAL_EACH_TIME, HEAD_DIM_QK, num_tile_v>(
-      s,
-      kv_smem,
-      stage_idx,
-      iter * DEAL_EACH_TIME,
-      chunk_len,
-      tidx,
-      st
-    );
+    compute_sv<VEC_SIZE,
+               num_vec_per_head_v,
+               bdx,
+               DEAL_EACH_TIME,
+               HEAD_DIM_QK,
+               num_tile_v>(
+        s, kv_smem, stage_idx, iter * DEAL_EACH_TIME, chunk_len, tidx, st);
     __syncthreads();
 
 #pragma unroll
     for (int j = 0; j < loop_times; ++j) {
       const uint32_t k_seq_offset = j * bdy + gid;
-      produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
-        kv_smem,
-        cache_k,
-        block_table_now,
-        chunk_start + k_seq_offset + (iter + NUM_STAGES) * DEAL_EACH_TIME,
-        stage_idx * DEAL_EACH_TIME + k_seq_offset,
-        kv_head_idx,
-        kv_num_heads,
-        tidx,
-        chunk_start,
-        chunk_end
-      );
+      produce_kv<SharedMemFillMode::kNoFill,
+                 HEAD_DIM_QK,
+                 VEC_SIZE,
+                 num_vec_per_head_qk,
+                 bdx,
+                 BLOCK_SIZE,
+                 CACHE_VEC_SIZE>(
+          kv_smem,
+          cache_k,
+          block_table_now,
+          chunk_start + k_seq_offset + (iter + NUM_STAGES) * DEAL_EACH_TIME,
+          stage_idx * DEAL_EACH_TIME + k_seq_offset,
+          kv_head_idx,
+          kv_num_heads,
+          tidx,
+          chunk_start,
+          chunk_end);
     }
     commit_group();
     stage_idx = (stage_idx + 1) % NUM_STAGES;
@@ -290,45 +322,59 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   __syncthreads();
 
   // normize if not partition_kv
-  for(uint32_t vid = tidx; vid < num_vec_per_head_v; vid += bdx) {
+  for (uint32_t vid = tidx; vid < num_vec_per_head_v; vid += bdx) {
     const uint32_t tile_id = vid / bdx;
     if (!partition_kv || num_chunk_this_seq == 1) {
       st.normalize(tile_id);
     }
     if (partition_kv && num_chunk_this_seq > 1) {
-      const uint32_t head_idx = (bid * num_chunks + chunk_id) * q_num_heads + q_head_idx;
-      Store<T, VEC_SIZE>(st.o[tile_id], tmp_workspace + head_idx * HEAD_DIM_V + vid * VEC_SIZE);
+      const uint32_t head_idx =
+          (bid * num_chunks + chunk_id) * q_num_heads + q_head_idx;
+      Store<T, VEC_SIZE>(
+          st.o[tile_id],
+          tmp_workspace + head_idx * HEAD_DIM_V + vid * VEC_SIZE);
       tmp_m[head_idx] = st.m;
       tmp_d[head_idx] = st.d;
     } else {
-      Store<OutT, VEC_SIZE>(st.o[tile_id], out + (q_write_idx * q_num_heads + q_head_idx) * HEAD_DIM_V + vid * VEC_SIZE);
+      Store<OutT, VEC_SIZE>(
+          st.o[tile_id],
+          out + (q_write_idx * q_num_heads + q_head_idx) * HEAD_DIM_V +
+              vid * VEC_SIZE);
     }
   }
 }
 
-
-template <typename T, uint32_t GROUP_SIZE, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_V, uint32_t BLOCK_SIZE, bool CAUSAL, uint32_t NUM_STAGE, uint32_t cache_bytes, uint32_t DEAL_EACH_TIME>
+template <typename T,
+          uint32_t GROUP_SIZE,
+          uint32_t HEAD_DIM_QK,
+          uint32_t HEAD_DIM_V,
+          uint32_t BLOCK_SIZE,
+          bool CAUSAL,
+          uint32_t NUM_STAGE,
+          uint32_t cache_bytes,
+          uint32_t DEAL_EACH_TIME>
 void MultiQueryDecoderAttention(
-  const AppendAttnMetaData& meta_data,
-  cudaStream_t &stream,
-  const paddle::Tensor &q,
-  const paddle::Tensor &cache_k, // [max_block_num, num_kv_heads, block_size, head_dim]
-  const paddle::Tensor &cache_v, // [num_kv_heads, head_dim]
-  const paddle::optional<paddle::Tensor>& attn_mask,
-  const paddle::optional<paddle::Tensor>& shift_bias,
-  const paddle::optional<paddle::Tensor>& smooth_weight,
-  const paddle::Tensor &seq_lens_q,
-  const paddle::Tensor &seq_lens_kv,
-  const paddle::Tensor &batch_id_per_token,
-  const paddle::Tensor &cu_seqlens_q,
-  const paddle::Tensor &block_table,
-  const int max_seq_len,
-  const int max_dec_len,
-  const float rope_scale,
-  const float rope_theta,
-  const float softmax_scale,
-  const float in_scale,
-  paddle::Tensor *out) {
+    const AppendAttnMetaData &meta_data,
+    cudaStream_t &stream,
+    const paddle::Tensor &q,
+    const paddle::Tensor
+        &cache_k,  // [max_block_num, num_kv_heads, block_size, head_dim]
+    const paddle::Tensor &cache_v,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor> &attn_mask,
+    const paddle::optional<paddle::Tensor> &shift_bias,
+    const paddle::optional<paddle::Tensor> &smooth_weight,
+    const paddle::Tensor &seq_lens_q,
+    const paddle::Tensor &seq_lens_kv,
+    const paddle::Tensor &batch_id_per_token,
+    const paddle::Tensor &cu_seqlens_q,
+    const paddle::Tensor &block_table,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float rope_scale,
+    const float rope_theta,
+    const float softmax_scale,
+    const float in_scale,
+    paddle::Tensor *out) {
   using NV_TYPE = typename cascade_attn_type_traits<T>::type;
 
   auto num_heads = meta_data.q_num_heads;
@@ -338,8 +384,8 @@ void MultiQueryDecoderAttention(
   auto max_block_num_per_seq = meta_data.max_blocks_per_seq;
   constexpr int num_stages = NUM_STAGE;
 
-  constexpr int vec_size = 16 / sizeof(T); // 8 16 32
-  constexpr int cache_vec_size = 128 / cache_bytes; // 8 16 32
+  constexpr int vec_size = 16 / sizeof(T);           // 8 16 32
+  constexpr int cache_vec_size = 128 / cache_bytes;  // 8 16 32
   constexpr int blockxc = HEAD_DIM_QK / cache_vec_size;
   constexpr int num_vec_per_head = HEAD_DIM_QK / vec_size;
   constexpr int blockx = num_vec_per_head < 32 ? num_vec_per_head : 32;
@@ -349,12 +395,25 @@ void MultiQueryDecoderAttention(
 
   constexpr int num_threads = blockx * blocky;
 
-  auto splitkv_kernel = multi_query_decode_attention_kernel<true, NV_TYPE, NV_TYPE, NV_TYPE, num_stages, DEAL_EACH_TIME, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V,
-                                                                        BLOCK_SIZE, vec_size, cache_vec_size, blockx, blocky>;
+  auto splitkv_kernel = multi_query_decode_attention_kernel<true,
+                                                            NV_TYPE,
+                                                            NV_TYPE,
+                                                            NV_TYPE,
+                                                            num_stages,
+                                                            DEAL_EACH_TIME,
+                                                            GROUP_SIZE,
+                                                            HEAD_DIM_QK,
+                                                            HEAD_DIM_V,
+                                                            BLOCK_SIZE,
+                                                            vec_size,
+                                                            cache_vec_size,
+                                                            blockx,
+                                                            blocky>;
   uint32_t cache_smem_bytes = 0;
 
   const T *shift_bias_ptr = shift_bias ? shift_bias.get().data<T>() : nullptr;
-  const T *smooth_weight_ptr = smooth_weight ? smooth_weight.get().data<T>() : nullptr;
+  const T *smooth_weight_ptr =
+      smooth_weight ? smooth_weight.get().data<T>() : nullptr;
   cache_smem_bytes = num_stages * DEAL_EACH_TIME * HEAD_DIM_QK * sizeof(T);
 
   const uint32_t chunk_size = get_max_partition_size(bsz);
@@ -363,51 +422,64 @@ void MultiQueryDecoderAttention(
 
   if (smem_size >= 48 * 1024) {
     cudaFuncSetAttribute(
-      splitkv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+        splitkv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
   }
   const int dev_id = 0;
   int sm_count;
   int act_blocks_per_sm;
   cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
   cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-    &act_blocks_per_sm, splitkv_kernel, num_threads, smem_size);
+      &act_blocks_per_sm, splitkv_kernel, num_threads, smem_size);
   assert(act_blocks_per_sm > 1);
 
   const int num_blocks_per_wave = sm_count * act_blocks_per_sm;
   const int num_blocks_need = gridx * num_chunks * kv_num_heads;
   const int max_num_chunks = div_up(num_blocks_per_wave, num_blocks_need);
-  const float ratio = static_cast<float>(num_blocks_need) / static_cast<float>(num_blocks_per_wave);
+  const float ratio = static_cast<float>(num_blocks_need) /
+                      static_cast<float>(num_blocks_per_wave);
 
   dim3 grids(gridx, num_chunks, kv_num_heads);
   dim3 blocks(blockx, blocky);
   if (num_chunks <= 1) {
-    auto no_splitkv_kernel = multi_query_decode_attention_kernel<false, NV_TYPE, NV_TYPE, NV_TYPE, num_stages, DEAL_EACH_TIME, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V, BLOCK_SIZE, vec_size,
-                                                                             cache_vec_size, blockx, blocky>;
+    auto no_splitkv_kernel = multi_query_decode_attention_kernel<false,
+                                                                 NV_TYPE,
+                                                                 NV_TYPE,
+                                                                 NV_TYPE,
+                                                                 num_stages,
+                                                                 DEAL_EACH_TIME,
+                                                                 GROUP_SIZE,
+                                                                 HEAD_DIM_QK,
+                                                                 HEAD_DIM_V,
+                                                                 BLOCK_SIZE,
+                                                                 vec_size,
+                                                                 cache_vec_size,
+                                                                 blockx,
+                                                                 blocky>;
     if (smem_size >= 48 * 1024) {
-      cudaFuncSetAttribute(
-        no_splitkv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+      cudaFuncSetAttribute(no_splitkv_kernel,
+                           cudaFuncAttributeMaxDynamicSharedMemorySize,
+                           smem_size);
     }
     no_splitkv_kernel<<<grids, blocks, smem_size, stream>>>(
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(q.data<T>())),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(cache_k.data<T>())),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(cache_v.data<T>())),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(shift_bias_ptr)),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
-      seq_lens_q.data<int>(),
-      seq_lens_kv.data<int>(),
-      cu_seqlens_q.data<int>(),
-      block_table.data<int>(),
-      max_seq_len,
-      max_dec_len,
-      max_block_num_per_seq,
-      softmax_scale,
-      in_scale,
-      chunk_size,
-      nullptr,
-      nullptr,
-      nullptr,
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>()))
-    );
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(q.data<T>())),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_k.data<T>())),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_v.data<T>())),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(shift_bias_ptr)),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(smooth_weight_ptr)),
+        seq_lens_q.data<int>(),
+        seq_lens_kv.data<int>(),
+        cu_seqlens_q.data<int>(),
+        block_table.data<int>(),
+        max_seq_len,
+        max_dec_len,
+        max_block_num_per_seq,
+        softmax_scale,
+        in_scale,
+        chunk_size,
+        nullptr,
+        nullptr,
+        nullptr,
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(out->data<T>())));
 
     // CHECK(cudaGetLastError());
     // CHECK(cudaDeviceSynchronize());
@@ -417,34 +489,33 @@ void MultiQueryDecoderAttention(
     tmp_workspace = allocator->Allocate(
         phi::SizeOf(q.dtype()) *
         static_cast<size_t>(bsz * num_chunks * num_heads * HEAD_DIM_V));
-    tmp_m = allocator->Allocate(
-        phi::SizeOf(q.dtype()) *
-        static_cast<size_t>(bsz * num_chunks * num_heads));
-    tmp_d = allocator->Allocate(
-        phi::SizeOf(q.dtype()) *
-        static_cast<size_t>(bsz * num_chunks * num_heads));
+    tmp_m =
+        allocator->Allocate(phi::SizeOf(q.dtype()) *
+                            static_cast<size_t>(bsz * num_chunks * num_heads));
+    tmp_d =
+        allocator->Allocate(phi::SizeOf(q.dtype()) *
+                            static_cast<size_t>(bsz * num_chunks * num_heads));
 
     splitkv_kernel<<<grids, blocks, smem_size, stream>>>(
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(q.data<T>())),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(cache_k.data<T>())),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(cache_v.data<T>())),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(shift_bias_ptr)),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
-      seq_lens_q.data<int>(),
-      seq_lens_kv.data<int>(),
-      cu_seqlens_q.data<int>(),
-      block_table.data<int>(),
-      max_seq_len,
-      max_dec_len,
-      max_block_num_per_seq,
-      softmax_scale,
-      in_scale,
-      chunk_size,
-      reinterpret_cast<NV_TYPE*>(tmp_workspace->ptr()),
-      reinterpret_cast<NV_TYPE*>(tmp_m->ptr()),
-      reinterpret_cast<NV_TYPE*>(tmp_d->ptr()),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>()))
-    );
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(q.data<T>())),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_k.data<T>())),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(cache_v.data<T>())),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(shift_bias_ptr)),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(smooth_weight_ptr)),
+        seq_lens_q.data<int>(),
+        seq_lens_kv.data<int>(),
+        cu_seqlens_q.data<int>(),
+        block_table.data<int>(),
+        max_seq_len,
+        max_dec_len,
+        max_block_num_per_seq,
+        softmax_scale,
+        in_scale,
+        chunk_size,
+        reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
+        reinterpret_cast<NV_TYPE *>(tmp_m->ptr()),
+        reinterpret_cast<NV_TYPE *>(tmp_d->ptr()),
+        reinterpret_cast<NV_TYPE *>(const_cast<T *>(out->data<T>())));
     // CHECK(cudaGetLastError());
     // CHECK(cudaDeviceSynchronize());
 
@@ -452,23 +523,27 @@ void MultiQueryDecoderAttention(
     constexpr int bdy = 256 / mblockx;
     dim3 grids_merge(bsz, num_heads);
     dim3 blocks_merge(mblockx, bdy);
-    merge_varlen_multi_chunks_v2_kernel<NV_TYPE, NV_TYPE, vec_size, bdy, HEAD_DIM_V><<<grids_merge, blocks_merge, 0, stream>>>(
-      reinterpret_cast<NV_TYPE*>(tmp_workspace->ptr()),
-      reinterpret_cast<NV_TYPE*>(tmp_m->ptr()),
-      reinterpret_cast<NV_TYPE*>(tmp_d->ptr()),
-      seq_lens_q.data<int>(),
-      seq_lens_kv.data<int>(),
-      cu_seqlens_q.data<int>(),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(shift_bias_ptr)),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
-      reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>())),
-      in_scale,
-      num_chunks,
-      chunk_size,
-      max_seq_len,
-      num_heads,
-      HEAD_DIM_V
-    );
+    merge_varlen_multi_chunks_v2_kernel<NV_TYPE,
+                                        NV_TYPE,
+                                        vec_size,
+                                        bdy,
+                                        HEAD_DIM_V>
+        <<<grids_merge, blocks_merge, 0, stream>>>(
+            reinterpret_cast<NV_TYPE *>(tmp_workspace->ptr()),
+            reinterpret_cast<NV_TYPE *>(tmp_m->ptr()),
+            reinterpret_cast<NV_TYPE *>(tmp_d->ptr()),
+            seq_lens_q.data<int>(),
+            seq_lens_kv.data<int>(),
+            cu_seqlens_q.data<int>(),
+            reinterpret_cast<NV_TYPE *>(const_cast<T *>(shift_bias_ptr)),
+            reinterpret_cast<NV_TYPE *>(const_cast<T *>(smooth_weight_ptr)),
+            reinterpret_cast<NV_TYPE *>(const_cast<T *>(out->data<T>())),
+            in_scale,
+            num_chunks,
+            chunk_size,
+            max_seq_len,
+            num_heads,
+            HEAD_DIM_V);
   }
   // CHECK(cudaGetLastError());
   // CHECK(cudaDeviceSynchronize());
diff --git a/custom_ops/gpu_ops/append_attn/multiquery_decoder_attention_kernel.h b/custom_ops/gpu_ops/append_attn/multiquery_decoder_attention_kernel.h
index a1f8aa7cc8..457f383e5e 100644
--- a/custom_ops/gpu_ops/append_attn/multiquery_decoder_attention_kernel.h
+++ b/custom_ops/gpu_ops/append_attn/multiquery_decoder_attention_kernel.h
@@ -15,25 +15,34 @@
 
 #include "decode_attention_func.cuh"
 
-template <typename T, uint32_t GROUP_SIZE, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_V, uint32_t BLOCK_SIZE, bool CAUSAL, uint32_t NUM_STAGE, uint32_t cache_bytes, uint32_t DEAL_EACH_TIME>
+template <typename T,
+          uint32_t GROUP_SIZE,
+          uint32_t HEAD_DIM_QK,
+          uint32_t HEAD_DIM_V,
+          uint32_t BLOCK_SIZE,
+          bool CAUSAL,
+          uint32_t NUM_STAGE,
+          uint32_t cache_bytes,
+          uint32_t DEAL_EACH_TIME>
 void MultiQueryDecoderAttention(
-  const AppendAttnMetaData& meta_data,
-  cudaStream_t &stream,
-  const paddle::Tensor &q,
-  const paddle::Tensor &cache_k, // [max_block_num, num_kv_heads, block_size, head_dim]
-  const paddle::Tensor &cache_v, // [num_kv_heads, head_dim]
-  const paddle::optional<paddle::Tensor>& attn_mask,
-  const paddle::optional<paddle::Tensor>& shift_bias,
-  const paddle::optional<paddle::Tensor>& smooth_weight,
-  const paddle::Tensor &seq_lens_q,
-  const paddle::Tensor &seq_lens_kv,
-  const paddle::Tensor &batch_id_per_token,
-  const paddle::Tensor &cu_seqlens_q,
-  const paddle::Tensor &block_table,
-  const int max_seq_len,
-  const int max_dec_len,
-  const float rope_scale,
-  const float rope_theta,
-  const float softmax_scale,
-  const float in_scale,
-  paddle::Tensor *out);
+    const AppendAttnMetaData &meta_data,
+    cudaStream_t &stream,
+    const paddle::Tensor &q,
+    const paddle::Tensor
+        &cache_k,  // [max_block_num, num_kv_heads, block_size, head_dim]
+    const paddle::Tensor &cache_v,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor> &attn_mask,
+    const paddle::optional<paddle::Tensor> &shift_bias,
+    const paddle::optional<paddle::Tensor> &smooth_weight,
+    const paddle::Tensor &seq_lens_q,
+    const paddle::Tensor &seq_lens_kv,
+    const paddle::Tensor &batch_id_per_token,
+    const paddle::Tensor &cu_seqlens_q,
+    const paddle::Tensor &block_table,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float rope_scale,
+    const float rope_theta,
+    const float softmax_scale,
+    const float in_scale,
+    paddle::Tensor *out);
diff --git a/custom_ops/gpu_ops/append_attn/utils.cuh b/custom_ops/gpu_ops/append_attn/utils.cuh
index 35af00822f..fa5b3bca17 100644
--- a/custom_ops/gpu_ops/append_attn/utils.cuh
+++ b/custom_ops/gpu_ops/append_attn/utils.cuh
@@ -27,7 +27,7 @@ struct AppendAttnMetaData {
   int head_dims;
   int head_dims_v;
   int max_blocks_per_seq;
-  const int *mask_offset = nullptr;
+  const int* mask_offset = nullptr;
 };
 
 __forceinline__ __host__ __device__ int div_up(int a, int b) {
@@ -110,29 +110,33 @@ __device__ __forceinline__ uint32_t sub_if_greater_or_zero(uint32_t x,
 
 /******************************FASTER CAST*********************************/
 
-inline __device__ static void convert_fp8(__nv_bfloat16* result, const uint32_t& source) {
-
+inline __device__ static void convert_fp8(__nv_bfloat16* result,
+                                          const uint32_t& source) {
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
-    uint32_t dest0;
-    uint32_t dest1;
-    asm volatile( \
-        "{\n" \
-        ".reg .b16 lo, hi;\n" \
-        "mov.b32 {lo, hi}, %2;\n" \
-        "cvt.rn.f16x2.e4m3x2 %0, lo;\n" \
-        "cvt.rn.f16x2.e4m3x2 %1, hi;\n" \
-        "}\n" : "=r"(dest0), "=r"(dest1) : "r"(source));
+  uint32_t dest0;
+  uint32_t dest1;
+  asm volatile(
+      "{\n"
+      ".reg .b16 lo, hi;\n"
+      "mov.b32 {lo, hi}, %2;\n"
+      "cvt.rn.f16x2.e4m3x2 %0, lo;\n"
+      "cvt.rn.f16x2.e4m3x2 %1, hi;\n"
+      "}\n"
+      : "=r"(dest0), "=r"(dest1)
+      : "r"(source));
 
-    ((nv_bfloat162*)(result))[0] = __float22bfloat162_rn(__half22float2(((half2*)(&dest0))[0]));
-    ((nv_bfloat162*)(result))[1] = __float22bfloat162_rn(__half22float2(((half2*)(&dest1))[0]));
+  ((nv_bfloat162*)(result))[0] =
+      __float22bfloat162_rn(__half22float2(((half2*)(&dest0))[0]));
+  ((nv_bfloat162*)(result))[1] =
+      __float22bfloat162_rn(__half22float2(((half2*)(&dest1))[0]));
 #else
-    printf("Do not support fp8 in arch < 890\n");
-    asm("trap;");
+  printf("Do not support fp8 in arch < 890\n");
+  asm("trap;");
 #endif
-
 }
 
-inline __device__ static void convert_fp8(half* result, const uint32_t& source) {
+inline __device__ static void convert_fp8(half* result,
+                                          const uint32_t& source) {
   printf("Do not support fp8 to half although it's very easy.\n");
 }
 
@@ -301,8 +305,8 @@ __forceinline__ __host__ __device__ void vec_cast<nv_bfloat16, float>(
 
 #define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...) \
   switch (head_dim) {                              \
-      case 64: {                                  \
-      constexpr size_t HEAD_DIM = 64;             \
+    case 64: {                                     \
+      constexpr size_t HEAD_DIM = 64;              \
       __VA_ARGS__                                  \
       break;                                       \
     }                                              \
@@ -385,9 +389,8 @@ __forceinline__ __host__ __device__ void vec_cast<nv_bfloat16, float>(
     PD_THROW("not support the cache_type: ", cache_type);                 \
   }
 
-
 #define DISPATCH_DEAL_EACH_TIME(deal_each_time, DEAL_EACH_TIME, ...) \
-  if (deal_each_time == 32) {                                 \
+  if (deal_each_time == 32) {                                        \
     constexpr size_t DEAL_EACH_TIME = 32;                            \
     __VA_ARGS__                                                      \
   } else if (deal_each_time == 64) {                                 \
@@ -404,7 +407,7 @@ __forceinline__ __host__ __device__ void vec_cast<nv_bfloat16, float>(
   } else if (num_threads == 256) {                          \
     constexpr size_t NUM_THREADS = 256;                     \
     __VA_ARGS__                                             \
-   } else {                                                 \
+  } else {                                                  \
     PD_THROW("not support the num_threads", num_threads);   \
   }
 
@@ -456,7 +459,7 @@ __forceinline__ __host__ __device__ void vec_cast<nv_bfloat16, float>(
   }
 
 #define DISPATCH_MLA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \
-  if (group_size == 8) {                              \
+  if (group_size == 8) {                                     \
     constexpr size_t GROUP_SIZE = 8;                         \
     __VA_ARGS__                                              \
   } else if (group_size == 16) {                             \
@@ -538,9 +541,11 @@ inline HOSTDEVICE T roundWithTiesToEven(T x) {
           : xUpper);
 }
 
-
 template <typename T, bool is_need_kv_quant, bool IsFP8, int RoundType = 0>
-__host__ __device__ __forceinline__ uint8_t QuantToC8(const T scale, const T value, const float max_bound, const float min_bound) {
+__host__ __device__ __forceinline__ uint8_t QuantToC8(const T scale,
+                                                      const T value,
+                                                      const float max_bound,
+                                                      const float min_bound) {
   uint8_t eight_bits;
   float quant_value;
   if constexpr (is_need_kv_quant) {
@@ -572,8 +577,8 @@ __host__ __device__ __forceinline__ uint8_t QuantToC8(const T scale, const T val
   return eight_bits;
 }
 
-
-template <typename T, bool IsFP8>inline __device__ static void convert_c8(T * result, const uint32_t& source){
+template <typename T, bool IsFP8>
+inline __device__ static void convert_c8(T* result, const uint32_t& source) {
   if constexpr (IsFP8) {
     convert_fp8(result, source);
   } else {
@@ -583,12 +588,12 @@ template <typename T, bool IsFP8>inline __device__ static void convert_c8(T * re
 
 constexpr int kWarpSize = 32;
 
-template<typename T>
+template <typename T>
 inline __device__ void WelfordCombine1(T b_m2, T* m2) {
   *m2 += b_m2;
 }
 
-template<typename T, int thread_group_width = kWarpSize>
+template <typename T, int thread_group_width = kWarpSize>
 __inline__ __device__ void WelfordWarpReduce(T thread_m2, T* m2) {
   *m2 = thread_m2;
   for (int mask = thread_group_width / 2; mask > 0; mask >>= 1) {
@@ -597,7 +602,7 @@ __inline__ __device__ void WelfordWarpReduce(T thread_m2, T* m2) {
   }
 }
 
-template<typename T, int thread_group_width = kWarpSize>
+template <typename T, int thread_group_width = kWarpSize>
 __inline__ __device__ void WelfordWarpAllReduce(T thread_m2, T* m2) {
   WelfordWarpReduce<T, thread_group_width>(thread_m2, m2);
 }
diff --git a/custom_ops/gpu_ops/beam_search_softmax.cu b/custom_ops/gpu_ops/beam_search_softmax.cu
index a12f2ca17b..e78d7000e9 100644
--- a/custom_ops/gpu_ops/beam_search_softmax.cu
+++ b/custom_ops/gpu_ops/beam_search_softmax.cu
@@ -36,8 +36,8 @@ namespace cub = hipcub;
 static constexpr int kBlockSizeForSmallBeamWidth = 256;
 static constexpr int kMaxVocabPartForStage1FastKernel = 128;
 
-#define CASE_K(K)                                        \
-  case K:                                                \
+#define CASE_K(K)                               \
+  case K:                                       \
     invokeTopKSoftMaxLauncher<T, 2 * K, GROUP>( \
         params, beam_group_idx, stream);        \
     break
@@ -1336,24 +1336,25 @@ adding while op and without affecting the speed. Use a 'fake inplace' method
 here. Not elegant but useful ︸_︸.
 *****/
 
-std::vector<paddle::Tensor> BeamSearchSoftmax(const paddle::Tensor &logits,
-                             const paddle::Tensor &seq_lens,
-                             const paddle::Tensor &stop_flags,       // inplace
-                             const paddle::Tensor &end_ids,
-                             const paddle::Tensor &step_ids,
-                             const paddle::Tensor &max_dec_lens,
-                             const paddle::Tensor &block_tables,     // inplace
-                             const paddle::Tensor &cum_scores,       // inplace
-                             const paddle::Tensor &beam_cache_ids,   // inplace
-                             const paddle::Tensor &beam_hyps,        // inplace
-                             const paddle::Tensor &beam_hyps_score,  // inplace
-                             const paddle::Tensor &beam_finished,    // inplace
-                             const paddle::Tensor &beam_width,
-                             const paddle::Tensor &beam_group_num,
-                             const paddle::Tensor &length_penalty,
-                             const paddle::Tensor &diversity_penalty,
-                             bool fuse_softmax,
-                             bool early_stop) {
+std::vector<paddle::Tensor> BeamSearchSoftmax(
+    const paddle::Tensor &logits,
+    const paddle::Tensor &seq_lens,
+    const paddle::Tensor &stop_flags,  // inplace
+    const paddle::Tensor &end_ids,
+    const paddle::Tensor &step_ids,
+    const paddle::Tensor &max_dec_lens,
+    const paddle::Tensor &block_tables,     // inplace
+    const paddle::Tensor &cum_scores,       // inplace
+    const paddle::Tensor &beam_cache_ids,   // inplace
+    const paddle::Tensor &beam_hyps,        // inplace
+    const paddle::Tensor &beam_hyps_score,  // inplace
+    const paddle::Tensor &beam_finished,    // inplace
+    const paddle::Tensor &beam_width,
+    const paddle::Tensor &beam_group_num,
+    const paddle::Tensor &length_penalty,
+    const paddle::Tensor &diversity_penalty,
+    bool fuse_softmax,
+    bool early_stop) {
   std::vector<int64_t> logits_shape = logits.shape();
   // logits_shape
   auto cu_stream = logits.stream();
@@ -1380,43 +1381,43 @@ std::vector<paddle::Tensor> BeamSearchSoftmax(const paddle::Tensor &logits,
   const int end_ids_len = end_ids.dims()[0];
   const int beam_group_size = beam_width_scalar / beam_group_num_scalar;
 
-  auto next_tokens = paddle::full({logits_shape[0], 1}, 0,  end_ids.type(),
-                                 paddle::GPUPlace());
+  auto next_tokens =
+      paddle::full({logits_shape[0], 1}, 0, end_ids.type(), paddle::GPUPlace());
 
-  auto parent_ids = paddle::full({logits_shape[0], 1}, 0,  end_ids.type(),
-                                 paddle::GPUPlace());
+  auto parent_ids =
+      paddle::full({logits_shape[0], 1}, 0, end_ids.type(), paddle::GPUPlace());
 
-  auto cum_scores_ori = paddle::empty(cum_scores.shape(), logits.type(),
-                                 paddle::GPUPlace());
+  auto cum_scores_ori =
+      paddle::empty(cum_scores.shape(), logits.type(), paddle::GPUPlace());
 
-  auto beam_cache_ids_ori = paddle::empty(beam_cache_ids.shape(), end_ids.type(),
-                                 paddle::GPUPlace());
+  auto beam_cache_ids_ori =
+      paddle::empty(beam_cache_ids.shape(), end_ids.type(), paddle::GPUPlace());
 
-  auto block_tables_ori = paddle::empty(block_tables.shape(), end_ids.type(),
-                                 paddle::GPUPlace());
+  auto block_tables_ori =
+      paddle::empty(block_tables.shape(), end_ids.type(), paddle::GPUPlace());
   cudaMemcpyAsync(cum_scores_ori.mutable_data<float>(),
                   cum_scores.data<float>(),
-                  sizeof(float)*cum_scores.numel(),
+                  sizeof(float) * cum_scores.numel(),
                   cudaMemcpyDeviceToDevice,
                   cu_stream);
   cudaMemcpyAsync(beam_cache_ids_ori.mutable_data<int>(),
                   beam_cache_ids.data<int>(),
-                  sizeof(int)*beam_cache_ids.numel(),
+                  sizeof(int) * beam_cache_ids.numel(),
                   cudaMemcpyDeviceToDevice,
                   cu_stream);
   cudaMemcpyAsync(block_tables_ori.mutable_data<int>(),
                   block_tables.data<int>(),
-                  sizeof(int)*block_tables.numel(),
+                  sizeof(int) * block_tables.numel(),
                   cudaMemcpyDeviceToDevice,
                   cu_stream);
 
   const int tmp_size = batch_size * beam_group_size * beam_group_size * 2;
 
-  auto tmp_topk_id = paddle::full({tmp_size}, 0, end_ids.type(),
-                                 paddle::GPUPlace());
+  auto tmp_topk_id =
+      paddle::full({tmp_size}, 0, end_ids.type(), paddle::GPUPlace());
 
-  auto tmp_topk_val = paddle::full({tmp_size}, 0.0, logits.type(),
-                                 paddle::GPUPlace());
+  auto tmp_topk_val =
+      paddle::full({tmp_size}, 0.0, logits.type(), paddle::GPUPlace());
 
   BeamSearchParams<float> params;
   params.batch_size = batch_size;
@@ -1449,7 +1450,8 @@ std::vector<paddle::Tensor> BeamSearchSoftmax(const paddle::Tensor &logits,
   params.block_tables_out = const_cast<int *>(block_tables.data<int>());
   params.cum_scores_out = const_cast<float *>(cum_scores.data<float>());
   params.beam_hyps_out = const_cast<int *>(beam_hyps.data<int>());
-  params.beam_hyps_score_out = const_cast<float *>(beam_hyps_score.data<float>());
+  params.beam_hyps_score_out =
+      const_cast<float *>(beam_hyps_score.data<float>());
   params.beam_finished = const_cast<bool *>(beam_finished.data<bool>());
   params.stop_flags = const_cast<bool *>(stop_flags.data<bool>());
 
@@ -1470,8 +1472,8 @@ std::vector<paddle::Tensor> BeamSearchSoftmax(const paddle::Tensor &logits,
 
   const int workspace_size = tmp_id_val_size * 2 + tmp_stage1_to_stage2_size;
 
-  auto wsp_buffer_tensor = paddle::full({workspace_size}, 0, logits.type(),
-                                 paddle::GPUPlace());
+  auto wsp_buffer_tensor =
+      paddle::full({workspace_size}, 0, logits.type(), paddle::GPUPlace());
 
   params.tmp_ids = reinterpret_cast<int *>(wsp_buffer_tensor.data<float>());
   params.tmp_vals = wsp_buffer_tensor.data<float>() + tmp_id_val_size;
@@ -1480,11 +1482,9 @@ std::vector<paddle::Tensor> BeamSearchSoftmax(const paddle::Tensor &logits,
   for (int beam_group_idx = 0; beam_group_idx < beam_group_num_scalar;
        ++beam_group_idx) {
     if (beam_group_num_scalar == 1) {
-      invokeTopkSoftMax<float, false>(
-          &params, beam_group_idx, cu_stream);
+      invokeTopkSoftMax<float, false>(&params, beam_group_idx, cu_stream);
     } else {
-      invokeTopkSoftMax<float, true>(
-          &params, beam_group_idx, cu_stream);
+      invokeTopkSoftMax<float, true>(&params, beam_group_idx, cu_stream);
     }
   }
   updateBeamSearchParams<float>(&params, cu_stream);
@@ -1492,54 +1492,66 @@ std::vector<paddle::Tensor> BeamSearchSoftmax(const paddle::Tensor &logits,
 }
 
 std::vector<std::vector<int64_t>> BeamSearchSoftmaxShape(
-                             const std::vector<int64_t> &logits,
-                             const std::vector<int64_t> &seq_lens,
-                             const std::vector<int64_t> &stop_flags,       // inplace
-                             const std::vector<int64_t> &end_ids,
-                             const std::vector<int64_t> &step_ids,
-                             const std::vector<int64_t> &max_dec_lens,
-                             const std::vector<int64_t> &block_tables,     // inplace
-                             const std::vector<int64_t> &cum_scores,       // inplace
-                             const std::vector<int64_t> &beam_cache_ids,   // inplace
-                             const std::vector<int64_t> &beam_hyps,        // inplace
-                             const std::vector<int64_t> &beam_hyps_score,  // inplace
-                             const std::vector<int64_t> &beam_finished,    // inplace
-                             const std::vector<int64_t> &beam_width,
-                             const std::vector<int64_t> &beam_group_num,
-                             const std::vector<int64_t> &length_penalty,
-                             const std::vector<int64_t> &diversity_penalty) {
-    std::vector<int64_t> next_tokens = {logits[0],1};
-    std::vector<int64_t> parent_ids = {logits[0],1};
-    return {next_tokens,parent_ids};
+    const std::vector<int64_t> &logits,
+    const std::vector<int64_t> &seq_lens,
+    const std::vector<int64_t> &stop_flags,  // inplace
+    const std::vector<int64_t> &end_ids,
+    const std::vector<int64_t> &step_ids,
+    const std::vector<int64_t> &max_dec_lens,
+    const std::vector<int64_t> &block_tables,     // inplace
+    const std::vector<int64_t> &cum_scores,       // inplace
+    const std::vector<int64_t> &beam_cache_ids,   // inplace
+    const std::vector<int64_t> &beam_hyps,        // inplace
+    const std::vector<int64_t> &beam_hyps_score,  // inplace
+    const std::vector<int64_t> &beam_finished,    // inplace
+    const std::vector<int64_t> &beam_width,
+    const std::vector<int64_t> &beam_group_num,
+    const std::vector<int64_t> &length_penalty,
+    const std::vector<int64_t> &diversity_penalty) {
+  std::vector<int64_t> next_tokens = {logits[0], 1};
+  std::vector<int64_t> parent_ids = {logits[0], 1};
+  return {next_tokens, parent_ids};
 }
 
 std::vector<paddle::DataType> BeamSearchSoftmaxDtype(
-                             const paddle::DataType &logits,
-                             const paddle::DataType &seq_lens,
-                             const paddle::DataType &stop_flags,       // inplace
-                             const paddle::DataType &end_ids,
-                             const paddle::DataType &step_ids,
-                             const paddle::DataType &max_dec_lens,
-                             const paddle::DataType &block_tables,     // inplace
-                             const paddle::DataType &cum_scores,       // inplace
-                             const paddle::DataType &beam_cache_ids,   // inplace
-                             const paddle::DataType &beam_hyps,        // inplace
-                             const paddle::DataType &beam_hyps_score,  // inplace
-                             const paddle::DataType &beam_finished,    // inplace
-                             const paddle::DataType &beam_width,
-                             const paddle::DataType &beam_group_num,
-                             const paddle::DataType &length_penalty,
-                             const paddle::DataType &diversity_penalty) {
-    return {paddle::DataType::INT32, paddle::DataType::INT32};
+    const paddle::DataType &logits,
+    const paddle::DataType &seq_lens,
+    const paddle::DataType &stop_flags,  // inplace
+    const paddle::DataType &end_ids,
+    const paddle::DataType &step_ids,
+    const paddle::DataType &max_dec_lens,
+    const paddle::DataType &block_tables,     // inplace
+    const paddle::DataType &cum_scores,       // inplace
+    const paddle::DataType &beam_cache_ids,   // inplace
+    const paddle::DataType &beam_hyps,        // inplace
+    const paddle::DataType &beam_hyps_score,  // inplace
+    const paddle::DataType &beam_finished,    // inplace
+    const paddle::DataType &beam_width,
+    const paddle::DataType &beam_group_num,
+    const paddle::DataType &length_penalty,
+    const paddle::DataType &diversity_penalty) {
+  return {paddle::DataType::INT32, paddle::DataType::INT32};
 }
 
 PD_BUILD_STATIC_OP(beam_search_softmax)
-    .Inputs({"logits", "seq_lens", "stop_flags", "end_ids", "step_ids", "max_dec_lens", "block_tables"
-    , "cum_scores", "beam_cache_ids", "beam_hyps", "beam_hyps_score", "beam_finished"
-    , "beam_width", "beam_group_num", "length_penalty", "diversity_penalty"})
+    .Inputs({"logits",
+             "seq_lens",
+             "stop_flags",
+             "end_ids",
+             "step_ids",
+             "max_dec_lens",
+             "block_tables",
+             "cum_scores",
+             "beam_cache_ids",
+             "beam_hyps",
+             "beam_hyps_score",
+             "beam_finished",
+             "beam_width",
+             "beam_group_num",
+             "length_penalty",
+             "diversity_penalty"})
     .Outputs({"next_tokens", "parent_ids"})
-    .Attrs({"fuse_softmax: bool",
-            "early_stop: bool"})
+    .Attrs({"fuse_softmax: bool", "early_stop: bool"})
     .SetKernelFn(PD_KERNEL(BeamSearchSoftmax))
     .SetInferShapeFn(PD_INFER_SHAPE(BeamSearchSoftmaxShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(BeamSearchSoftmaxDtype));
diff --git a/custom_ops/gpu_ops/common/configManager.h b/custom_ops/gpu_ops/common/configManager.h
index d0bb751e97..960df1259b 100644
--- a/custom_ops/gpu_ops/common/configManager.h
+++ b/custom_ops/gpu_ops/common/configManager.h
@@ -22,87 +22,96 @@
 #include <limits>
 
 class ConfigManager {
-public:
-    static ConfigManager& get_instance(const std::string& config_path = "fastdeploy_op_configs.json") {
-        static ConfigManager instance(config_path);
-        return instance;
+ public:
+  static ConfigManager& get_instance(
+      const std::string& config_path = "fastdeploy_op_configs.json") {
+    static ConfigManager instance(config_path);
+    return instance;
+  }
+
+  std::string get_best_config(const std::string& op_name,
+                              const size_t m,
+                              const size_t n,
+                              const size_t k) {
+    initialize();
+    std::string mnk_string = op_name + "-" + std::to_string(update_m(m)) + "x" +
+                             std::to_string(n) + "x" + std::to_string(k);
+    if (configs_.contains(mnk_string)) {
+      return configs_.at(mnk_string);
     }
+    return "";
+  }
 
-    std::string get_best_config(const std::string& op_name, const size_t m, const size_t n, const size_t k) {
-        initialize();
-        std::string mnk_string = op_name + "-" +
-                            std::to_string(update_m(m)) + "x" + std::to_string(n) + "x" + std::to_string(k);
-        if (configs_.contains(mnk_string)) {
-            return configs_.at(mnk_string);
-        }
-        return "";
+  int64_t update_m(const size_t m) {
+    size_t new_m = m;
+    if (m < 4) {
+      return m;
+    } else if (m < 16) {
+      return (m + 3) / 4 * 4;
+    } else if (m < 64) {
+      return (m + 15) / 16 * 16;
+    } else if (m < 256) {
+      return (m + 31) / 32 * 32;
+    } else if (m < 512) {
+      return (m + 63) / 64 * 64;
+    } else if (m < 1024) {
+      return (m + 127) / 128 * 128;
+    } else if (m < 8192) {
+      return (m + 1023) / 1024 * 1024;
+    } else if (m < 32768) {
+      return (m + 4095) / 4096 * 4096;
+    } else {
+      return 32768;
     }
+  }
 
-    int64_t update_m(const size_t m) {
-        size_t new_m = m;
-        if (m < 4) {
-            return m;
-        } else if (m < 16) {
-            return  (m + 3) / 4 * 4;
-        } else if (m < 64) {
-            return (m + 15) / 16 * 16;
-        } else if (m < 256) {
-            return (m + 31) / 32 * 32;
-        } else if (m < 512) {
-            return (m + 63) / 64 * 64;
-        } else if (m < 1024) {
-            return (m + 127) / 128 * 128;
-        } else if (m < 8192) {
-            return (m + 1023) / 1024 * 1024;
-        } else if (m < 32768) {
-            return (m + 4095) / 4096 * 4096;
-        } else {
-            return 32768;
-        }
+  void update(const std::string& op_name,
+              const size_t m,
+              const size_t n,
+              const size_t k,
+              const std::string& config) {
+    initialize();
+    std::string mnk_string = op_name + "-" + std::to_string(update_m(m)) + "x" +
+                             std::to_string(n) + "x" + std::to_string(k);
+    configs_[mnk_string] = config;
+  }
+
+  void print() const {
+    std::cout << configs_.dump(4) << std::endl;  // Pretty print with 4 spaces
+  }
+
+  ~ConfigManager() {
+    std::ofstream file(config_path_);
+    if (file.is_open()) {
+      file << configs_.dump(4);  // Pretty print with 4 spaces
+      file.close();
     }
+  }
 
-    void update(const std::string& op_name, const size_t m, const size_t n, const size_t k, const std::string& config) {
-        initialize();
-        std::string mnk_string = op_name + "-" +
-                            std::to_string(update_m(m)) + "x" + std::to_string(n) + "x" + std::to_string(k);
-        configs_[mnk_string] = config;
+ private:
+  void initialize() {
+    if (initialized_) return;
+    std::ifstream file(config_path_);
+    if (file.is_open()) {
+      try {
+        file >> configs_;
+      } catch (const std::exception& e) {
+        std::cerr << "Error reading configs from " << config_path_ << " : "
+                  << e.what() << std::endl;
+        configs_ = nlohmann::json::object();  // Create an empty JSON object
+      }
+      file.close();
+    } else {
+      configs_ = nlohmann::json::object();  // Create an empty JSON object
     }
+    initialized_ = true;
+  }
 
-    void print() const {
-        std::cout << configs_.dump(4) << std::endl; // Pretty print with 4 spaces
-    }
+  ConfigManager(const std::string& config_path) : config_path_(config_path) {}
+  ConfigManager(const ConfigManager&) = delete;
+  ConfigManager& operator=(const ConfigManager&) = delete;
 
-    ~ConfigManager() {
-        std::ofstream file(config_path_);
-        if (file.is_open()) {
-            file << configs_.dump(4); // Pretty print with 4 spaces
-            file.close();
-        }
-    }
-
-private:
-    void initialize() {
-        if (initialized_) return;
-        std::ifstream file(config_path_);
-        if (file.is_open()) {
-            try {
-                file >> configs_;
-            } catch (const std::exception& e) {
-                std::cerr << "Error reading configs from " << config_path_ << " : " << e.what() << std::endl;
-                configs_ = nlohmann::json::object(); // Create an empty JSON object
-            }
-            file.close();
-        } else {
-            configs_ = nlohmann::json::object(); // Create an empty JSON object
-        }
-        initialized_ = true;
-    }
-
-    ConfigManager(const std::string& config_path) : config_path_(config_path) {}
-    ConfigManager(const ConfigManager&) = delete;
-    ConfigManager& operator=(const ConfigManager&) = delete;
-
-    nlohmann::json configs_;
-    std::string config_path_;
-    bool initialized_{false};
+  nlohmann::json configs_;
+  std::string config_path_;
+  bool initialized_{false};
 };
diff --git a/custom_ops/gpu_ops/common/cudaUtils.h b/custom_ops/gpu_ops/common/cudaUtils.h
index 9bbd1f6e80..7123e33ebb 100644
--- a/custom_ops/gpu_ops/common/cudaUtils.h
+++ b/custom_ops/gpu_ops/common/cudaUtils.h
@@ -16,18 +16,18 @@
 #include <cuda_runtime.h>
 #include "paddle/phi/core/enforce.h"
 
-namespace common
-{
+namespace common {
 
-inline int getSMVersion()
-{
-    int device{-1};
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&device));
-    int sm_major = 0;
-    int sm_minor = 0;
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device));
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device));
-    return sm_major * 10 + sm_minor;
+inline int getSMVersion() {
+  int device{-1};
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&device));
+  int sm_major = 0;
+  int sm_minor = 0;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(
+      &sm_major, cudaDevAttrComputeCapabilityMajor, device));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(
+      &sm_minor, cudaDevAttrComputeCapabilityMinor, device));
+  return sm_major * 10 + sm_minor;
 }
 
-}
+}  // namespace common
diff --git a/custom_ops/gpu_ops/common/quantization.h b/custom_ops/gpu_ops/common/quantization.h
index e6a74760bb..433e795311 100644
--- a/custom_ops/gpu_ops/common/quantization.h
+++ b/custom_ops/gpu_ops/common/quantization.h
@@ -20,312 +20,240 @@
 #include <optional>
 #include <string>
 
-namespace common
-{
+namespace common {
 
-class QuantMode
-{
-    // [WARNING] KEEP BELOW DEFINITION IN SYNC WITH tensorrt_llm/quantization/mode.py
-public:
-    using BaseType = std::uint32_t;
+class QuantMode {
+  // [WARNING] KEEP BELOW DEFINITION IN SYNC WITH
+  // tensorrt_llm/quantization/mode.py
+ public:
+  using BaseType = std::uint32_t;
 
-    explicit constexpr QuantMode(BaseType value) noexcept
-        : mValue{value}
-    {
+  explicit constexpr QuantMode(BaseType value) noexcept : mValue{value} {}
+
+  QuantMode() noexcept = default;
+
+  constexpr QuantMode(QuantMode const&) noexcept = default;
+
+  constexpr QuantMode& operator=(QuantMode const& other) noexcept = default;
+
+  static constexpr QuantMode none() noexcept { return QuantMode(BaseType(0)); }
+
+  static constexpr QuantMode int4Weights() noexcept {
+    return QuantMode(BaseType(1u) << 0);
+  }
+
+  static constexpr QuantMode int8Weights() noexcept {
+    return QuantMode(BaseType(1u) << 1);
+  }
+
+  static constexpr QuantMode activations() noexcept {
+    return QuantMode(BaseType(1u) << 2);
+  }
+
+  static constexpr QuantMode perChannelScaling() noexcept {
+    return QuantMode(BaseType(1u) << 3);
+  }
+
+  static constexpr QuantMode perTokenScaling() noexcept {
+    return QuantMode(BaseType(1u) << 4);
+  }
+
+  static constexpr QuantMode perGroupScaling() noexcept {
+    return QuantMode(BaseType(1u) << 5);
+  }
+
+  static constexpr QuantMode int8KvCache() noexcept {
+    return QuantMode(BaseType(1u) << 6);
+  }
+
+  static constexpr QuantMode fp8KvCache() noexcept {
+    return QuantMode(BaseType(1u) << 7);
+  }
+
+  static constexpr QuantMode fp8Qdq() noexcept {
+    return QuantMode(BaseType(1u) << 8);
+  }
+
+  static constexpr QuantMode fp8RowWise() noexcept {
+    return QuantMode(BaseType(1u) << 3 | BaseType(1u) << 4 | BaseType(1u) << 9);
+  }
+
+  constexpr BaseType value() const noexcept { return mValue; }
+
+  constexpr bool isSet(QuantMode const& mode) const noexcept {
+    return (mValue & mode.value()) == mode.value();
+  }
+
+  constexpr bool hasInt4Weights() const noexcept {
+    return isSet(int4Weights());
+  }
+
+  constexpr bool hasInt8Weights() const noexcept {
+    return isSet(int8Weights());
+  }
+
+  constexpr bool hasActivations() const noexcept {
+    return isSet(activations());
+  }
+
+  constexpr bool hasPerChannelScaling() const noexcept {
+    return isSet(perChannelScaling());
+  }
+
+  constexpr bool hasPerTokenScaling() const noexcept {
+    return isSet(perTokenScaling());
+  }
+
+  constexpr bool hasPerGroupScaling() const noexcept {
+    return isSet(perGroupScaling());
+  }
+
+  constexpr bool hasStaticActivationScaling() const noexcept {
+    return !hasPerTokenScaling();
+  }
+
+  constexpr bool hasInt8KvCache() const noexcept {
+    return isSet(int8KvCache());
+  }
+
+  constexpr bool hasFp8KvCache() const noexcept { return isSet(fp8KvCache()); }
+
+  constexpr bool hasFp8Qdq() const noexcept { return isSet(fp8Qdq()); }
+
+  constexpr bool hasFp8RowWise() const noexcept { return isSet(fp8RowWise()); }
+
+  constexpr bool hasKvCacheQuant() const noexcept {
+    return hasInt8KvCache() || hasFp8KvCache();
+  }
+
+  static constexpr QuantMode fromDescription(bool quantizeWeights = false,
+                                             bool quantizeActivations = false,
+                                             bool perToken = false,
+                                             bool perChannel = false,
+                                             bool perGroup = false,
+                                             bool useInt4Weights = false,
+                                             bool useInt8KvCache = false,
+                                             bool useFp8KvCache = false,
+                                             bool useFp8Qdq = false,
+                                             bool useFp8RowWise = false) {
+    QuantMode quantMode{};
+    if (quantizeWeights) {
+      if (useInt4Weights)
+        quantMode += int4Weights();
+      else
+        quantMode += int8Weights();
     }
 
-    QuantMode() noexcept = default;
-
-    constexpr QuantMode(QuantMode const&) noexcept = default;
-
-    constexpr QuantMode& operator=(QuantMode const& other) noexcept = default;
-
-    static constexpr QuantMode none() noexcept
-    {
-        return QuantMode(BaseType(0));
+    if (quantizeActivations) {
+      quantMode += activations();
     }
 
-    static constexpr QuantMode int4Weights() noexcept
-    {
-        return QuantMode(BaseType(1u) << 0);
+    if (perChannel) {
+      quantMode += QuantMode::perChannelScaling();
+    }
+    if (perToken) {
+      quantMode += QuantMode::perTokenScaling();
+    }
+    if (perGroup) {
+      quantMode += QuantMode::perGroupScaling();
     }
 
-    static constexpr QuantMode int8Weights() noexcept
-    {
-        return QuantMode(BaseType(1u) << 1);
+    if (useInt8KvCache) {
+      quantMode += int8KvCache();
     }
 
-    static constexpr QuantMode activations() noexcept
-    {
-        return QuantMode(BaseType(1u) << 2);
+    if (useFp8KvCache) {
+      quantMode += fp8KvCache();
     }
 
-    static constexpr QuantMode perChannelScaling() noexcept
-    {
-        return QuantMode(BaseType(1u) << 3);
+    if (useFp8Qdq) {
+      quantMode += fp8Qdq();
     }
 
-    static constexpr QuantMode perTokenScaling() noexcept
-    {
-        return QuantMode(BaseType(1u) << 4);
+    if (useFp8RowWise) {
+      quantMode += fp8RowWise();
     }
 
-    static constexpr QuantMode perGroupScaling() noexcept
-    {
-        return QuantMode(BaseType(1u) << 5);
+    return quantMode;
+  }
+
+  static constexpr QuantMode useSmoothQuant(bool perToken = false,
+                                            bool perChannel = false) {
+    return fromDescription(true, true, perToken, perChannel);
+  }
+
+  static constexpr QuantMode useWeightOnly(bool useInt4Weights = false,
+                                           bool perGroup = false) {
+    return fromDescription(true, false, false, false, perGroup, useInt4Weights);
+  }
+
+  static const QuantMode fromQuantAlgo(
+      std::optional<std::string> quantAlgo = std::nullopt,
+      std::optional<std::string> kvCacheQuantAlgo = std::nullopt) {
+    QuantMode quantMode{};
+    if (quantAlgo == "W8A16") {
+      quantMode = useWeightOnly(false, false);
+    } else if (quantAlgo == "W4A16") {
+      quantMode = useWeightOnly(true, false);
+    } else if (quantAlgo == "W4A16_AWQ") {
+      quantMode = useWeightOnly(true, true);
+    } else if (quantAlgo == "W4A8_AWQ") {
+      quantMode = useWeightOnly(true, true);
+    } else if (quantAlgo == "W4A16_GPTQ") {
+      quantMode = useWeightOnly(true, true);
+    } else if (quantAlgo == "W8A8_SQ_PER_CHANNEL") {
+      quantMode = useSmoothQuant(false, true);
+    } else if (quantAlgo == "W8A8_SQ_PER_TENSOR_PLUGIN") {
+      quantMode = useSmoothQuant(false, false);
+    } else if (quantAlgo == "W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN") {
+      quantMode = useSmoothQuant(true, true);
+    } else if (quantAlgo == "W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN") {
+      quantMode = useSmoothQuant(false, true);
+    } else if (quantAlgo == "W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN") {
+      quantMode = useSmoothQuant(true, false);
+    } else if (quantAlgo == "FP8") {
+      quantMode = fromDescription(
+          false, false, false, false, false, false, false, false, true);
+    } else if (quantAlgo == "FP8_ROWWISE") {
+      quantMode = fromDescription(
+          false, false, true, true, false, false, false, false, false, true);
     }
 
-    static constexpr QuantMode int8KvCache() noexcept
-    {
-        return QuantMode(BaseType(1u) << 6);
+    if (kvCacheQuantAlgo == "INT8") {
+      quantMode += int8KvCache();
+    } else if (kvCacheQuantAlgo == "FP8") {
+      quantMode += fp8KvCache();
     }
 
-    static constexpr QuantMode fp8KvCache() noexcept
-    {
-        return QuantMode(BaseType(1u) << 7);
-    }
+    return quantMode;
+  }
 
-    static constexpr QuantMode fp8Qdq() noexcept
-    {
-        return QuantMode(BaseType(1u) << 8);
-    }
+  constexpr QuantMode operator+(QuantMode const& other) const noexcept {
+    return QuantMode(mValue | other.mValue);
+  }
 
-    static constexpr QuantMode fp8RowWise() noexcept
-    {
-        return QuantMode(BaseType(1u) << 3 | BaseType(1u) << 4 | BaseType(1u) << 9);
-    }
+  constexpr QuantMode& operator+=(QuantMode const& other) noexcept {
+    return *this = *this + other;
+  }
 
-    constexpr BaseType value() const noexcept
-    {
-        return mValue;
-    }
+  constexpr QuantMode operator-(QuantMode const& other) const noexcept {
+    return QuantMode(mValue & ~other.mValue);
+  }
 
-    constexpr bool isSet(QuantMode const& mode) const noexcept
-    {
-        return (mValue & mode.value()) == mode.value();
-    }
+  constexpr QuantMode& operator-=(QuantMode const& other) noexcept {
+    return *this = *this - other;
+  }
 
-    constexpr bool hasInt4Weights() const noexcept
-    {
-        return isSet(int4Weights());
-    }
+  constexpr bool operator==(QuantMode const& other) const noexcept {
+    return mValue == other.mValue;
+  }
 
-    constexpr bool hasInt8Weights() const noexcept
-    {
-        return isSet(int8Weights());
-    }
+  constexpr bool operator!=(QuantMode const& other) const noexcept {
+    return !(*this == other);
+  }
 
-    constexpr bool hasActivations() const noexcept
-    {
-        return isSet(activations());
-    }
-
-    constexpr bool hasPerChannelScaling() const noexcept
-    {
-        return isSet(perChannelScaling());
-    }
-
-    constexpr bool hasPerTokenScaling() const noexcept
-    {
-        return isSet(perTokenScaling());
-    }
-
-    constexpr bool hasPerGroupScaling() const noexcept
-    {
-        return isSet(perGroupScaling());
-    }
-
-    constexpr bool hasStaticActivationScaling() const noexcept
-    {
-        return !hasPerTokenScaling();
-    }
-
-    constexpr bool hasInt8KvCache() const noexcept
-    {
-        return isSet(int8KvCache());
-    }
-
-    constexpr bool hasFp8KvCache() const noexcept
-    {
-        return isSet(fp8KvCache());
-    }
-
-    constexpr bool hasFp8Qdq() const noexcept
-    {
-        return isSet(fp8Qdq());
-    }
-
-    constexpr bool hasFp8RowWise() const noexcept
-    {
-        return isSet(fp8RowWise());
-    }
-
-    constexpr bool hasKvCacheQuant() const noexcept
-    {
-        return hasInt8KvCache() || hasFp8KvCache();
-    }
-
-    static constexpr QuantMode fromDescription(bool quantizeWeights = false, bool quantizeActivations = false,
-        bool perToken = false, bool perChannel = false, bool perGroup = false, bool useInt4Weights = false,
-        bool useInt8KvCache = false, bool useFp8KvCache = false, bool useFp8Qdq = false, bool useFp8RowWise = false)
-    {
-        QuantMode quantMode{};
-        if (quantizeWeights)
-        {
-            if (useInt4Weights)
-                quantMode += int4Weights();
-            else
-                quantMode += int8Weights();
-        }
-
-        if (quantizeActivations)
-        {
-            quantMode += activations();
-        }
-
-        if (perChannel)
-        {
-            quantMode += QuantMode::perChannelScaling();
-        }
-        if (perToken)
-        {
-            quantMode += QuantMode::perTokenScaling();
-        }
-        if (perGroup)
-        {
-            quantMode += QuantMode::perGroupScaling();
-        }
-
-        if (useInt8KvCache)
-        {
-            quantMode += int8KvCache();
-        }
-
-        if (useFp8KvCache)
-        {
-            quantMode += fp8KvCache();
-        }
-
-        if (useFp8Qdq)
-        {
-            quantMode += fp8Qdq();
-        }
-
-        if (useFp8RowWise)
-        {
-            quantMode += fp8RowWise();
-        }
-
-        return quantMode;
-    }
-
-    static constexpr QuantMode useSmoothQuant(bool perToken = false, bool perChannel = false)
-    {
-        return fromDescription(true, true, perToken, perChannel);
-    }
-
-    static constexpr QuantMode useWeightOnly(bool useInt4Weights = false, bool perGroup = false)
-    {
-        return fromDescription(true, false, false, false, perGroup, useInt4Weights);
-    }
-
-    static const QuantMode fromQuantAlgo(
-        std::optional<std::string> quantAlgo = std::nullopt, std::optional<std::string> kvCacheQuantAlgo = std::nullopt)
-    {
-        QuantMode quantMode{};
-        if (quantAlgo == "W8A16")
-        {
-            quantMode = useWeightOnly(false, false);
-        }
-        else if (quantAlgo == "W4A16")
-        {
-            quantMode = useWeightOnly(true, false);
-        }
-        else if (quantAlgo == "W4A16_AWQ")
-        {
-            quantMode = useWeightOnly(true, true);
-        }
-        else if (quantAlgo == "W4A8_AWQ")
-        {
-            quantMode = useWeightOnly(true, true);
-        }
-        else if (quantAlgo == "W4A16_GPTQ")
-        {
-            quantMode = useWeightOnly(true, true);
-        }
-        else if (quantAlgo == "W8A8_SQ_PER_CHANNEL")
-        {
-            quantMode = useSmoothQuant(false, true);
-        }
-        else if (quantAlgo == "W8A8_SQ_PER_TENSOR_PLUGIN")
-        {
-            quantMode = useSmoothQuant(false, false);
-        }
-        else if (quantAlgo == "W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN")
-        {
-            quantMode = useSmoothQuant(true, true);
-        }
-        else if (quantAlgo == "W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN")
-        {
-            quantMode = useSmoothQuant(false, true);
-        }
-        else if (quantAlgo == "W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN")
-        {
-            quantMode = useSmoothQuant(true, false);
-        }
-        else if (quantAlgo == "FP8")
-        {
-            quantMode = fromDescription(false, false, false, false, false, false, false, false, true);
-        }
-        else if (quantAlgo == "FP8_ROWWISE")
-        {
-            quantMode = fromDescription(false, false, true, true, false, false, false, false, false, true);
-        }
-
-        if (kvCacheQuantAlgo == "INT8")
-        {
-            quantMode += int8KvCache();
-        }
-        else if (kvCacheQuantAlgo == "FP8")
-        {
-            quantMode += fp8KvCache();
-        }
-
-        return quantMode;
-    }
-
-    constexpr QuantMode operator+(QuantMode const& other) const noexcept
-    {
-        return QuantMode(mValue | other.mValue);
-    }
-
-    constexpr QuantMode& operator+=(QuantMode const& other) noexcept
-    {
-        return *this = *this + other;
-    }
-
-    constexpr QuantMode operator-(QuantMode const& other) const noexcept
-    {
-        return QuantMode(mValue & ~other.mValue);
-    }
-
-    constexpr QuantMode& operator-=(QuantMode const& other) noexcept
-    {
-        return *this = *this - other;
-    }
-
-    constexpr bool operator==(QuantMode const& other) const noexcept
-    {
-        return mValue == other.mValue;
-    }
-
-    constexpr bool operator!=(QuantMode const& other) const noexcept
-    {
-        return !(*this == other);
-    }
-
-private:
-    BaseType mValue{0};
+ private:
+  BaseType mValue{0};
 };
 
-} // namespace common
+}  // namespace common
diff --git a/custom_ops/gpu_ops/cuda_multiprocess.h b/custom_ops/gpu_ops/cuda_multiprocess.h
index c4b3c84109..c8a138e133 100644
--- a/custom_ops/gpu_ops/cuda_multiprocess.h
+++ b/custom_ops/gpu_ops/cuda_multiprocess.h
@@ -52,35 +52,36 @@ namespace cub = hipcub;
 #define GPU(str) cuda##str
 #endif
 
-#define checkCudaErrors(call)                   \
-    do {                                        \
-        GPU(Error_t) err = call;                \
-        if (err != GPU(Success)) {              \
-            printf("CUDA error at %s %d: %s\n", \
-                   __FILE__,                    \
-                   __LINE__,                    \
-                   GPU(GetErrorString)(err));   \
-            exit(EXIT_FAILURE);                 \
-        }                                       \
-    } while (0)
+#define checkCudaErrors(call)             \
+  do {                                    \
+    GPU(Error_t) err = call;              \
+    if (err != GPU(Success)) {            \
+      printf("CUDA error at %s %d: %s\n", \
+             __FILE__,                    \
+             __LINE__,                    \
+             GPU(GetErrorString)(err));   \
+      exit(EXIT_FAILURE);                 \
+    }                                     \
+  } while (0)
 
 typedef struct shmStruct_st {
-    size_t nprocesses;
-    GPU(IpcMemHandle_t) memHandle;
+  size_t nprocesses;
+  GPU(IpcMemHandle_t) memHandle;
 } shmStruct;
 
 typedef struct sharedMemoryInfo_st {
-    void *addr;
-    size_t size;
+  void *addr;
+  size_t size;
 #if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
-    HANDLE shmHandle;
+  HANDLE shmHandle;
 #else
-    int shmFd;
+  int shmFd;
 #endif
 } sharedMemoryInfo;
 
-
-inline int sharedMemoryOpen(const char *name, size_t sz, sharedMemoryInfo *info) {
+inline int sharedMemoryOpen(const char *name,
+                            size_t sz,
+                            sharedMemoryInfo *info) {
 #if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
   info->size = sz;
 
diff --git a/custom_ops/gpu_ops/cutlass_extensions/arch/copy_red_global.hpp b/custom_ops/gpu_ops/cutlass_extensions/arch/copy_red_global.hpp
index 61a41031bf..4b99d9651b 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/arch/copy_red_global.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/arch/copy_red_global.hpp
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 #pragma once
@@ -38,315 +39,331 @@
 
 // Config
 
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700) && (__CUDACC_VER_MAJOR__ >= 10))
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700) && \
+     (__CUDACC_VER_MAJOR__ >= 10))
 #define CUTE_ARCH_RED_F16_SM70_ENABLED
 #endif
 
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12))
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && \
+     (__CUDACC_VER_MAJOR__ >= 12))
 #define CUTE_ARCH_RED_VEC_SM90_ENABLED
 #define CUTE_ARCH_RED_BF16_SM90_ENABLED
 #endif
 
-namespace cute
-{
+namespace cute {
 
 //////////////////////////////////
 // Wrapper around CUDA's atomicAdd
 //////////////////////////////////
 
 template <class T>
-struct TypedAtomicAdd
-{
-    using SRegisters = T[1];
-    using DRegisters = T[1];
+struct TypedAtomicAdd {
+  using SRegisters = T[1];
+  using DRegisters = T[1];
 
-    CUTE_HOST_DEVICE static constexpr void copy(T const& src, T& dst)
-    {
-        atomicAdd(&dst, src);
-    }
+  CUTE_HOST_DEVICE static constexpr void copy(T const& src, T& dst) {
+    atomicAdd(&dst, src);
+  }
 };
 
 template <class T>
-struct Copy_Traits<TypedAtomicAdd<T>>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<TypedAtomicAdd<T>> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<T>::value>>>;
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, Int<sizeof_bits<T>::value>>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<T>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<T>::value>>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
 //////////////////////////////////
 // F16 ADD PTX
 //////////////////////////////////
 
-struct SM70_RED_ADD_NOFTZ_F16
-{
-    using SRegisters = uint16_t[1];
-    using DRegisters = uint16_t[1];
+struct SM70_RED_ADD_NOFTZ_F16 {
+  using SRegisters = uint16_t[1];
+  using DRegisters = uint16_t[1];
 
-    CUTE_HOST_DEVICE static void copy(uint16_t const& src0, uint16_t& gmem_dst)
-    {
+  CUTE_HOST_DEVICE static void copy(uint16_t const& src0, uint16_t& gmem_dst) {
 #if defined(CUTE_ARCH_RED_F16_SM70_ENABLED)
-        asm volatile("red.global.add.noftz.f16 [%0], %1;\n" ::"l"(&gmem_dst), "h"(src0));
+    asm volatile("red.global.add.noftz.f16 [%0], %1;\n" ::"l"(&gmem_dst),
+                 "h"(src0));
 #else
-        CUTE_INVALID_CONTROL_PATH("Trying to use red.global.f16 without CUTE_ARCH_RED_F16_SM70_ENABLED.");
+    CUTE_INVALID_CONTROL_PATH(
+        "Trying to use red.global.f16 without CUTE_ARCH_RED_F16_SM70_ENABLED.");
 #endif
-    }
+  }
 };
 
 template <>
-struct Copy_Traits<SM70_RED_ADD_NOFTZ_F16>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<SM70_RED_ADD_NOFTZ_F16> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, _16>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, _16>>;
 
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, _16>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, _16>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
-struct SM70_RED_ADD_NOFTZ_F16x2
-{
-    using SRegisters = uint32_t[1];
-    using DRegisters = uint32_t[1];
+struct SM70_RED_ADD_NOFTZ_F16x2 {
+  using SRegisters = uint32_t[1];
+  using DRegisters = uint32_t[1];
 
-    CUTE_HOST_DEVICE static void copy(uint32_t const& src0, uint32_t& gmem_dst)
-    {
+  CUTE_HOST_DEVICE static void copy(uint32_t const& src0, uint32_t& gmem_dst) {
 #if defined(CUTE_ARCH_RED_F16_SM70_ENABLED)
-        asm volatile("red.global.add.noftz.f16x2 [%0], %1;\n" ::"l"(&gmem_dst), "r"(src0));
+    asm volatile("red.global.add.noftz.f16x2 [%0], %1;\n" ::"l"(&gmem_dst),
+                 "r"(src0));
 #else
-        CUTE_INVALID_CONTROL_PATH("Trying to use red.global.f16 without CUTE_ARCH_RED_F16_SM70_ENABLED.");
+    CUTE_INVALID_CONTROL_PATH(
+        "Trying to use red.global.f16 without CUTE_ARCH_RED_F16_SM70_ENABLED.");
 #endif
-    }
+  }
 };
 
 template <>
-struct Copy_Traits<SM70_RED_ADD_NOFTZ_F16x2>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<SM70_RED_ADD_NOFTZ_F16x2> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, _32>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, _32>>;
 
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, _32>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, _32>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
-struct SM90_RED_ADD_NOFTZ_F16x2_V2
-{
-    using SRegisters = uint32_t[2];
-    using DRegisters = uint64_t[1];
+struct SM90_RED_ADD_NOFTZ_F16x2_V2 {
+  using SRegisters = uint32_t[2];
+  using DRegisters = uint64_t[1];
 
-    CUTE_HOST_DEVICE static void copy(uint32_t const& src0, uint32_t const& src1, uint64_t& gmem_dst)
-    {
+  CUTE_HOST_DEVICE static void copy(uint32_t const& src0,
+                                    uint32_t const& src1,
+                                    uint64_t& gmem_dst) {
 #if defined(CUTE_ARCH_RED_VEC_SM90_ENABLED)
-        asm volatile("red.global.add.noftz.v2.f16x2 [%0], {%1, %2};\n" ::"l"(&gmem_dst), "r"(src0), "r"(src1));
+    asm volatile(
+        "red.global.add.noftz.v2.f16x2 [%0], {%1, %2};\n" ::"l"(&gmem_dst),
+        "r"(src0),
+        "r"(src1));
 #else
-        CUTE_INVALID_CONTROL_PATH("Trying to use red.global.vX without CUTE_ARCH_RED_VEC_SM90_ENABLED.");
+    CUTE_INVALID_CONTROL_PATH(
+        "Trying to use red.global.vX without CUTE_ARCH_RED_VEC_SM90_ENABLED.");
 #endif
-    }
+  }
 };
 
 template <>
-struct Copy_Traits<SM90_RED_ADD_NOFTZ_F16x2_V2>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<SM90_RED_ADD_NOFTZ_F16x2_V2> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, _64>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, _64>>;
 
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, _64>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, _64>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
-struct SM90_RED_ADD_NOFTZ_F16x2_V4
-{
-    using SRegisters = uint32_t[4];
-    using DRegisters = uint128_t[1];
+struct SM90_RED_ADD_NOFTZ_F16x2_V4 {
+  using SRegisters = uint32_t[4];
+  using DRegisters = uint128_t[1];
 
-    CUTE_HOST_DEVICE static void copy(
-        uint32_t const& src0, uint32_t const& src1, uint32_t const& src2, uint32_t const& src3, uint128_t& gmem_dst)
-    {
+  CUTE_HOST_DEVICE static void copy(uint32_t const& src0,
+                                    uint32_t const& src1,
+                                    uint32_t const& src2,
+                                    uint32_t const& src3,
+                                    uint128_t& gmem_dst) {
 #if defined(CUTE_ARCH_RED_VEC_SM90_ENABLED)
-        asm volatile("red.global.add.noftz.v4.f16x2 [%0], {%1, %2, %3, %4};\n" ::"l"(&gmem_dst), "r"(src0), "r"(src1),
-            "r"(src2), "r"(src3));
+    asm volatile(
+        "red.global.add.noftz.v4.f16x2 [%0], {%1, %2, %3, %4};\n" ::"l"(
+            &gmem_dst),
+        "r"(src0),
+        "r"(src1),
+        "r"(src2),
+        "r"(src3));
 #else
-        CUTE_INVALID_CONTROL_PATH("Trying to use red.global.vX without CUTE_ARCH_RED_VEC_SM90_ENABLED.");
+    CUTE_INVALID_CONTROL_PATH(
+        "Trying to use red.global.vX without CUTE_ARCH_RED_VEC_SM90_ENABLED.");
 #endif
-    }
+  }
 };
 
 template <>
-struct Copy_Traits<SM90_RED_ADD_NOFTZ_F16x2_V4>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<SM90_RED_ADD_NOFTZ_F16x2_V4> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, _128>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, _128>>;
 
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, _128>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, _128>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
 //////////////////////////////////
 // BF16 ADD PTX
 //////////////////////////////////
 
-struct SM90_RED_ADD_NOFTZ_BF16
-{
-    using SRegisters = uint16_t[1];
-    using DRegisters = uint16_t[1];
+struct SM90_RED_ADD_NOFTZ_BF16 {
+  using SRegisters = uint16_t[1];
+  using DRegisters = uint16_t[1];
 
-    CUTE_HOST_DEVICE static void copy(uint16_t const& src0, uint16_t& gmem_dst)
-    {
+  CUTE_HOST_DEVICE static void copy(uint16_t const& src0, uint16_t& gmem_dst) {
 #if defined(CUTE_ARCH_RED_BF16_SM90_ENABLED)
-        asm volatile("red.global.add.noftz.bf16 [%0], %1;\n" ::"l"(&gmem_dst), "h"(src0));
+    asm volatile("red.global.add.noftz.bf16 [%0], %1;\n" ::"l"(&gmem_dst),
+                 "h"(src0));
 #else
-        CUTE_INVALID_CONTROL_PATH("Trying to use red.global.bf16 without CUTE_ARCH_RED_BF16_SM90_ENABLED.");
+    CUTE_INVALID_CONTROL_PATH(
+        "Trying to use red.global.bf16 without "
+        "CUTE_ARCH_RED_BF16_SM90_ENABLED.");
 #endif
-    }
+  }
 };
 
 template <>
-struct Copy_Traits<SM90_RED_ADD_NOFTZ_BF16>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<SM90_RED_ADD_NOFTZ_BF16> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, _16>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, _16>>;
 
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, _16>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, _16>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
 //////////////////////////////////
 
-struct SM90_RED_ADD_NOFTZ_BF16x2
-{
-    using SRegisters = uint32_t[1];
-    using DRegisters = uint32_t[1];
+struct SM90_RED_ADD_NOFTZ_BF16x2 {
+  using SRegisters = uint32_t[1];
+  using DRegisters = uint32_t[1];
 
-    CUTE_HOST_DEVICE static void copy(uint32_t const& src0, uint32_t& gmem_dst)
-    {
+  CUTE_HOST_DEVICE static void copy(uint32_t const& src0, uint32_t& gmem_dst) {
 #if defined(CUTE_ARCH_RED_BF16_SM90_ENABLED)
-        asm volatile("red.global.add.noftz.bf16x2 [%0], %1;\n" ::"l"(&gmem_dst), "r"(src0));
+    asm volatile("red.global.add.noftz.bf16x2 [%0], %1;\n" ::"l"(&gmem_dst),
+                 "r"(src0));
 #else
-        CUTE_INVALID_CONTROL_PATH("Trying to use red.global.bf16 without CUTE_ARCH_RED_BF16_SM90_ENABLED.");
+    CUTE_INVALID_CONTROL_PATH(
+        "Trying to use red.global.bf16 without "
+        "CUTE_ARCH_RED_BF16_SM90_ENABLED.");
 #endif
-    }
+  }
 };
 
 template <>
-struct Copy_Traits<SM90_RED_ADD_NOFTZ_BF16x2>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<SM90_RED_ADD_NOFTZ_BF16x2> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, _32>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, _32>>;
 
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, _32>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, _32>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
 //////////////////////////////////
 
-struct SM90_RED_ADD_NOFTZ_BF16x2_V2
-{
-    using SRegisters = uint32_t[2];
-    using DRegisters = uint64_t[1];
+struct SM90_RED_ADD_NOFTZ_BF16x2_V2 {
+  using SRegisters = uint32_t[2];
+  using DRegisters = uint64_t[1];
 
-    CUTE_HOST_DEVICE static void copy(uint32_t const& src0, uint32_t const& src1, uint64_t& gmem_dst)
-    {
+  CUTE_HOST_DEVICE static void copy(uint32_t const& src0,
+                                    uint32_t const& src1,
+                                    uint64_t& gmem_dst) {
 #if defined(CUTE_ARCH_RED_BF16_SM90_ENABLED)
-        asm volatile("red.global.add.noftz.v2.bf16x2 [%0], {%1, %2};\n" ::"l"(&gmem_dst), "r"(src0), "r"(src1));
+    asm volatile(
+        "red.global.add.noftz.v2.bf16x2 [%0], {%1, %2};\n" ::"l"(&gmem_dst),
+        "r"(src0),
+        "r"(src1));
 #else
-        CUTE_INVALID_CONTROL_PATH("Trying to use red.global.bf16 without CUTE_ARCH_RED_BF16_SM90_ENABLED.");
+    CUTE_INVALID_CONTROL_PATH(
+        "Trying to use red.global.bf16 without "
+        "CUTE_ARCH_RED_BF16_SM90_ENABLED.");
 #endif
-    }
+  }
 };
 
 template <>
-struct Copy_Traits<SM90_RED_ADD_NOFTZ_BF16x2_V2>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<SM90_RED_ADD_NOFTZ_BF16x2_V2> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, _64>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, _64>>;
 
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, _64>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, _64>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
 //////////////////////////////////
 
-struct SM90_RED_ADD_NOFTZ_BF16x2_V4
-{
-    using SRegisters = uint32_t[4];
-    using DRegisters = uint128_t[1];
+struct SM90_RED_ADD_NOFTZ_BF16x2_V4 {
+  using SRegisters = uint32_t[4];
+  using DRegisters = uint128_t[1];
 
-    CUTE_HOST_DEVICE static void copy(
-        uint32_t const& src0, uint32_t const& src1, uint32_t const& src2, uint32_t const& src3, uint128_t& gmem_dst)
-    {
+  CUTE_HOST_DEVICE static void copy(uint32_t const& src0,
+                                    uint32_t const& src1,
+                                    uint32_t const& src2,
+                                    uint32_t const& src3,
+                                    uint128_t& gmem_dst) {
 #if defined(CUTE_ARCH_RED_BF16_SM90_ENABLED)
-        asm volatile("red.global.add.noftz.v4.bf16x2 [%0], {%1, %2, %3, %4};\n" ::"l"(&gmem_dst), "r"(src0), "r"(src1),
-            "r"(src2), "r"(src3));
+    asm volatile(
+        "red.global.add.noftz.v4.bf16x2 [%0], {%1, %2, %3, %4};\n" ::"l"(
+            &gmem_dst),
+        "r"(src0),
+        "r"(src1),
+        "r"(src2),
+        "r"(src3));
 #else
-        CUTE_INVALID_CONTROL_PATH("Trying to use red.global.bf16 without CUTE_ARCH_RED_BF16_SM90_ENABLED.");
+    CUTE_INVALID_CONTROL_PATH(
+        "Trying to use red.global.bf16 without "
+        "CUTE_ARCH_RED_BF16_SM90_ENABLED.");
 #endif
-    }
+  }
 };
 
 template <>
-struct Copy_Traits<SM90_RED_ADD_NOFTZ_BF16x2_V4>
-{
-    // Logical thread id to thread idx (one-thread)
-    using ThrID = Layout<_1>;
+struct Copy_Traits<SM90_RED_ADD_NOFTZ_BF16x2_V4> {
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
 
-    // Map from (src-thr,src-val) to bit
-    using SrcLayout = Layout<Shape<_1, _128>>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, _128>>;
 
-    // Map from (dst-thr,dst-val) to bit
-    using DstLayout = Layout<Shape<_1, _128>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, _128>>;
 
-    // Reference map from (thr,val) to bit
-    using RefLayout = SrcLayout;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
 };
 
 //////////////////////////////////
 
-} // end namespace cute
+}  // end namespace cute
diff --git a/custom_ops/gpu_ops/cutlass_extensions/arch/memory_copy_sm80.h b/custom_ops/gpu_ops/cutlass_extensions/arch/memory_copy_sm80.h
index a9975c0138..0e9fa1b151 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/arch/memory_copy_sm80.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/arch/memory_copy_sm80.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
@@ -59,8 +60,9 @@ template <
     bool GlobalToShared = true>
 struct copy;
 
-/// Initiates an asynchronous copy from global memory to shared memory. Rather than predicate
-/// the entire transfer, zeros are written to SMEM if the guard predicate is false.
+/// Initiates an asynchronous copy from global memory to shared memory. Rather
+/// than predicate the entire transfer, zeros are written to SMEM if the guard
+/// predicate is false.
 ///
 /// cp.async
 ///
@@ -72,7 +74,8 @@ template <
     bool GlobalToShared = true>
 struct copy_zfill;
 
-/// Blocks until all but <N> previous cp.async.commit_group operations have committed.
+/// Blocks until all but <N> previous cp.async.commit_group operations have
+/// committed.
 ///
 /// cp.async
 ///
@@ -86,11 +89,11 @@ template <
     /// Size of the access in bytes
     int SizeInBytes>
 struct copy<SizeInBytes, CacheOperation::Always, true> {
-
   /// Copy
   CUTLASS_DEVICE
   copy(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
-      cp_async<SizeInBytes, CacheOperation::Always>(smem_ptr, global_ptr, pred_guard);
+    cp_async<SizeInBytes, CacheOperation::Always>(
+        smem_ptr, global_ptr, pred_guard);
   }
 };
 
@@ -99,15 +102,15 @@ template <
     /// Size of the access in bytes
     int SizeInBytes>
 struct copy<SizeInBytes, CacheOperation::Always, false> {
-
   /// Copy
   CUTLASS_DEVICE
   copy(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
-      using AccessType  = Array<uint8_t, SizeInBytes>;
+    using AccessType = Array<uint8_t, SizeInBytes>;
 
-      if (pred_guard) {
-        *static_cast<AccessType *>(smem_ptr) = *static_cast<AccessType const *>(global_ptr);
-      }
+    if (pred_guard) {
+      *static_cast<AccessType *>(smem_ptr) =
+          *static_cast<AccessType const *>(global_ptr);
+    }
   }
 };
 
@@ -116,11 +119,11 @@ template <
     /// Size of the access in bytes
     int SizeInBytes>
 struct copy_zfill<SizeInBytes, CacheOperation::Always, true> {
-
   /// Copy with zero fill
   CUTLASS_DEVICE
   copy_zfill(void *smem_ptr, void const *global_ptr, bool pred_guard) {
-      cp_async_zfill<SizeInBytes, CacheOperation::Always>(smem_ptr, global_ptr, pred_guard);
+    cp_async_zfill<SizeInBytes, CacheOperation::Always>(
+        smem_ptr, global_ptr, pred_guard);
   }
 };
 
@@ -129,20 +132,19 @@ template <
     /// Size of the access in bytes
     int SizeInBytes>
 struct copy_zfill<SizeInBytes, CacheOperation::Always, false> {
-
   /// Copy with zero fill
   CUTLASS_DEVICE
   copy_zfill(void *smem_ptr, void const *global_ptr, bool pred_guard) {
-      using AccessType  = Array<uint8_t, SizeInBytes>;
+    using AccessType = Array<uint8_t, SizeInBytes>;
 
-      if (pred_guard) {
-        *static_cast<AccessType *>(smem_ptr) = *static_cast<AccessType const *>(global_ptr);
-      }
-      else {
-        AccessType zeros;
-        zeros.clear();
-        *static_cast<AccessType *>(smem_ptr) = zeros;
-      }
+    if (pred_guard) {
+      *static_cast<AccessType *>(smem_ptr) =
+          *static_cast<AccessType const *>(global_ptr);
+    } else {
+      AccessType zeros;
+      zeros.clear();
+      *static_cast<AccessType *>(smem_ptr) = zeros;
+    }
   }
 };
 
@@ -153,11 +155,11 @@ template <
     /// Size of the access in bytes
     int SizeInBytes>
 struct copy<SizeInBytes, CacheOperation::Global, true> {
-
   /// Copy
   CUTLASS_DEVICE
   copy(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
-      cp_async<SizeInBytes, CacheOperation::Global>(smem_ptr, global_ptr, pred_guard);
+    cp_async<SizeInBytes, CacheOperation::Global>(
+        smem_ptr, global_ptr, pred_guard);
   }
 };
 
@@ -166,15 +168,15 @@ template <
     /// Size of the access in bytes
     int SizeInBytes>
 struct copy<SizeInBytes, CacheOperation::Global, false> {
-
   /// Copy
   CUTLASS_DEVICE
   copy(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
-      using AccessType  = Array<uint8_t, SizeInBytes>;
+    using AccessType = Array<uint8_t, SizeInBytes>;
 
-      if (pred_guard) {
-        *static_cast<AccessType *>(smem_ptr) = *static_cast<AccessType const *>(global_ptr);
-      }
+    if (pred_guard) {
+      *static_cast<AccessType *>(smem_ptr) =
+          *static_cast<AccessType const *>(global_ptr);
+    }
   }
 };
 
@@ -183,11 +185,11 @@ template <
     /// Size of the access in bytes
     int SizeInBytes>
 struct copy_zfill<SizeInBytes, CacheOperation::Global, true> {
-
   /// Copy with zero fill
   CUTLASS_DEVICE
   copy_zfill(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
-      cp_async_zfill<SizeInBytes, CacheOperation::Global>(smem_ptr, global_ptr, pred_guard);
+    cp_async_zfill<SizeInBytes, CacheOperation::Global>(
+        smem_ptr, global_ptr, pred_guard);
   }
 };
 
@@ -196,31 +198,29 @@ template <
     /// Size of the access in bytes
     int SizeInBytes>
 struct copy_zfill<SizeInBytes, CacheOperation::Global, false> {
-
   /// Copy with zero fill
   CUTLASS_DEVICE
   copy_zfill(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
-      using AccessType  = Array<uint8_t, SizeInBytes>;
+    using AccessType = Array<uint8_t, SizeInBytes>;
 
-      if (pred_guard) {
-        *static_cast<AccessType *>(smem_ptr) = *static_cast<AccessType const *>(global_ptr);
-      }
-      else {
-        AccessType zeros;
-        zeros.clear();
-        *static_cast<AccessType *>(smem_ptr) = zeros;
-      }
+    if (pred_guard) {
+      *static_cast<AccessType *>(smem_ptr) =
+          *static_cast<AccessType const *>(global_ptr);
+    } else {
+      AccessType zeros;
+      zeros.clear();
+      *static_cast<AccessType *>(smem_ptr) = zeros;
+    }
   }
 };
 
-/// Establishes an ordering w.r.t previously issued cp.async instructions. Does not block.
+/// Establishes an ordering w.r.t previously issued cp.async instructions. Does
+/// not block.
 template <bool GlobalToShared>
-CUTLASS_DEVICE
-void copy_fence() {}
+CUTLASS_DEVICE void copy_fence() {}
 
 template <>
-CUTLASS_DEVICE
-void copy_fence<true>() {
+CUTLASS_DEVICE void copy_fence<true>() {
   cp_async_fence();
 }
 
@@ -229,7 +229,6 @@ void copy_fence<true>() {
 /// Partial specialization
 template <int N>
 struct copy_wait<N, false> {
-
   CUTLASS_DEVICE
   copy_wait() {}
 };
@@ -237,7 +236,6 @@ struct copy_wait<N, false> {
 /// Partial specialization
 template <int N>
 struct copy_wait<N, true> {
-
   CUTLASS_DEVICE
   copy_wait() { cp_async_wait<N>(); }
 };
diff --git a/custom_ops/gpu_ops/cutlass_extensions/arch/mma.h b/custom_ops/gpu_ops/cutlass_extensions/arch/mma.h
index 2362da4f7f..2ab2981518 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/arch/mma.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/arch/mma.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -37,10 +38,8 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace arch
-{
+namespace cutlass {
+namespace arch {
 
 // Tag which triggers MMA which will trigger
 struct OpMultiplyAddDequantizeInterleavedBToA;
@@ -52,8 +51,8 @@ struct OpMultiplyAddDequantizeInterleavedBToA;
   split out the template below into OpMultiplyAddDequantizeInterleavedBToA along
   with the quantization op before instantiating the GEMM pieces.
 
-  Note that this is somewhat of a hack, but it SIGNIFICANTLY reduces the amount of
-  code we need to duplicate.
+  Note that this is somewhat of a hack, but it SIGNIFICANTLY reduces the amount
+  of code we need to duplicate.
  */
 struct OpMultiplyAddDequantizeInterleavedBToA_percol_scale;
 struct OpMultiplyAddDequantizeInterleavedBToA_fine_scale;
@@ -61,60 +60,59 @@ struct OpMultiplyAddDequantizeInterleavedBToA_fine_scalebias;
 
 // The default just forwards the original operator
 template <typename MmaOp, WeightOnlyQuantOp QuantOp_>
-struct TagOperator
-{
-    using TaggedOperator = MmaOp;
+struct TagOperator {
+  using TaggedOperator = MmaOp;
 };
 
 // Specializations below attach more information to the operator
 template <>
-struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA, WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>
-{
-    using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_percol_scale;
+struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA,
+                   WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY> {
+  using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_percol_scale;
 };
 
 template <>
-struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA, WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY>
-{
-    using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_fine_scale;
+struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA,
+                   WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY> {
+  using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_fine_scale;
 };
 
 template <>
-struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA, WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS>
-{
-    using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_fine_scalebias;
+struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA,
+                   WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS> {
+  using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_fine_scalebias;
 };
 
-// Here we instantiate some structs to "detag" the tagged operator. It splits it back to the original
-// operator + the extra information. If no extra info was tagged, the dequant op per column scaling
-// as a default.
+// Here we instantiate some structs to "detag" the tagged operator. It splits it
+// back to the original operator + the extra information. If no extra info was
+// tagged, the dequant op per column scaling as a default.
 template <typename TaggedMmaOp>
-struct DetagOperator
-{
-    using Operator = TaggedMmaOp;
-    static constexpr WeightOnlyQuantOp QuantOp = WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY;
+struct DetagOperator {
+  using Operator = TaggedMmaOp;
+  static constexpr WeightOnlyQuantOp QuantOp =
+      WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY;
 };
 
 template <>
-struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_percol_scale>
-{
-    using Operator = OpMultiplyAddDequantizeInterleavedBToA;
-    static constexpr WeightOnlyQuantOp QuantOp = WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY;
+struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_percol_scale> {
+  using Operator = OpMultiplyAddDequantizeInterleavedBToA;
+  static constexpr WeightOnlyQuantOp QuantOp =
+      WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY;
 };
 
 template <>
-struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_fine_scale>
-{
-    using Operator = OpMultiplyAddDequantizeInterleavedBToA;
-    static constexpr WeightOnlyQuantOp QuantOp = WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY;
+struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_fine_scale> {
+  using Operator = OpMultiplyAddDequantizeInterleavedBToA;
+  static constexpr WeightOnlyQuantOp QuantOp =
+      WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY;
 };
 
 template <>
-struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_fine_scalebias>
-{
-    using Operator = OpMultiplyAddDequantizeInterleavedBToA;
-    static constexpr WeightOnlyQuantOp QuantOp = WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS;
+struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_fine_scalebias> {
+  using Operator = OpMultiplyAddDequantizeInterleavedBToA;
+  static constexpr WeightOnlyQuantOp QuantOp =
+      WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS;
 };
 
-} // namespace arch
-} // namespace cutlass
+}  // namespace arch
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/compute_occupancy.h b/custom_ops/gpu_ops/cutlass_extensions/compute_occupancy.h
index 29ee976691..7f7cce414e 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/compute_occupancy.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/compute_occupancy.h
@@ -20,66 +20,65 @@
 #include "cutlass/device_kernel.h"
 #include "common/cudaUtils.h"
 
-namespace cutlass_extensions
-{
+namespace cutlass_extensions {
 
 template <typename GemmKernel, bool enable_cutlass_3x = false>
-inline int compute_occupancy_for_kernel()
-{
+inline int compute_occupancy_for_kernel() {
+  int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
 
-    int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
-
-    if (smem_size > (48 << 10))
-    {
-        cudaFuncAttributes attr;
-        int device = 0;
-        int max_smem_per_block = 0;
-        PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&device));
-        PADDLE_ENFORCE_GPU_SUCCESS(
-            cudaDeviceGetAttribute(&max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device));
-        if constexpr (enable_cutlass_3x)
-        {
-            PADDLE_ENFORCE_GPU_SUCCESS(cudaFuncGetAttributes(&attr, cutlass::device_kernel<GemmKernel>));
-        }
-        else
-        {
-            PADDLE_ENFORCE_GPU_SUCCESS(cudaFuncGetAttributes(&attr, cutlass::Kernel<GemmKernel>));
-        }
-        if (smem_size + attr.sharedSizeBytes >= static_cast<size_t>(max_smem_per_block))
-        {
-            // This should mean that
-            // cudaFuncSetAttribute(cutlass::Kernel<GemmKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)
-            // wouldn't work. In that case, we return an occupancy of 0. This will cause the heuristic to ignore this
-            // configuration.
-            return 0;
-        }
-
-        if constexpr (enable_cutlass_3x)
-        {
-            PADDLE_ENFORCE_GPU_SUCCESS(cudaFuncSetAttribute(
-                cutlass::device_kernel<GemmKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
-        }
-        else
-        {
-            PADDLE_ENFORCE_GPU_SUCCESS(cudaFuncSetAttribute(
-                cutlass::Kernel<GemmKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
-        }
+  if (smem_size > (48 << 10)) {
+    cudaFuncAttributes attr;
+    int device = 0;
+    int max_smem_per_block = 0;
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&device));
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(
+        &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device));
+    if constexpr (enable_cutlass_3x) {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaFuncGetAttributes(&attr, cutlass::device_kernel<GemmKernel>));
+    } else {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaFuncGetAttributes(&attr, cutlass::Kernel<GemmKernel>));
+    }
+    if (smem_size + attr.sharedSizeBytes >=
+        static_cast<size_t>(max_smem_per_block)) {
+      // This should mean that
+      // cudaFuncSetAttribute(cutlass::Kernel<GemmKernel>,
+      // cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) wouldn't work.
+      // In that case, we return an occupancy of 0. This will cause the
+      // heuristic to ignore this configuration.
+      return 0;
     }
 
-    int max_active_blocks = -1;
-    if constexpr (enable_cutlass_3x)
-    {
-        PADDLE_ENFORCE_GPU_SUCCESS(
-            cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_active_blocks, cutlass::device_kernel<GemmKernel>,
-                128 * (GemmKernel::NumLoadWarpGroups + GemmKernel::NumMmaWarpGroups), smem_size));
-    }
-    else
-    {
-        PADDLE_ENFORCE_GPU_SUCCESS(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-            &max_active_blocks, cutlass::Kernel<GemmKernel>, GemmKernel::kThreadCount, smem_size));
+    if constexpr (enable_cutlass_3x) {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaFuncSetAttribute(cutlass::device_kernel<GemmKernel>,
+                               cudaFuncAttributeMaxDynamicSharedMemorySize,
+                               smem_size));
+    } else {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaFuncSetAttribute(cutlass::Kernel<GemmKernel>,
+                               cudaFuncAttributeMaxDynamicSharedMemorySize,
+                               smem_size));
     }
+  }
 
-    return max_active_blocks;
+  int max_active_blocks = -1;
+  if constexpr (enable_cutlass_3x) {
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &max_active_blocks,
+        cutlass::device_kernel<GemmKernel>,
+        128 * (GemmKernel::NumLoadWarpGroups + GemmKernel::NumMmaWarpGroups),
+        smem_size));
+  } else {
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &max_active_blocks,
+        cutlass::Kernel<GemmKernel>,
+        GemmKernel::kThreadCount,
+        smem_size));
+  }
+
+  return max_active_blocks;
 }
 
-} // namespace cutlass_extensions
+}  // namespace cutlass_extensions
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp b/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp
index 3e5aa4b038..60c0dc6302 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp
@@ -47,7 +47,8 @@
 // breaks when moving scales to the CPU.
 //
 
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp
 
 #pragma once
 
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp b/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp
index fa1df1fb1e..dbc6a718f3 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp
@@ -47,7 +47,8 @@
 // breaks when moving scales to the CPU.
 //
 
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp
 
 #pragma once
 
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp b/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
index 7b56c3c1ac..aee2fbaeb2 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
@@ -47,7 +47,8 @@
 // breaks when moving scales to the CPU.
 //
 
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
 
 #pragma once
 
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp b/custom_ops/gpu_ops/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
index 513d3741fb..42db1bd158 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
 
 #pragma once
 
@@ -24,31 +25,41 @@ using namespace cute;
  */
 template <typename ElementD, typename OutputTileThreadMap>
 struct ScaledEpilogueBase {
-protected:
+ protected:
   using Accum = cutlass::epilogue::threadblock::VisitorAccFetch;
 
   template <typename T>
   using ColOrScalarLoad =
       cutlass::epilogue::threadblock::VisitorColOrScalarBroadcast<
-          OutputTileThreadMap, T, Stride<Int<1>, Int<0>, Int<0>>>;
+          OutputTileThreadMap,
+          T,
+          Stride<Int<1>, Int<0>, Int<0>>>;
 
   template <typename T>
   using RowOrScalarLoad =
       cutlass::epilogue::threadblock::VisitorRowOrScalarBroadcast<
-          OutputTileThreadMap, T, Stride<Int<0>, Int<1>, Int<0>>>;
+          OutputTileThreadMap,
+          T,
+          Stride<Int<0>, Int<1>, Int<0>>>;
 
   template <typename T>
   using ColLoad = cutlass::epilogue::threadblock::VisitorColBroadcast<
-      OutputTileThreadMap, T, Stride<Int<1>, Int<0>, Int<0>>>;
+      OutputTileThreadMap,
+      T,
+      Stride<Int<1>, Int<0>, Int<0>>>;
 
   template <typename T>
   using RowLoad = cutlass::epilogue::threadblock::VisitorRowBroadcast<
-      OutputTileThreadMap, T, Stride<Int<0>, Int<1>, Int<0>>>;
+      OutputTileThreadMap,
+      T,
+      Stride<Int<0>, Int<1>, Int<0>>>;
 
   template <typename T>
   using RowOrZeroLoad =
       cutlass::epilogue::threadblock::VisitorRowOrZeroBroadcast<
-          OutputTileThreadMap, T, Stride<Int<0>, Int<1>, Int<0>>>;
+          OutputTileThreadMap,
+          T,
+          Stride<Int<0>, Int<1>, Int<0>>>;
 
   // This utility function constructs the arguments for the load descriptors
   // from a tensor. It can handle both row and column, as well as row/column or
@@ -56,15 +67,11 @@ protected:
   template <typename Descriptor, typename T>
   static auto args_from_tensor(paddle::Tensor const &tensor) {
     using Arguments = typename Descriptor::Arguments;
-    auto *data_ptr = static_cast<T *>(const_cast<void *>(
-        tensor.data()));
-     if constexpr (std::is_same_v<Descriptor,
-                                                     ColOrScalarLoad<T>> ||
-                                      std::is_same_v<Descriptor,
-                                                     RowOrScalarLoad<T>>) {
+    auto *data_ptr = static_cast<T *>(const_cast<void *>(tensor.data()));
+    if constexpr (std::is_same_v<Descriptor, ColOrScalarLoad<T>> ||
+                  std::is_same_v<Descriptor, RowOrScalarLoad<T>>) {
       return Arguments{data_ptr, tensor.numel() != 1};
-    }
-    else {
+    } else {
       // it would technically work but no use case as data_ptr is never nullptr
       static_assert(!std::is_same_v<Descriptor, RowOrZeroLoad<T>>);
       return Arguments{data_ptr};
@@ -102,24 +109,28 @@ protected:
 template <typename ElementD, typename OutputTileThreadMap>
 struct ScaledEpilogue
     : private ScaledEpilogueBase<ElementD, OutputTileThreadMap> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementD, OutputTileThreadMap>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
   using ScaleB = typename SUPER::template RowOrScalarLoad<float>;
 
   using Compute0 = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTCompute0 =
       cutlass::epilogue::threadblock::Sm80EVT<Compute0, ScaleB, Accum>;
 
   using Compute1 = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiplies, ElementD, float,
+      cutlass::multiplies,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
+ public:
   using EVTCompute =
       cutlass::epilogue::threadblock::Sm80EVT<Compute1, ScaleA, EVTCompute0>;
   using ArgumentType = typename EVTCompute::Arguments;
@@ -146,26 +157,30 @@ public:
 template <typename ElementD, typename OutputTileThreadMap>
 struct ScaledEpilogueBias
     : protected ScaledEpilogueBase<ElementD, OutputTileThreadMap> {
-protected:
+ protected:
   using SUPER = ScaledEpilogueBase<ElementD, OutputTileThreadMap>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
   using ScaleB = typename SUPER::template RowOrScalarLoad<float>;
   using Bias = typename SUPER::template RowLoad<ElementD>;
   using Compute0 = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTCompute0 =
       cutlass::epilogue::threadblock::Sm80EVT<Compute0, ScaleB, Accum>;
 
   using Compute1 = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiply_add, ElementD, float,
+      cutlass::multiply_add,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
-  using EVTCompute = cutlass::epilogue::threadblock::Sm80EVT<Compute1, ScaleA,
-                                                             EVTCompute0, Bias>;
+ public:
+  using EVTCompute = cutlass::epilogue::threadblock::
+      Sm80EVT<Compute1, ScaleA, EVTCompute0, Bias>;
   using ArgumentType = typename EVTCompute::Arguments;
   static ArgumentType prepare_args(paddle::Tensor const &a_scales,
                                    paddle::Tensor const &b_scales,
@@ -190,7 +205,7 @@ public:
 template <typename ElementD, typename OutputTileThreadMap>
 struct ScaledEpilogueBiasAzp
     : protected ScaledEpilogueBase<ElementD, OutputTileThreadMap> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementD, OutputTileThreadMap>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
@@ -202,35 +217,40 @@ private:
 
   // Compute float(accum - azp_adj), both operands are int32_t
   using ComputeAzp = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::minus, float, int32_t,
+      cutlass::minus,
+      float,
+      int32_t,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTComputeAzp =
       cutlass::epilogue::threadblock::Sm80EVT<ComputeAzp, Accum, AzpWithAdj>;
 
   using ComputeScaleB = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-  using EVTComputeScaleB =
-      cutlass::epilogue::threadblock::Sm80EVT<ComputeScaleB, ScaleB,
-                                              EVTComputeAzp>;
+  using EVTComputeScaleB = cutlass::epilogue::threadblock::
+      Sm80EVT<ComputeScaleB, ScaleB, EVTComputeAzp>;
 
   using ComputeScaleBiasA = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiply_add, ElementD, float,
+      cutlass::multiply_add,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
-  using EVTCompute =
-      cutlass::epilogue::threadblock::Sm80EVT<ComputeScaleBiasA, ScaleA,
-                                              EVTComputeScaleB, Bias>;
+ public:
+  using EVTCompute = cutlass::epilogue::threadblock::
+      Sm80EVT<ComputeScaleBiasA, ScaleA, EVTComputeScaleB, Bias>;
 
   using ArgumentType = typename EVTCompute::Arguments;
 
-  static ArgumentType
-  prepare_args(paddle::Tensor const &a_scales, paddle::Tensor const &b_scales,
-               paddle::Tensor const &azp_adj,
-               paddle::optional<paddle::Tensor> const &bias) {
+  static ArgumentType prepare_args(
+      paddle::Tensor const &a_scales,
+      paddle::Tensor const &b_scales,
+      paddle::Tensor const &azp_adj,
+      paddle::optional<paddle::Tensor> const &bias) {
     auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
     auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
     auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -257,7 +277,7 @@ public:
 template <typename ElementD, typename OutputTileThreadMap>
 struct ScaledEpilogueBiasAzpToken
     : protected ScaledEpilogueBase<ElementD, OutputTileThreadMap> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementD, OutputTileThreadMap>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
@@ -272,7 +292,9 @@ private:
 
   // Compute azp * azp_adj
   using ComputeAzp = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiplies, int32_t, int32_t,
+      cutlass::multiplies,
+      int32_t,
+      int32_t,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTComputeAzp =
@@ -280,35 +302,41 @@ private:
 
   // Compute float(accum - azp*azp_adj), all operands are int32_t
   using ComputeAcc = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::minus, float, int32_t,
+      cutlass::minus,
+      float,
+      int32_t,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTComputeAcc =
       cutlass::epilogue::threadblock::Sm80EVT<ComputeAcc, Accum, EVTComputeAzp>;
 
   using ComputeScaleB = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-  using EVTComputeScaleB =
-      cutlass::epilogue::threadblock::Sm80EVT<ComputeScaleB, ScaleB,
-                                              EVTComputeAcc>;
+  using EVTComputeScaleB = cutlass::epilogue::threadblock::
+      Sm80EVT<ComputeScaleB, ScaleB, EVTComputeAcc>;
 
   using ComputeScaleBiasA = cutlass::epilogue::threadblock::VisitorCompute<
-      cutlass::multiply_add, ElementD, float,
+      cutlass::multiply_add,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
-  using EVTCompute =
-      cutlass::epilogue::threadblock::Sm80EVT<ComputeScaleBiasA, ScaleA,
-                                              EVTComputeScaleB, Bias>;
+ public:
+  using EVTCompute = cutlass::epilogue::threadblock::
+      Sm80EVT<ComputeScaleBiasA, ScaleA, EVTComputeScaleB, Bias>;
 
   using ArgumentType = typename EVTCompute::Arguments;
 
-  static ArgumentType
-  prepare_args(paddle::Tensor const &a_scales, paddle::Tensor const &b_scales,
-               paddle::Tensor const &azp_adj, paddle::Tensor const &azp,
-               paddle::optional<paddle::Tensor> const &bias) {
+  static ArgumentType prepare_args(
+      paddle::Tensor const &a_scales,
+      paddle::Tensor const &b_scales,
+      paddle::Tensor const &azp_adj,
+      paddle::Tensor const &azp,
+      paddle::optional<paddle::Tensor> const &bias) {
     auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
     auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
     auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -324,4 +352,4 @@ public:
   }
 };
 
-}; // namespace fastdeploy::c2x
+};  // namespace fastdeploy::c2x
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp b/custom_ops/gpu_ops/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
index 90a5e342b5..abb73ce84c 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
 
 #pragma once
 
@@ -24,24 +25,28 @@ namespace fastdeploy::c3x {
 
 using namespace cute;
 
-template <typename T> struct identity {
+template <typename T>
+struct identity {
   CUTLASS_HOST_DEVICE
   T operator()(T lhs) const { return lhs; }
 };
 
 template <typename ElementAcc, typename ElementD, typename TileShape>
 struct TrivialEpilogue {
-private:
+ private:
   using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
   using Compute = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::epilogue::thread::Identity, ElementD, ElementAcc,
+      cutlass::epilogue::thread::Identity,
+      ElementD,
+      ElementAcc,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
+ public:
   using EVTCompute = cutlass::epilogue::fusion::Sm90EVT<Compute, Accum>;
   using ArgumentType = typename EVTCompute::Arguments;
 
-  template <typename... Args> static ArgumentType prepare_args(Args... args) {
+  template <typename... Args>
+  static ArgumentType prepare_args(Args... args) {
     return {};
   }
 };
@@ -52,38 +57,60 @@ public:
  */
 template <typename ElementAcc, typename ElementD, typename TileShape>
 struct ScaledEpilogueBase {
-protected:
+ protected:
   using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
 
   template <typename T>
   using ColOrScalarLoad = cutlass::epilogue::fusion::Sm90ColOrScalarBroadcast<
-      0 /*Stages*/, TileShape, T, Stride<Int<1>, Int<0>, Int<0>>>;
+      0 /*Stages*/,
+      TileShape,
+      T,
+      Stride<Int<1>, Int<0>, Int<0>>>;
 
   template <typename T>
   using RowOrScalarLoad = cutlass::epilogue::fusion::Sm90RowOrScalarBroadcast<
-      0 /*Stages*/, TileShape, T, Stride<Int<0>, Int<1>, Int<0>>>;
+      0 /*Stages*/,
+      TileShape,
+      T,
+      Stride<Int<0>, Int<1>, Int<0>>>;
 
   // Don't want to support nullptr by default
   template <typename T, bool EnableNullPtr = false>
   using ColLoad = cutlass::epilogue::fusion::Sm90ColBroadcast<
-      0 /*Stages*/, TileShape, T, T, Stride<Int<1>, Int<0>, Int<0>>,
-      128 / sizeof_bits_v<T>, EnableNullPtr>;
+      0 /*Stages*/,
+      TileShape,
+      T,
+      T,
+      Stride<Int<1>, Int<0>, Int<0>>,
+      128 / sizeof_bits_v<T>,
+      EnableNullPtr>;
 
   // Don't want to support nullptr by default
   template <typename T, bool EnableNullPtr = false>
   using RowLoad = cutlass::epilogue::fusion::Sm90RowBroadcast<
-      0 /*Stages*/, TileShape, T, T, Stride<Int<0>, Int<1>, Int<0>>,
-      128 / sizeof_bits_v<T>, EnableNullPtr>;
+      0 /*Stages*/,
+      TileShape,
+      T,
+      T,
+      Stride<Int<0>, Int<1>, Int<0>>,
+      128 / sizeof_bits_v<T>,
+      EnableNullPtr>;
 
   template <typename T>
   using ColOrScalarLoadArray =
       cutlass::epilogue::fusion::Sm90ColOrScalarBroadcastArray<
-          0 /*Stages*/, TileShape, T, Stride<Int<1>, Int<0>, Int<0>>>;
+          0 /*Stages*/,
+          TileShape,
+          T,
+          Stride<Int<1>, Int<0>, Int<0>>>;
 
   template <typename T>
   using RowOrScalarLoadArray =
       cutlass::epilogue::fusion::Sm90RowOrScalarBroadcastArray<
-          0 /*Stages*/, TileShape, T, Stride<Int<0>, Int<1>, Int<0>>>;
+          0 /*Stages*/,
+          TileShape,
+          T,
+          Stride<Int<0>, Int<1>, Int<0>>>;
 
   // This utility function constructs the arguments for the load descriptors
   // from a tensor. It can handle both row and column, as well as row/column or
@@ -142,24 +169,28 @@ protected:
 template <typename ElementAcc, typename ElementD, typename TileShape>
 struct ScaledEpilogue
     : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
   using ScaleB = typename SUPER::template RowOrScalarLoad<float>;
 
   using Compute0 = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTCompute0 =
       cutlass::epilogue::fusion::Sm90EVT<Compute0, ScaleB, Accum>;
 
   using Compute1 = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, ElementD, float,
+      cutlass::multiplies,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
+ public:
   using EVTCompute =
       cutlass::epilogue::fusion::Sm90EVT<Compute1, ScaleA, EVTCompute0>;
   using ArgumentType = typename EVTCompute::Arguments;
@@ -186,7 +217,7 @@ public:
 template <typename ElementAcc, typename ElementD, typename TileShape>
 struct ScaledEpilogueBias
     : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
@@ -194,17 +225,21 @@ private:
   using Bias = typename SUPER::template RowLoad<ElementD>;
 
   using Compute0 = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTCompute0 =
       cutlass::epilogue::fusion::Sm90EVT<Compute0, ScaleB, Accum>;
 
   using Compute1 = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiply_add, ElementD, float,
+      cutlass::multiply_add,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
+ public:
   using EVTCompute =
       cutlass::epilogue::fusion::Sm90EVT<Compute1, ScaleA, EVTCompute0, Bias>;
 
@@ -229,7 +264,7 @@ public:
 template <typename ElementAcc, typename ElementD, typename TileShape>
 struct ScaledEpilogueColumnBias
     : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
@@ -237,17 +272,21 @@ private:
   using Bias = typename SUPER::template ColLoad<ElementD>;
 
   using Compute0 = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTCompute0 =
       cutlass::epilogue::fusion::Sm90EVT<Compute0, ScaleB, Accum>;
 
   using Compute1 = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiply_add, ElementD, float,
+      cutlass::multiply_add,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
+ public:
   using EVTCompute =
       cutlass::epilogue::fusion::Sm90EVT<Compute1, ScaleA, EVTCompute0, Bias>;
 
@@ -275,7 +314,7 @@ public:
 template <typename ElementAcc, typename ElementD, typename TileShape>
 struct ScaledEpilogueBiasAzp
     : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
@@ -287,33 +326,39 @@ private:
 
   // Compute float(accum - azp_adj), both operands are int32_t
   using ComputeAzp = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::minus, float, int32_t,
+      cutlass::minus,
+      float,
+      int32_t,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTComputeAzp =
       cutlass::epilogue::fusion::Sm90EVT<ComputeAzp, Accum, AzpWithAdj>;
 
   using ComputeScaleB = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTComputeScaleB =
       cutlass::epilogue::fusion::Sm90EVT<ComputeScaleB, ScaleB, EVTComputeAzp>;
 
   using ComputeScaleBiasA = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiply_add, ElementD, float,
+      cutlass::multiply_add,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
-  using EVTCompute =
-      cutlass::epilogue::fusion::Sm90EVT<ComputeScaleBiasA, ScaleA,
-                                         EVTComputeScaleB, Bias>;
+ public:
+  using EVTCompute = cutlass::epilogue::fusion::
+      Sm90EVT<ComputeScaleBiasA, ScaleA, EVTComputeScaleB, Bias>;
   using ArgumentType = typename EVTCompute::Arguments;
 
-  static ArgumentType
-  prepare_args(paddle::Tensor const &a_scales, paddle::Tensor const &b_scales,
-               paddle::Tensor const &azp_adj,
-               paddle::optional<paddle::Tensor> const &bias) {
+  static ArgumentType prepare_args(
+      paddle::Tensor const &a_scales,
+      paddle::Tensor const &b_scales,
+      paddle::Tensor const &azp_adj,
+      paddle::optional<paddle::Tensor> const &bias) {
     auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
     auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
     auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -340,7 +385,7 @@ public:
 template <typename ElementAcc, typename ElementD, typename TileShape>
 struct ScaledEpilogueBiasAzpToken
     : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
@@ -355,7 +400,9 @@ private:
 
   // Compute azp * azp_adj
   using ComputeAzp = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, int32_t, int32_t,
+      cutlass::multiplies,
+      int32_t,
+      int32_t,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTComputeAzp =
@@ -363,33 +410,40 @@ private:
 
   // Compute float(accum - azp*azp_adj), all operands are int32_t
   using ComputeAcc = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::minus, float, int32_t,
+      cutlass::minus,
+      float,
+      int32_t,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTComputeAcc =
       cutlass::epilogue::fusion::Sm90EVT<ComputeAcc, Accum, EVTComputeAzp>;
 
   using ComputeScaleB = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTComputeScaleB =
       cutlass::epilogue::fusion::Sm90EVT<ComputeScaleB, ScaleB, EVTComputeAcc>;
 
   using ComputeScaleBiasA = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiply_add, ElementD, float,
+      cutlass::multiply_add,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
-  using EVTCompute =
-      cutlass::epilogue::fusion::Sm90EVT<ComputeScaleBiasA, ScaleA,
-                                         EVTComputeScaleB, Bias>;
+ public:
+  using EVTCompute = cutlass::epilogue::fusion::
+      Sm90EVT<ComputeScaleBiasA, ScaleA, EVTComputeScaleB, Bias>;
   using ArgumentType = typename EVTCompute::Arguments;
 
-  static ArgumentType
-  prepare_args(paddle::Tensor const &a_scales, paddle::Tensor const &b_scales,
-               paddle::Tensor const &azp_adj, paddle::Tensor const &azp,
-               paddle::optional<paddle::Tensor> const &bias) {
+  static ArgumentType prepare_args(
+      paddle::Tensor const &a_scales,
+      paddle::Tensor const &b_scales,
+      paddle::Tensor const &azp_adj,
+      paddle::Tensor const &azp,
+      paddle::optional<paddle::Tensor> const &bias) {
     auto a_args = SUPER::template args_from_tensor<ScaleA, float>(a_scales);
     auto b_args = SUPER::template args_from_tensor<ScaleB, float>(b_scales);
     auto bias_args = SUPER::template args_from_tensor<Bias, ElementD>(bias);
@@ -414,24 +468,28 @@ public:
 template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
 struct ScaledEpilogueArray
     : private ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor> {
-private:
+ private:
   using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor>;
   using Accum = typename SUPER::Accum;
   using ScaleA = typename SUPER::template ColOrScalarLoadArray<float>;
   using ScaleB = typename SUPER::template RowOrScalarLoadArray<float>;
 
   using Compute0 = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, float, float,
+      cutlass::multiplies,
+      float,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
   using EVTCompute0 =
       cutlass::epilogue::fusion::Sm90EVT<Compute0, ScaleB, Accum>;
 
   using Compute1 = cutlass::epilogue::fusion::Sm90Compute<
-      cutlass::multiplies, ElementD, float,
+      cutlass::multiplies,
+      ElementD,
+      float,
       cutlass::FloatRoundStyle::round_to_nearest>;
 
-public:
+ public:
   using EVTCompute =
       cutlass::epilogue::fusion::Sm90EVT<Compute1, ScaleA, EVTCompute0>;
   using ArgumentType = typename EVTCompute::Arguments;
@@ -441,7 +499,8 @@ public:
 
   static ArgumentType prepare_args(float const *const *a_scales_ptr,
                                    float const *const *b_scales_ptr,
-                                   bool a_col_broadcast, bool b_row_broadcast) {
+                                   bool a_col_broadcast,
+                                   bool b_row_broadcast) {
     auto a_args = SUPER::template args_from_tensor<ScaleAArray, float>(
         a_scales_ptr, a_col_broadcast);
     auto b_args = SUPER::template args_from_tensor<ScaleBArray, float>(
@@ -452,4 +511,4 @@ public:
   }
 };
 
-}; // namespace fastdeploy::c3x
+};  // namespace fastdeploy::c3x
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue/thread/fused_activations.h b/custom_ops/gpu_ops/cutlass_extensions/epilogue/thread/fused_activations.h
index f3c622b88a..0aef590a10 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue/thread/fused_activations.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue/thread/fused_activations.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,18 +18,20 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
-  \brief Functor performing linear combination with a maximum operation used by epilogues.
+  \brief Functor performing linear combination with a maximum operation used by
+  epilogues.
 */
 
 #pragma once
@@ -46,60 +48,53 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace epilogue
-{
-namespace thread
-{
+namespace cutlass {
+namespace epilogue {
+namespace thread {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-__forceinline__ __device__ float copysignf_pos(float a, float b)
-{
-    float r;
-    r = __int_as_float(__float_as_int(a) | (__float_as_int(b) & 0x80000000));
-    return r;
+__forceinline__ __device__ float copysignf_pos(float a, float b) {
+  float r;
+  r = __int_as_float(__float_as_int(a) | (__float_as_int(b) & 0x80000000));
+  return r;
 }
 
-__forceinline__ __device__ float tanh_opt(float x)
-{
+__forceinline__ __device__ float tanh_opt(float x) {
 #if (__CUDACC_VER_MAJOR__ < 11) || (__CUDA_ARCH__ < 750)
-    float const exp_val = -1.f * fabs(2 * x);
-    return copysignf_pos((1.0f - __expf(exp_val)) / (__expf(exp_val) + 1.0f), x);
+  float const exp_val = -1.f * fabs(2 * x);
+  return copysignf_pos((1.0f - __expf(exp_val)) / (__expf(exp_val) + 1.0f), x);
 #else
-    return fast_tanh(x);
+  return fast_tanh(x);
 #endif
 }
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 template <>
-struct GELU_taylor<float>
-{
-    static bool const kIsHeavy = true;
+struct GELU_taylor<float> {
+  static bool const kIsHeavy = true;
 
-    CUTLASS_DEVICE
-    float operator()(float const& z) const
-    {
+  CUTLASS_DEVICE
+  float operator()(float const& z) const {
+    float k0 = float(0.7978845608028654);
+    float k1 = float(0.044715);
 
-        float k0 = float(0.7978845608028654);
-        float k1 = float(0.044715);
+    return float(
+        cutlass::constants::half<float>() * z *
+        (cutlass::constants::one<float>() +
+         tanh_opt(k0 * z * (cutlass::constants::one<float>() + k1 * z * z))));
+  }
 
-        return float(cutlass::constants::half<float>() * z
-            * (cutlass::constants::one<float>() + tanh_opt(k0 * z * (cutlass::constants::one<float>() + k1 * z * z))));
-    }
+  using Params = LinearCombinationGenericParams<float>;
 
-    using Params = LinearCombinationGenericParams<float>;
-
-    CUTLASS_DEVICE
-    float operator()(float const& scalar, Params const& params_) const
-    {
-        return this->operator()(scalar);
-    }
+  CUTLASS_DEVICE
+  float operator()(float const& scalar, Params const& params_) const {
+    return this->operator()(scalar);
+  }
 };
 
-} // namespace thread
-} // namespace epilogue
-} // namespace cutlass
+}  // namespace thread
+}  // namespace epilogue
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale.h b/custom_ops/gpu_ops/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale.h
index aeec5e5d0b..031f3e7048 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,20 +18,23 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
-  \brief Epilogue visitor for threadblock scoped INT8 GEMMs that uses one scaling factor per row, and one per column.
+  \brief Epilogue visitor for threadblock scoped INT8 GEMMs that uses one
+  scaling factor per row, and one per column.
 
-  original file: 3rdparty/cutlass/include/cutlass/epilogue/threadblock/epilogue_visitor_with_softmax.h
+  original file:
+  3rdparty/cutlass/include/cutlass/epilogue/threadblock/epilogue_visitor_with_softmax.h
 
 */
 
@@ -46,305 +49,312 @@
 #include "cutlass/numeric_conversion.h"
 #include "common/quantization.h"
 
-namespace cutlass
-{
-namespace epilogue
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
 
-template <typename ThreadblockShape_, int ThreadCount, typename ScaleTileIterator_, typename OutputTileIterator_,
-    typename ElementAccumulator_, typename ElementCompute_, typename ElementwiseFunctor_, bool UseMasking_ = false>
-class EpilogueVisitorPerRowPerCol
-{
-public:
-    using ThreadblockShape = ThreadblockShape_;
-    static int const kThreadCount = ThreadCount;
+template <typename ThreadblockShape_,
+          int ThreadCount,
+          typename ScaleTileIterator_,
+          typename OutputTileIterator_,
+          typename ElementAccumulator_,
+          typename ElementCompute_,
+          typename ElementwiseFunctor_,
+          bool UseMasking_ = false>
+class EpilogueVisitorPerRowPerCol {
+ public:
+  using ThreadblockShape = ThreadblockShape_;
+  static int const kThreadCount = ThreadCount;
 
-    using ScaleTileIterator = ScaleTileIterator_;
-    using OutputTileIterator = OutputTileIterator_;
-    using ElementwiseFunctor = ElementwiseFunctor_;
+  using ScaleTileIterator = ScaleTileIterator_;
+  using OutputTileIterator = OutputTileIterator_;
+  using ElementwiseFunctor = ElementwiseFunctor_;
 
-    static int const kIterations = OutputTileIterator::kIterations;
-    static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
+  static int const kIterations = OutputTileIterator::kIterations;
+  static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
 
-    using ElementOutput = typename OutputTileIterator::Element;
-    using LayoutOutput = cutlass::layout::RowMajor;
-    using ElementAccumulator = ElementAccumulator_;
+  using ElementOutput = typename OutputTileIterator::Element;
+  using LayoutOutput = cutlass::layout::RowMajor;
+  using ElementAccumulator = ElementAccumulator_;
 
-    using AlphaScaleElementType = typename ScaleTileIterator::Element;
+  using AlphaScaleElementType = typename ScaleTileIterator::Element;
 
-    using ElementCompute = ElementCompute_;
-    using AccumulatorFragment = Array<ElementAccumulator, kElementsPerAccess>;
-    using ComputeFragment = Array<ElementCompute_, kElementsPerAccess>;
-    using OutputVector = Array<ElementOutput, kElementsPerAccess>;
+  using ElementCompute = ElementCompute_;
+  using AccumulatorFragment = Array<ElementAccumulator, kElementsPerAccess>;
+  using ComputeFragment = Array<ElementCompute_, kElementsPerAccess>;
+  using OutputVector = Array<ElementOutput, kElementsPerAccess>;
 
-    static int const kThreadsPerRow = OutputTileIterator::ThreadMap::Detail::kAccessWidth;
-    static bool const kHasMultiStepsInRow = (OutputTileIterator::ThreadMap::Iterations::kColumn > 1);
+  static int const kThreadsPerRow =
+      OutputTileIterator::ThreadMap::Detail::kAccessWidth;
+  static bool const kHasMultiStepsInRow =
+      (OutputTileIterator::ThreadMap::Iterations::kColumn > 1);
 
-    /// Argument structure
-    struct Arguments
-    {
+  /// Argument structure
+  struct Arguments {
+    typename ElementwiseFunctor::Params elementwise;
+    int64_t batch_stride_alpha;
+    int64_t batch_stride_C;
+    int64_t batch_stride_D;
 
-        typename ElementwiseFunctor::Params elementwise;
-        int64_t batch_stride_alpha;
-        int64_t batch_stride_C;
-        int64_t batch_stride_D;
+    //
+    // Methods
+    //
+    Arguments() : batch_stride_alpha(0), batch_stride_C(0), batch_stride_D(0) {}
 
-        //
-        // Methods
-        //
-        Arguments()
-            : batch_stride_alpha(0)
-            , batch_stride_C(0)
-            , batch_stride_D(0)
-        {
-        }
+    Arguments(typename ElementwiseFunctor::Params elementwise_)
+        : elementwise(elementwise_),
+          batch_stride_alpha(0),
+          batch_stride_C(0),
+          batch_stride_D(0) {}
 
-        Arguments(typename ElementwiseFunctor::Params elementwise_)
-            : elementwise(elementwise_)
-            , batch_stride_alpha(0)
-            , batch_stride_C(0)
-            , batch_stride_D(0)
-        {
-        }
+    Arguments(typename ElementwiseFunctor::Params elementwise_,
+              int64_t batch_stride_alpha_,
+              int64_t batch_stride_C_,
+              int64_t batch_stride_D_)
+        : elementwise(elementwise_),
+          batch_stride_alpha(batch_stride_alpha_),
+          batch_stride_C(batch_stride_C_),
+          batch_stride_D(batch_stride_D_) {}
+  };
 
-        Arguments(typename ElementwiseFunctor::Params elementwise_, int64_t batch_stride_alpha_,
-            int64_t batch_stride_C_, int64_t batch_stride_D_)
-            : elementwise(elementwise_)
-            , batch_stride_alpha(batch_stride_alpha_)
-            , batch_stride_C(batch_stride_C_)
-            , batch_stride_D(batch_stride_D_)
-        {
-        }
-    };
+  struct Params {
+    typename ElementwiseFunctor::Params elementwise;
+    int64_t batch_stride_alpha;
+    int64_t batch_stride_C;
+    int64_t batch_stride_D;
 
-    struct Params
-    {
+    //
+    // Methods
+    //
+    CUTLASS_HOST_DEVICE
+    Params() {}
 
-        typename ElementwiseFunctor::Params elementwise;
-        int64_t batch_stride_alpha;
-        int64_t batch_stride_C;
-        int64_t batch_stride_D;
+    CUTLASS_HOST_DEVICE
+    Params(Arguments const& args)
+        : elementwise(args.elementwise),
+          batch_stride_alpha(args.batch_stride_alpha),
+          batch_stride_C(args.batch_stride_C),
+          batch_stride_D(args.batch_stride_D) {}
+  };
 
-        //
-        // Methods
-        //
-        CUTLASS_HOST_DEVICE
-        Params() {}
+  /// Shared storage
+  struct SharedStorage {};
 
-        CUTLASS_HOST_DEVICE
-        Params(Arguments const& args)
-            : elementwise(args.elementwise)
-            , batch_stride_alpha(args.batch_stride_alpha)
-            , batch_stride_C(args.batch_stride_C)
-            , batch_stride_D(args.batch_stride_D)
-        {
-        }
-    };
+ private:
+  Params const& params_;
+  SharedStorage& shared_storage_;
+  MatrixCoord extent_;
+  MatrixCoord extent_real_;
+  ElementwiseFunctor elementwise_;
 
-    /// Shared storage
-    struct SharedStorage
-    {
-    };
+  bool const per_token_quant_;
+  bool const per_channel_quant_;
 
-private:
-    Params const& params_;
-    SharedStorage& shared_storage_;
-    MatrixCoord extent_;
-    MatrixCoord extent_real_;
-    ElementwiseFunctor elementwise_;
+  AlphaScaleElementType* ptr_alpha_row_;
+  AlphaScaleElementType* ptr_alpha_col_;
+  ScaleTileIterator iterator_alpha_col_;
+  OutputTileIterator iterator_C_;
+  OutputTileIterator iterator_D_;
 
-    bool const per_token_quant_;
-    bool const per_channel_quant_;
+  AlphaScaleElementType element_alpha_row_ = 1.0f;
+  AlphaScaleElementType element_alpha_col_ = 1.0f;
+  typename ScaleTileIterator::Fragment fragment_alpha_col_;
+  typename OutputTileIterator::Fragment fragment_C_;
+  typename OutputTileIterator::Fragment fragment_D_;
 
-    AlphaScaleElementType* ptr_alpha_row_;
-    AlphaScaleElementType* ptr_alpha_col_;
-    ScaleTileIterator iterator_alpha_col_;
-    OutputTileIterator iterator_C_;
-    OutputTileIterator iterator_D_;
+  ElementAccumulator beta_;
 
-    AlphaScaleElementType element_alpha_row_ = 1.0f;
-    AlphaScaleElementType element_alpha_col_ = 1.0f;
-    typename ScaleTileIterator::Fragment fragment_alpha_col_;
-    typename OutputTileIterator::Fragment fragment_C_;
-    typename OutputTileIterator::Fragment fragment_D_;
+  int column_offset_;
 
-    ElementAccumulator beta_;
+  MatrixCoord thread_offset_;
 
-    int column_offset_;
+ public:
+  CUTLASS_DEVICE
+  EpilogueVisitorPerRowPerCol(
+      Params const& params,
+      SharedStorage& shared_storage,
+      cutlass::MatrixCoord const& problem_size,
+      int thread_idx,
+      int warp_idx,
+      int lane_idx,
+      typename ScaleTileIterator::Params params_alpha_col,
+      typename OutputTileIterator::Params params_C,
+      typename OutputTileIterator::Params params_D,
+      common::QuantMode quant_option,
+      AlphaScaleElementType* ptr_alpha_row,
+      AlphaScaleElementType* ptr_alpha_col,
+      typename OutputTileIterator::Element* ptr_C,
+      typename OutputTileIterator::Element* ptr_D,
+      cutlass::MatrixCoord const& threadblock_offset = cutlass::MatrixCoord(0,
+                                                                            0),
+      int column_offset = 0,
+      cutlass::MatrixCoord const& problem_size_real = cutlass::MatrixCoord(0,
+                                                                           0))
+      : params_(params),
+        shared_storage_(shared_storage),
+        extent_(problem_size),
+        elementwise_(params.elementwise),
+        per_token_quant_(quant_option.hasPerTokenScaling()),
+        per_channel_quant_(quant_option.hasPerChannelScaling()),
+        ptr_alpha_row_(ptr_alpha_row),
+        ptr_alpha_col_(ptr_alpha_col),
+        iterator_alpha_col_(params_alpha_col,
+                            ptr_alpha_col,
+                            problem_size,
+                            thread_idx,
+                            threadblock_offset),
+        iterator_C_(
+            params_C, ptr_C, problem_size, thread_idx, threadblock_offset),
+        iterator_D_(
+            params_D, ptr_D, problem_size, thread_idx, threadblock_offset),
+        extent_real_(problem_size_real) {
+    beta_ = (params.elementwise.beta_ptr ? *params.elementwise.beta_ptr
+                                         : params.elementwise.beta);
 
-    MatrixCoord thread_offset_;
-
-public:
-    CUTLASS_DEVICE
-    EpilogueVisitorPerRowPerCol(Params const& params, SharedStorage& shared_storage,
-        cutlass::MatrixCoord const& problem_size, int thread_idx, int warp_idx, int lane_idx,
-        typename ScaleTileIterator::Params params_alpha_col, typename OutputTileIterator::Params params_C,
-        typename OutputTileIterator::Params params_D, common::QuantMode quant_option, AlphaScaleElementType* ptr_alpha_row,
-        AlphaScaleElementType* ptr_alpha_col, typename OutputTileIterator::Element* ptr_C,
-        typename OutputTileIterator::Element* ptr_D,
-        cutlass::MatrixCoord const& threadblock_offset = cutlass::MatrixCoord(0, 0), int column_offset = 0,
-        cutlass::MatrixCoord const& problem_size_real = cutlass::MatrixCoord(0, 0))
-        : params_(params)
-        , shared_storage_(shared_storage)
-        , extent_(problem_size)
-        , elementwise_(params.elementwise)
-        , per_token_quant_(quant_option.hasPerTokenScaling())
-        , per_channel_quant_(quant_option.hasPerChannelScaling())
-        , ptr_alpha_row_(ptr_alpha_row)
-        , ptr_alpha_col_(ptr_alpha_col)
-        , iterator_alpha_col_(params_alpha_col, ptr_alpha_col, problem_size, thread_idx, threadblock_offset)
-        , iterator_C_(params_C, ptr_C, problem_size, thread_idx, threadblock_offset)
-        , iterator_D_(params_D, ptr_D, problem_size, thread_idx, threadblock_offset)
-        , extent_real_(problem_size_real)
-    {
-        beta_ = (params.elementwise.beta_ptr ? *params.elementwise.beta_ptr : params.elementwise.beta);
-
-        if (beta_ == ElementAccumulator())
-        {
-            iterator_C_.clear_mask();
-        }
-
-        if (!per_channel_quant_ && (ptr_alpha_col_ != nullptr))
-        {
-            element_alpha_col_ = *ptr_alpha_col_;
-        }
-
-        if (!per_token_quant_ && (ptr_alpha_row_ != nullptr))
-        {
-            element_alpha_row_ = *ptr_alpha_row_;
-        }
+    if (beta_ == ElementAccumulator()) {
+      iterator_C_.clear_mask();
     }
 
-    /// Helper to indicate split-K behavior
-    CUTLASS_DEVICE
-    void set_k_partition(int split_k_index, ///< Index of this threadblock within split-K partitioned scheme
-        int split_k_slices)
-    {                                       ///< Total number of split-K slices
+    if (!per_channel_quant_ && (ptr_alpha_col_ != nullptr)) {
+      element_alpha_col_ = *ptr_alpha_col_;
     }
 
-    /// Called to set the batch index
-    CUTLASS_DEVICE
-    void set_batch_index(int batch_idx)
-    {
-        iterator_alpha_col_.add_pointer_offset(batch_idx * params_.batch_stride_alpha);
-        iterator_C_.add_pointer_offset(batch_idx * params_.batch_stride_C);
-        iterator_D_.add_pointer_offset(batch_idx * params_.batch_stride_D);
+    if (!per_token_quant_ && (ptr_alpha_row_ != nullptr)) {
+      element_alpha_row_ = *ptr_alpha_row_;
+    }
+  }
+
+  /// Helper to indicate split-K behavior
+  CUTLASS_DEVICE
+  void set_k_partition(
+      int split_k_index,     ///< Index of this threadblock within split-K
+                             ///< partitioned scheme
+      int split_k_slices) {  ///< Total number of split-K slices
+  }
+
+  /// Called to set the batch index
+  CUTLASS_DEVICE
+  void set_batch_index(int batch_idx) {
+    iterator_alpha_col_.add_pointer_offset(batch_idx *
+                                           params_.batch_stride_alpha);
+    iterator_C_.add_pointer_offset(batch_idx * params_.batch_stride_C);
+    iterator_D_.add_pointer_offset(batch_idx * params_.batch_stride_D);
+  }
+
+  /// Called at the start of the epilogue just before iterating over accumulator
+  /// slices
+  CUTLASS_DEVICE
+  void begin_epilogue() {
+    if (per_channel_quant_) {
+      iterator_alpha_col_.load(fragment_alpha_col_);
+    }
+  }
+
+  /// Called at the start of one step before starting accumulator exchange
+  CUTLASS_DEVICE
+  void begin_step(int step_idx) {
+    fragment_D_.clear();
+    fragment_C_.clear();
+
+    if (elementwise_.kScale !=
+        cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling) {
+      iterator_C_.load(fragment_C_);
+      ++iterator_C_;
+    }
+  }
+
+  /// Called at the start of a row
+  CUTLASS_DEVICE
+  void begin_row(int row_idx) {
+    // load alpha_row in begin_step only when per token(row) scaling is used
+    if (per_token_quant_) {
+      int thread_offset_row =
+          iterator_D_.thread_start_row() +
+          OutputTileIterator::ThreadMap::iteration_offset(row_idx).row();
+
+      arch::global_load<AlphaScaleElementType, sizeof(AlphaScaleElementType)>(
+          element_alpha_row_,
+          ptr_alpha_row_ + thread_offset_row,
+          thread_offset_row < extent_.row());
+    }
+  }
+
+  /// Called after accumulators have been exchanged for each accumulator vector
+  CUTLASS_DEVICE
+  void visit(int iter_idx,
+             int row_idx,
+             int column_idx,
+             int frag_idx,
+             AccumulatorFragment const& accum) {
+    NumericArrayConverter<ElementCompute,
+                          ElementAccumulator,
+                          kElementsPerAccess>
+        source_converter;
+
+    ComputeFragment result = source_converter(accum);
+    if (per_channel_quant_) {
+      ComputeFragment alpha_col =
+          reinterpret_cast<ComputeFragment*>(&fragment_alpha_col_)[column_idx];
+      result = per_token_channel_scale_accumulator_(
+          result, alpha_col, element_alpha_row_);
+    } else {
+      result = per_token_scale_accumulator_(
+          result, element_alpha_col_, element_alpha_row_);
     }
 
-    /// Called at the start of the epilogue just before iterating over accumulator slices
-    CUTLASS_DEVICE
-    void begin_epilogue()
-    {
-        if (per_channel_quant_)
-        {
-            iterator_alpha_col_.load(fragment_alpha_col_);
-        }
+    // Convert to the output
+    NumericArrayConverter<ElementOutput, ElementCompute, kElementsPerAccess>
+        output_converter;
+    OutputVector& output =
+        reinterpret_cast<OutputVector*>(&fragment_D_)[frag_idx];
+    output = output_converter(result);
+  }
+
+  /// Called at the end of a row
+  CUTLASS_DEVICE
+  void end_row(int row_idx) {}
+
+  /// Called after all accumulator elements have been visited
+  CUTLASS_DEVICE
+  void end_step(int step_idx) {
+    iterator_D_.store(fragment_D_);
+    ++iterator_D_;
+  }
+
+  /// Called after all steps have been completed
+  CUTLASS_DEVICE
+  void end_epilogue() {}
+
+ private:
+  CUTLASS_DEVICE
+  ComputeFragment per_token_channel_scale_accumulator_(
+      ComputeFragment const& accum,
+      ComputeFragment const& scale_col,
+      AlphaScaleElementType const& scale_row) {
+    ComputeFragment result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < ComputeFragment::kElements; ++i) {
+      result[i] = accum[i] * (scale_col[i] * scale_row);
     }
 
-    /// Called at the start of one step before starting accumulator exchange
-    CUTLASS_DEVICE
-    void begin_step(int step_idx)
-    {
-        fragment_D_.clear();
-        fragment_C_.clear();
+    return result;
+  }
 
-        if (elementwise_.kScale != cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling)
-        {
-            iterator_C_.load(fragment_C_);
-            ++iterator_C_;
-        }
+  CUTLASS_DEVICE
+  ComputeFragment per_token_scale_accumulator_(
+      ComputeFragment const& accum,
+      AlphaScaleElementType const& scale_col,
+      AlphaScaleElementType const& scale_row) {
+    ComputeFragment result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < ComputeFragment::kElements; ++i) {
+      result[i] = accum[i] * (scale_col * scale_row);
     }
 
-    /// Called at the start of a row
-    CUTLASS_DEVICE
-    void begin_row(int row_idx)
-    {
-        // load alpha_row in begin_step only when per token(row) scaling is used
-        if (per_token_quant_)
-        {
-            int thread_offset_row
-                = iterator_D_.thread_start_row() + OutputTileIterator::ThreadMap::iteration_offset(row_idx).row();
-
-            arch::global_load<AlphaScaleElementType, sizeof(AlphaScaleElementType)>(
-                element_alpha_row_, ptr_alpha_row_ + thread_offset_row, thread_offset_row < extent_.row());
-        }
-    }
-
-    /// Called after accumulators have been exchanged for each accumulator vector
-    CUTLASS_DEVICE
-    void visit(int iter_idx, int row_idx, int column_idx, int frag_idx, AccumulatorFragment const& accum)
-    {
-
-        NumericArrayConverter<ElementCompute, ElementAccumulator, kElementsPerAccess> source_converter;
-
-        ComputeFragment result = source_converter(accum);
-        if (per_channel_quant_)
-        {
-            ComputeFragment alpha_col = reinterpret_cast<ComputeFragment*>(&fragment_alpha_col_)[column_idx];
-            result = per_token_channel_scale_accumulator_(result, alpha_col, element_alpha_row_);
-        }
-        else
-        {
-            result = per_token_scale_accumulator_(result, element_alpha_col_, element_alpha_row_);
-        }
-
-        // Convert to the output
-        NumericArrayConverter<ElementOutput, ElementCompute, kElementsPerAccess> output_converter;
-        OutputVector& output = reinterpret_cast<OutputVector*>(&fragment_D_)[frag_idx];
-        output = output_converter(result);
-    }
-
-    /// Called at the end of a row
-    CUTLASS_DEVICE
-    void end_row(int row_idx) {}
-
-    /// Called after all accumulator elements have been visited
-    CUTLASS_DEVICE
-    void end_step(int step_idx)
-    {
-
-        iterator_D_.store(fragment_D_);
-        ++iterator_D_;
-    }
-
-    /// Called after all steps have been completed
-    CUTLASS_DEVICE
-    void end_epilogue() {}
-
-private:
-    CUTLASS_DEVICE
-    ComputeFragment per_token_channel_scale_accumulator_(
-        ComputeFragment const& accum, ComputeFragment const& scale_col, AlphaScaleElementType const& scale_row)
-    {
-
-        ComputeFragment result;
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < ComputeFragment::kElements; ++i)
-        {
-            result[i] = accum[i] * (scale_col[i] * scale_row);
-        }
-
-        return result;
-    }
-
-    CUTLASS_DEVICE
-    ComputeFragment per_token_scale_accumulator_(
-        ComputeFragment const& accum, AlphaScaleElementType const& scale_col, AlphaScaleElementType const& scale_row)
-    {
-
-        ComputeFragment result;
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < ComputeFragment::kElements; ++i)
-        {
-            result[i] = accum[i] * (scale_col * scale_row);
-        }
-
-        return result;
-    }
+    return result;
+  }
 };
 
-} // namespace threadblock
-} // namespace epilogue
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace epilogue
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h b/custom_ops/gpu_ops/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h
index 6f26d79017..2f89cb3f21 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,23 +18,26 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
   \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
 
-  The epilogue rearranges the result of a matrix product through shared memory to match canonical
-  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+  The epilogue rearranges the result of a matrix product through shared memory
+  to match canonical tensor layouts in global memory. Epilogues support
+  conversion and reduction operations.
 
-  original file: 3rdparty/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op.h
+  original file:
+  3rdparty/cutlass/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op.h
 
 */
 
@@ -80,35 +83,45 @@
 
 ////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace epilogue
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace epilogue {
+namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-namespace detail
-{
+namespace detail {
 
-/// Partial specialization for bfloat16_t <= int32_t x 8 epilogues avoids shared memory bank conflicts.
-template <typename ThreadblockShape, typename WarpShape, typename InstructionShape, typename ThreadMap>
-struct DefaultIteratorsTensorOp<cutlass::bfloat16_t, int32_t, 8, ThreadblockShape, WarpShape, InstructionShape,
-    ThreadMap>
-{
-    using WarpTileIterator
-        = cutlass::epilogue::warp::TileIteratorTensorOpMixed<WarpShape, InstructionShape, int32_t, 32, 16, 8, 8>;
+/// Partial specialization for bfloat16_t <= int32_t x 8 epilogues avoids shared
+/// memory bank conflicts.
+template <typename ThreadblockShape,
+          typename WarpShape,
+          typename InstructionShape,
+          typename ThreadMap>
+struct DefaultIteratorsTensorOp<cutlass::bfloat16_t,
+                                int32_t,
+                                8,
+                                ThreadblockShape,
+                                WarpShape,
+                                InstructionShape,
+                                ThreadMap> {
+  using WarpTileIterator =
+      cutlass::epilogue::warp::TileIteratorTensorOpMixed<WarpShape,
+                                                         InstructionShape,
+                                                         int32_t,
+                                                         32,
+                                                         16,
+                                                         8,
+                                                         8>;
 
-    using SharedLoadIterator
-        = cutlass::epilogue::threadblock::SharedLoadIteratorMixed<ThreadMap, int32_t, 32, 16, 8, 8>;
+  using SharedLoadIterator = cutlass::epilogue::threadblock::
+      SharedLoadIteratorMixed<ThreadMap, int32_t, 32, 16, 8, 8>;
 
-    static int const kFragmentsPerIteration = 2;
+  static int const kFragmentsPerIteration = 2;
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace detail
+}  // namespace detail
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -116,167 +129,159 @@ struct DefaultIteratorsTensorOp<cutlass::bfloat16_t, int32_t, 8, ThreadblockShap
 ///
 /// Satisfies: ReadableTileIterator
 ///
-template <typename ThreadMap_ ///< Thread map (concept: OutputTileThreadMap)
-    >
-class SharedLoadIteratorMixed<ThreadMap_, int32_t, 32, 16, 8, 8>
-{
-public:
-    using ThreadMap = ThreadMap_;
-    using Shape = typename ThreadMap::Shape;
+template <typename ThreadMap_  ///< Thread map (concept: OutputTileThreadMap)
+          >
+class SharedLoadIteratorMixed<ThreadMap_, int32_t, 32, 16, 8, 8> {
+ public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
 
-    using Element = int32_t;
+  using Element = int32_t;
 
-    using Layout = layout::RowMajor;
-    using TensorRef = TensorRef<Element, Layout>;
-    using ConstTensorRef = typename TensorRef::ConstTensorRef;
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
 
-    using Index = typename Layout::Index;
-    using LongIndex = typename Layout::LongIndex;
-    using TensorCoord = MatrixCoord;
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
 
-    static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
 
-    static int const kAlignment = ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value / 8;
+  static int const kAlignment =
+      ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value / 8;
 
-    static int const kThreads = ThreadMap::kThreads;
+  static int const kThreads = ThreadMap::kThreads;
 
-    /// Fragment object
-    using Fragment = Array<Element,
-        ThreadMap::Iterations::kColumn * ThreadMap::Iterations::kRow * ThreadMap::Iterations::kGroup
-            * ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
+  /// Fragment object
+  using Fragment =
+      Array<Element,
+            ThreadMap::Iterations::kColumn * ThreadMap::Iterations::kRow *
+                ThreadMap::Iterations::kGroup *
+                ThreadMap::Iterations::kCluster *
+                ThreadMap::kElementsPerAccess>;
 
-    /// Memory access size
-    using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess, kAlignment>;
+  /// Memory access size
+  using AccessType =
+      AlignedArray<Element, ThreadMap::kElementsPerAccess, kAlignment>;
 
-    /// Vector type used for SMEM loads
-    using LoadType = AlignedArray<Element, const_min(128 / sizeof_bits<Element>::value, ThreadMap::kElementsPerAccess),
-        const_min(16, kAlignment)>;
+  /// Vector type used for SMEM loads
+  using LoadType = AlignedArray<Element,
+                                const_min(128 / sizeof_bits<Element>::value,
+                                          ThreadMap::kElementsPerAccess),
+                                const_min(16, kAlignment)>;
 
-    static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements;
+  static int const kLoadsPerAccess =
+      AccessType::kElements / LoadType::kElements;
 
-private:
-    //
-    // Data members
-    //
+ private:
+  //
+  // Data members
+  //
 
-    /// Byte-level pointer
-    LoadType const* pointers_[kLoadsPerAccess];
+  /// Byte-level pointer
+  LoadType const* pointers_[kLoadsPerAccess];
 
-    /// Stride along adjacent rows in units of LoadType
-    int stride_;
+  /// Stride along adjacent rows in units of LoadType
+  int stride_;
 
-public:
-    //
-    // Methods
-    //
+ public:
+  //
+  // Methods
+  //
 
-    /// Constructor
-    CUTLASS_DEVICE
-    SharedLoadIteratorMixed(TensorRef ref, int thread_idx)
-        : stride_((ref.stride(0) / LoadType::kElements))
-    {
+  /// Constructor
+  CUTLASS_DEVICE
+  SharedLoadIteratorMixed(TensorRef ref, int thread_idx)
+      : stride_((ref.stride(0) / LoadType::kElements)) {
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
 
-        TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
+    // Initialize pointers
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] = reinterpret_cast<LoadType const*>(ref.data());
 
-        // Initialize pointers
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < kLoadsPerAccess; ++i)
-        {
-            pointers_[i] = reinterpret_cast<LoadType const*>(ref.data());
+      int col_idx =
+          (thread_offset.column() / kElementsPerAccess) * kLoadsPerAccess;
+      int bank_offset =
+          (col_idx * int(sizeof(LoadType)) / 128) % kLoadsPerAccess;
 
-            int col_idx = (thread_offset.column() / kElementsPerAccess) * kLoadsPerAccess;
-            int bank_offset = (col_idx * int(sizeof(LoadType)) / 128) % kLoadsPerAccess;
+      col_idx += (bank_offset + i) % kLoadsPerAccess;
 
-            col_idx += (bank_offset + i) % kLoadsPerAccess;
-
-            pointers_[i] += thread_offset.row() * stride_ + col_idx;
-        }
+      pointers_[i] += thread_offset.row() * stride_ + col_idx;
     }
+  }
 
-    /// Adds a pointer offset in units of Element
-    CUTLASS_HOST_DEVICE
-    void add_pointer_offset(LongIndex pointer_offset)
-    {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < kLoadsPerAccess; ++i)
-        {
-            pointers_[i] += pointer_offset / LoadType::kElements;
-        }
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += pointer_offset / LoadType::kElements;
     }
+  }
 
-    CUTLASS_DEVICE
-    void add_tile_offset(TensorCoord const& offset)
-    {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < kLoadsPerAccess; ++i)
-        {
-            pointers_[i]
-                += offset.row() * Shape::kRow * stride_ + offset.column() * Shape::kColumn / LoadType::kElements;
-        }
+  CUTLASS_DEVICE
+  void add_tile_offset(TensorCoord const& offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += offset.row() * Shape::kRow * stride_ +
+                      offset.column() * Shape::kColumn / LoadType::kElements;
     }
+  }
 
-    /// Loads a fragment from memory
-    CUTLASS_DEVICE
-    void load_with_pointer_offset(Fragment& frag, Index pointer_offset) const
-    {
-
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_pointer_offset(Fragment& frag, Index pointer_offset) const {
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
+         ++cluster) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
         CUTLASS_PRAGMA_UNROLL
-        for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster)
-        {
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+          int row_ptr_offset = row * ThreadMap::Delta::kRow * stride_ +
+                               group * ThreadMap::Delta::kGroup * stride_ +
+                               cluster * ThreadMap::Delta::kCluster * stride_ +
+                               pointer_offset / LoadType::kElements;
+
+          int frag_row_idx =
+              (row + ThreadMap::Iterations::kRow *
+                         (group + ThreadMap::Iterations::kGroup * cluster));
+
+          LoadType* frag_ptr = reinterpret_cast<LoadType*>(&frag);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn;
+               ++column) {
+            int frag_idx =
+                frag_row_idx * ThreadMap::Iterations::kColumn + column;
 
             CUTLASS_PRAGMA_UNROLL
-            for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group)
-            {
+            for (int v = 0; v < kLoadsPerAccess; ++v) {
+              int vector_idx = (column * ThreadMap::Delta::kColumn /
+                                kElementsPerAccess * kLoadsPerAccess);
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int row = 0; row < ThreadMap::Iterations::kRow; ++row)
-                {
+              LoadType const* memory_pointer = pointers_[v] + row_ptr_offset;
 
-                    int row_ptr_offset = row * ThreadMap::Delta::kRow * stride_
-                        + group * ThreadMap::Delta::kGroup * stride_ + cluster * ThreadMap::Delta::kCluster * stride_
-                        + pointer_offset / LoadType::kElements;
-
-                    int frag_row_idx
-                        = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
-
-                    LoadType* frag_ptr = reinterpret_cast<LoadType*>(&frag);
-
-                    CUTLASS_PRAGMA_UNROLL
-                    for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column)
-                    {
-
-                        int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
-
-                        CUTLASS_PRAGMA_UNROLL
-                        for (int v = 0; v < kLoadsPerAccess; ++v)
-                        {
-
-                            int vector_idx
-                                = (column * ThreadMap::Delta::kColumn / kElementsPerAccess * kLoadsPerAccess);
-
-                            LoadType const* memory_pointer = pointers_[v] + row_ptr_offset;
-
-                            frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[vector_idx];
-                        }
-                    }
-                }
+              frag_ptr[frag_idx * kLoadsPerAccess + v] =
+                  memory_pointer[vector_idx];
             }
+          }
         }
+      }
     }
+  }
 
-    /// Loads a fragment
-    CUTLASS_DEVICE
-    void load(Fragment& frag) const
-    {
-
-        load_with_pointer_offset(frag, 0);
-    }
+  /// Loads a fragment
+  CUTLASS_DEVICE
+  void load(Fragment& frag) const { load_with_pointer_offset(frag, 0); }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace threadblock
-} // namespace epilogue
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace epilogue
+}  // namespace cutlass
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/epilogue_helpers.h b/custom_ops/gpu_ops/cutlass_extensions/epilogue_helpers.h
index 6ed5b9b920..ca8209d59c 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/epilogue_helpers.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/epilogue_helpers.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,9 +17,10 @@
 /**
  * @file epilogue_helpers.h
  *
- * This file includes types for the epilogues. The empty structs exist so we can signal to template
- * code the type of epilogue we want to run, and let the underlying code specify the details such as
- * element types, accumulator type and elements per vector access.
+ * This file includes types for the epilogues. The empty structs exist so we can
+ * signal to template code the type of epilogue we want to run, and let the
+ * underlying code specify the details such as element types, accumulator type
+ * and elements per vector access.
  *
  */
 
@@ -33,107 +34,161 @@
 
 // #include "cutlass/epilogue/fusion/operations.hpp"
 
-namespace cutlass_extensions
-{
+namespace cutlass_extensions {
 
-struct EpilogueOpBiasSilu
-{
+struct EpilogueOpBiasSilu {};
+
+struct EpilogueOpBiasReLU {};
+
+struct EpilogueOpBiasFtGelu {};
+
+struct EpilogueOpBias {};
+
+struct EpilogueOpDefaultSilu {};
+
+struct EpilogueOpDefaultReLU {};
+
+struct EpilogueOpDefaultFtGelu {};
+
+struct EpilogueOpDefault {};
+
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator,
+          typename Op>
+struct Epilogue {
+  static_assert(sizeof(ElementType) == 0, "Unrecognized Epilogue Tag");
 };
 
-struct EpilogueOpBiasReLU
-{
+constexpr auto BiasScaleMode =
+    cutlass::epilogue::thread::ScaleType::NoBetaScaling;
+
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator>
+struct Epilogue<ElementType,
+                ElementsPerVectorAccess,
+                ElementAccumulator,
+                EpilogueOpBiasSilu> {
+  using Op =
+      cutlass::epilogue::thread::LinearCombinationSilu<ElementType,
+                                                       ElementsPerVectorAccess,
+                                                       ElementAccumulator,
+                                                       ElementAccumulator,
+                                                       BiasScaleMode>;
 };
 
-struct EpilogueOpBiasFtGelu
-{
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator>
+struct Epilogue<ElementType,
+                ElementsPerVectorAccess,
+                ElementAccumulator,
+                EpilogueOpBiasReLU> {
+  using Op =
+      cutlass::epilogue::thread::LinearCombinationRelu<ElementType,
+                                                       ElementsPerVectorAccess,
+                                                       ElementAccumulator,
+                                                       ElementAccumulator,
+                                                       BiasScaleMode>;
 };
 
-struct EpilogueOpBias
-{
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator>
+struct Epilogue<ElementType,
+                ElementsPerVectorAccess,
+                ElementAccumulator,
+                EpilogueOpBiasFtGelu> {
+  using Op = cutlass::epilogue::thread::LinearCombinationGeneric<
+      cutlass::epilogue::thread::GELU_taylor,
+      ElementType,
+      ElementsPerVectorAccess,
+      ElementAccumulator,
+      ElementAccumulator,
+      BiasScaleMode,
+      cutlass::FloatRoundStyle::round_to_nearest,
+      true>;
 };
 
-struct EpilogueOpDefaultSilu
-{
-};
-
-struct EpilogueOpDefaultReLU
-{
-};
-
-struct EpilogueOpDefaultFtGelu
-{
-};
-
-struct EpilogueOpDefault
-{
-};
-
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator, typename Op>
-struct Epilogue
-{
-    static_assert(sizeof(ElementType) == 0, "Unrecognized Epilogue Tag");
-};
-
-constexpr auto BiasScaleMode = cutlass::epilogue::thread::ScaleType::NoBetaScaling;
-
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator>
-struct Epilogue<ElementType, ElementsPerVectorAccess, ElementAccumulator, EpilogueOpBiasSilu>
-{
-    using Op = cutlass::epilogue::thread::LinearCombinationSilu<ElementType, ElementsPerVectorAccess,
-        ElementAccumulator, ElementAccumulator, BiasScaleMode>;
-};
-
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator>
-struct Epilogue<ElementType, ElementsPerVectorAccess, ElementAccumulator, EpilogueOpBiasReLU>
-{
-    using Op = cutlass::epilogue::thread::LinearCombinationRelu<ElementType, ElementsPerVectorAccess,
-        ElementAccumulator, ElementAccumulator, BiasScaleMode>;
-};
-
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator>
-struct Epilogue<ElementType, ElementsPerVectorAccess, ElementAccumulator, EpilogueOpBiasFtGelu>
-{
-    using Op = cutlass::epilogue::thread::LinearCombinationGeneric<cutlass::epilogue::thread::GELU_taylor, ElementType,
-        ElementsPerVectorAccess, ElementAccumulator, ElementAccumulator, BiasScaleMode,
-        cutlass::FloatRoundStyle::round_to_nearest, true>;
-};
-
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator>
-struct Epilogue<ElementType, ElementsPerVectorAccess, ElementAccumulator, EpilogueOpBias>
-{
-    using Op = cutlass::epilogue::thread::LinearCombination<ElementType, ElementsPerVectorAccess, ElementAccumulator,
-        ElementAccumulator, BiasScaleMode>;
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator>
+struct Epilogue<ElementType,
+                ElementsPerVectorAccess,
+                ElementAccumulator,
+                EpilogueOpBias> {
+  using Op =
+      cutlass::epilogue::thread::LinearCombination<ElementType,
+                                                   ElementsPerVectorAccess,
+                                                   ElementAccumulator,
+                                                   ElementAccumulator,
+                                                   BiasScaleMode>;
 };
 
 constexpr auto DefaultScaleMode = cutlass::epilogue::thread::ScaleType::Default;
 
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator>
-struct Epilogue<ElementType, ElementsPerVectorAccess, ElementAccumulator, EpilogueOpDefaultSilu>
-{
-    using Op = cutlass::epilogue::thread::LinearCombinationSilu<ElementType, ElementsPerVectorAccess,
-        ElementAccumulator, ElementAccumulator, DefaultScaleMode>;
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator>
+struct Epilogue<ElementType,
+                ElementsPerVectorAccess,
+                ElementAccumulator,
+                EpilogueOpDefaultSilu> {
+  using Op =
+      cutlass::epilogue::thread::LinearCombinationSilu<ElementType,
+                                                       ElementsPerVectorAccess,
+                                                       ElementAccumulator,
+                                                       ElementAccumulator,
+                                                       DefaultScaleMode>;
 };
 
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator>
-struct Epilogue<ElementType, ElementsPerVectorAccess, ElementAccumulator, EpilogueOpDefaultReLU>
-{
-    using Op = cutlass::epilogue::thread::LinearCombinationRelu<ElementType, ElementsPerVectorAccess,
-        ElementAccumulator, ElementAccumulator, DefaultScaleMode>;
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator>
+struct Epilogue<ElementType,
+                ElementsPerVectorAccess,
+                ElementAccumulator,
+                EpilogueOpDefaultReLU> {
+  using Op =
+      cutlass::epilogue::thread::LinearCombinationRelu<ElementType,
+                                                       ElementsPerVectorAccess,
+                                                       ElementAccumulator,
+                                                       ElementAccumulator,
+                                                       DefaultScaleMode>;
 };
 
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator>
-struct Epilogue<ElementType, ElementsPerVectorAccess, ElementAccumulator, EpilogueOpDefaultFtGelu>
-{
-    using Op = cutlass::epilogue::thread::LinearCombinationGeneric<cutlass::epilogue::thread::GELU_taylor, ElementType,
-        ElementsPerVectorAccess, ElementAccumulator, ElementAccumulator, DefaultScaleMode,
-        cutlass::FloatRoundStyle::round_to_nearest, true>;
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator>
+struct Epilogue<ElementType,
+                ElementsPerVectorAccess,
+                ElementAccumulator,
+                EpilogueOpDefaultFtGelu> {
+  using Op = cutlass::epilogue::thread::LinearCombinationGeneric<
+      cutlass::epilogue::thread::GELU_taylor,
+      ElementType,
+      ElementsPerVectorAccess,
+      ElementAccumulator,
+      ElementAccumulator,
+      DefaultScaleMode,
+      cutlass::FloatRoundStyle::round_to_nearest,
+      true>;
 };
 
-template <typename ElementType, int ElementsPerVectorAccess, typename ElementAccumulator>
-struct Epilogue<ElementType, ElementsPerVectorAccess, ElementAccumulator, EpilogueOpDefault>
-{
-    using Op = cutlass::epilogue::thread::LinearCombination<ElementType, ElementsPerVectorAccess, ElementAccumulator,
-        ElementAccumulator, DefaultScaleMode>;
+template <typename ElementType,
+          int ElementsPerVectorAccess,
+          typename ElementAccumulator>
+struct Epilogue<ElementType,
+                ElementsPerVectorAccess,
+                ElementAccumulator,
+                EpilogueOpDefault> {
+  using Op =
+      cutlass::epilogue::thread::LinearCombination<ElementType,
+                                                   ElementsPerVectorAccess,
+                                                   ElementAccumulator,
+                                                   ElementAccumulator,
+                                                   DefaultScaleMode>;
 };
 
-} // namespace cutlass_extensions
+}  // namespace cutlass_extensions
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_builder.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_builder.hpp
index d327eb18ae..5d5b99b5b2 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_builder.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_builder.hpp
@@ -21,7 +21,6 @@
 
 #include "cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp"
 
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 namespace cutlass::gemm::collective {
@@ -29,19 +28,17 @@ namespace cutlass::gemm::collective {
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 // GMMA_TMA_WS_SS (BlockScaled Builders)
-template <
-  class ElementA,
-  class GmemLayoutATag,
-  int AlignmentA,
-  class ElementB,
-  class GmemLayoutBTag,
-  int AlignmentB,
-  class ElementAccumulator,
-  class TileShape_MNK,
-  class ClusterShape_MNK,
-  class StageCountType,
-  int ScaleGranularityM
->
+template <class ElementA,
+          class GmemLayoutATag,
+          int AlignmentA,
+          class ElementB,
+          class GmemLayoutBTag,
+          int AlignmentB,
+          class ElementAccumulator,
+          class TileShape_MNK,
+          class ClusterShape_MNK,
+          class StageCountType,
+          int ScaleGranularityM>
 struct CollectiveBuilder<
     arch::Sm90,
     arch::OpClassTensorOp,
@@ -55,82 +52,124 @@ struct CollectiveBuilder<
     TileShape_MNK,
     ClusterShape_MNK,
     StageCountType,
-    KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>,
-    cute::enable_if_t<
-      not detail::is_use_rmem_A<ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag>()>
-> {
-  using KernelScheduleType = KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>;
+    KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<
+        ScaleGranularityM>,
+    cute::enable_if_t<not detail::is_use_rmem_A<ElementA,
+                                                GmemLayoutATag,
+                                                ElementB,
+                                                GmemLayoutBTag>()>> {
+  using KernelScheduleType =
+      KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<
+          ScaleGranularityM>;
 
   static_assert(is_static<TileShape_MNK>::value);
   static_assert(is_static<ClusterShape_MNK>::value);
 #ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED
-  static_assert(cutlass::detail::dependent_false<ElementA>, "Unsupported Toolkit for SM90 Collective Builder\n");
+  static_assert(cutlass::detail::dependent_false<ElementA>,
+                "Unsupported Toolkit for SM90 Collective Builder\n");
 #endif
-  static_assert(detail::is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, detail::tma_alignment_bytes>(),
+  static_assert(detail::is_aligned<ElementA,
+                                   AlignmentA,
+                                   ElementB,
+                                   AlignmentB,
+                                   detail::tma_alignment_bytes>(),
                 "Should meet TMA alignment requirement\n");
 
-  static constexpr bool IsArrayOfPointersGemm = (cute::is_any_of_v<KernelScheduleType,
-                                                                   KernelPtrArrayTmaWarpSpecializedCooperative,
-                                                                   KernelPtrArrayTmaWarpSpecializedPingpong>);
+  static constexpr bool IsArrayOfPointersGemm =
+      (cute::is_any_of_v<KernelScheduleType,
+                         KernelPtrArrayTmaWarpSpecializedCooperative,
+                         KernelPtrArrayTmaWarpSpecializedPingpong>);
   static constexpr bool IsFP8Input = detail::is_input_fp8<ElementA, ElementB>();
   static_assert((!IsFP8Input || !IsArrayOfPointersGemm),
-                "KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum is only compatible with FP8 Blocked Scaled version right now.");
+                "KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum is "
+                "only compatible with FP8 Blocked Scaled version right now.");
 
   // For fp32 types, map to tf32 MMA value type
-  using ElementAMma = cute::conditional_t<cute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
-  using ElementBMma = cute::conditional_t<cute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
+  using ElementAMma = cute::
+      conditional_t<cute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
+  using ElementBMma = cute::
+      conditional_t<cute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
 
-  static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_ss_tag_to_major_A<ElementAMma, GmemLayoutATag>();
-  static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_ss_tag_to_major_B<ElementBMma, GmemLayoutBTag>();
+  static constexpr cute::GMMA::Major GmmaMajorA =
+      detail::gmma_ss_tag_to_major_A<ElementAMma, GmemLayoutATag>();
+  static constexpr cute::GMMA::Major GmmaMajorB =
+      detail::gmma_ss_tag_to_major_B<ElementBMma, GmemLayoutBTag>();
 
-  static constexpr bool IsCooperative = cute::is_any_of_v<KernelScheduleType,
-                                                          KernelTmaWarpSpecializedCooperative,
-                                                          KernelPtrArrayTmaWarpSpecializedCooperative,
-                                                          KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>>;
+  static constexpr bool IsCooperative = cute::is_any_of_v<
+      KernelScheduleType,
+      KernelTmaWarpSpecializedCooperative,
+      KernelPtrArrayTmaWarpSpecializedCooperative,
+      KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<
+          ScaleGranularityM>>;
   using AtomLayoutMNK = cute::conditional_t<IsCooperative,
-      Layout<Shape<_2,_1,_1>>, Layout<Shape<_1,_1,_1>>>;
+                                            Layout<Shape<_2, _1, _1>>,
+                                            Layout<Shape<_1, _1, _1>>>;
 
-  using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector<
-      ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{}));
+  using TiledMma = decltype(cute::make_tiled_mma(
+      cute::GMMA::ss_op_selector<ElementAMma,
+                                 ElementBMma,
+                                 ElementAccumulator,
+                                 TileShape_MNK,
+                                 GmmaMajorA,
+                                 GmmaMajorB>(),
+      AtomLayoutMNK{}));
 
-  using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{})));
-  using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{})));
+  using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(
+      shape<1>(ClusterShape_MNK{})));
+  using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(
+      shape<0>(ClusterShape_MNK{})));
 
-  using SmemLayoutAtomA = decltype(detail::ss_smem_selector<
-      GmmaMajorA, ElementAMma, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
-  using SmemLayoutAtomB = decltype(detail::ss_smem_selector<
-      GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutAtomA =
+      decltype(detail::ss_smem_selector<GmmaMajorA,
+                                        ElementAMma,
+                                        decltype(cute::get<0>(TileShape_MNK{})),
+                                        decltype(cute::get<2>(
+                                            TileShape_MNK{}))>());
+  using SmemLayoutAtomB =
+      decltype(detail::ss_smem_selector<GmmaMajorB,
+                                        ElementBMma,
+                                        decltype(cute::get<1>(TileShape_MNK{})),
+                                        decltype(cute::get<2>(
+                                            TileShape_MNK{}))>());
 
-  static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
+  static constexpr size_t TensorMapStorage =
+      IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */
+                            : 0;
   static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage);
 
-  static constexpr int PipelineStages = detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes - KernelSmemCarveout,
-      ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{});
-  using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<PipelineStages, ClusterShape_MNK, KernelScheduleType, ScaleGranularityM>;
+  static constexpr int PipelineStages =
+      detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes -
+                                                  KernelSmemCarveout,
+                                              ElementAMma,
+                                              ElementBMma,
+                                              TileShape_MNK>(StageCountType{});
+  using DispatchPolicy =
+      MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<
+          PipelineStages,
+          ClusterShape_MNK,
+          KernelScheduleType,
+          ScaleGranularityM>;
 
   using SmemCopyAtomA = void;
   using SmemCopyAtomB = void;
 
-  using CollectiveOp = CollectiveMma<
-      DispatchPolicy,
-      TileShape_MNK,
-      ElementA,
-      TagToStrideA_t<GmemLayoutATag>,
-      ElementB,
-      TagToStrideB_t<GmemLayoutBTag>,
-      TiledMma,
-      GmemTiledCopyA,
-      SmemLayoutAtomA,
-      SmemCopyAtomA,
-      cute::identity,
-      GmemTiledCopyB,
-      SmemLayoutAtomB,
-      SmemCopyAtomB,
-      cute::identity
-    >;
+  using CollectiveOp = CollectiveMma<DispatchPolicy,
+                                     TileShape_MNK,
+                                     ElementA,
+                                     TagToStrideA_t<GmemLayoutATag>,
+                                     ElementB,
+                                     TagToStrideB_t<GmemLayoutBTag>,
+                                     TiledMma,
+                                     GmemTiledCopyA,
+                                     SmemLayoutAtomA,
+                                     SmemCopyAtomA,
+                                     cute::identity,
+                                     GmemTiledCopyB,
+                                     SmemLayoutAtomB,
+                                     SmemCopyAtomB,
+                                     cute::identity>;
 };
 
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 }  // namespace cutlass::gemm::collective
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_builder_gated.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_builder_gated.hpp
index 227aee50fe..e3f7e03ef9 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_builder_gated.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_builder_gated.hpp
@@ -39,12 +39,23 @@ namespace cutlass::gemm::collective {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-template <class ArchTag, class OpClass, class ElementA, class GmemLayoutA,
-          int AlignmentA, class ElementB, class GmemLayoutB, int AlignmentB,
-          class ElementAccumulator, class TileShape_MNK, class ClusterShape_MNK,
-          class StageCountType, class KernelScheduleType,
-          template <class /* ElementCompute */> class Activation,
-          bool SwapAB = false, class Enable = void>
+template <class ArchTag,
+          class OpClass,
+          class ElementA,
+          class GmemLayoutA,
+          int AlignmentA,
+          class ElementB,
+          class GmemLayoutB,
+          int AlignmentB,
+          class ElementAccumulator,
+          class TileShape_MNK,
+          class ClusterShape_MNK,
+          class StageCountType,
+          class KernelScheduleType,
+          template <class /* ElementCompute */>
+          class Activation,
+          bool SwapAB = false,
+          class Enable = void>
 struct CollectiveBuilderGated {
   static_assert(sizeof(ElementA) == 0,
                 "Could not build a collective for given parameters.");
@@ -52,7 +63,7 @@ struct CollectiveBuilderGated {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm::collective
+}  // namespace cutlass::gemm::collective
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_mma_gated.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_mma_gated.hpp
index 56849ee56f..d2d06fed31 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_mma_gated.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/collective_mma_gated.hpp
@@ -39,12 +39,23 @@ namespace cutlass::gemm::collective {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-template <class DispatchPolicy, class TileShape, class ElementA, class StrideA,
-          class ElementB, class StrideB, class TiledMma, class GmemTiledCopyA,
-          class SmemLayoutAtomA, class SmemCopyAtomA, class TransformA,
-          class GmemTiledCopyB, class SmemLayoutAtomB, class SmemCopyAtomB,
+template <class DispatchPolicy,
+          class TileShape,
+          class ElementA,
+          class StrideA,
+          class ElementB,
+          class StrideB,
+          class TiledMma,
+          class GmemTiledCopyA,
+          class SmemLayoutAtomA,
+          class SmemCopyAtomA,
+          class TransformA,
+          class GmemTiledCopyB,
+          class SmemLayoutAtomB,
+          class SmemCopyAtomB,
           class TransformB,
-          template <class /* ElementCompute */> class Activation,
+          template <class /* ElementCompute */>
+          class Activation,
           bool SwapAB = false>
 struct CollectiveMmaGated {
   static_assert(cutlass::detail::dependent_false<ElementA>,
@@ -53,7 +64,7 @@ struct CollectiveMmaGated {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm::collective
+}  // namespace cutlass::gemm::collective
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/fp8_accumulation.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/fp8_accumulation.hpp
index 0a530e5c14..b22492fd9e 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/fp8_accumulation.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/fp8_accumulation.hpp
@@ -12,17 +12,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// adapted from: https://github.com/soundOfDestiny/cutlass/blob/a4208aa6958864923505cade9c63eb2a6daf16e5/include/cutlass/gemm/collective/fp8_accumulation.hpp
+// adapted from:
+// https://github.com/soundOfDestiny/cutlass/blob/a4208aa6958864923505cade9c63eb2a6daf16e5/include/cutlass/gemm/collective/fp8_accumulation.hpp
 
 /***************************************************************************************************
- * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -34,14 +35,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
@@ -61,24 +63,26 @@
 
 namespace cutlass::gemm::collective {
 
-template <
-    class EngineAccum,
-    class LayoutAccum>
+template <class EngineAccum, class LayoutAccum>
 struct GmmaFP8AccumulationWithScale {
   using TensorAccum = cute::Tensor<EngineAccum, LayoutAccum>;
   using ElementAccumulator = typename EngineAccum::value_type;
 
-  static_assert(is_static<LayoutAccum>::value, "Accumulator Layout should be static");
-  static_assert(is_rmem<TensorAccum>::value , "Accumulator tensor must be rmem resident.");
+  static_assert(is_static<LayoutAccum>::value,
+                "Accumulator Layout should be static");
+  static_assert(is_rmem<TensorAccum>::value,
+                "Accumulator tensor must be rmem resident.");
 
-private:
-  TensorAccum& accum_;
+ private:
+  TensorAccum &accum_;
   TensorAccum accum_temp_;
 
-  uint32_t accum_promotion_interval_;         // defines the max num of executed MMAs after which accum should be promoted.
-  uint32_t mma_count_per_mainloop_iteration_; // num of MMAs per k_tile of mainloop
-  uint32_t mma_count_;                        // current executed MMAs
-  uint32_t reset_accum_flag_;                 // accum needs to be zeroed or not.
+  uint32_t accum_promotion_interval_;  // defines the max num of executed MMAs
+                                       // after which accum should be promoted.
+  uint32_t
+      mma_count_per_mainloop_iteration_;  // num of MMAs per k_tile of mainloop
+  uint32_t mma_count_;                    // current executed MMAs
+  uint32_t reset_accum_flag_;             // accum needs to be zeroed or not.
 
   // promote or `add` the partial accumulators to main accumulator (FADD).
   CUTLASS_DEVICE
@@ -90,18 +94,20 @@ private:
     }
   }
 
-  // `multiply` scale the partial accumulators and `add` to main accumulator (FFMA).
-  template <
-    class EngineScale,
-    class LayoutScale>
-  CUTLASS_DEVICE
-  void scale_core(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+  // `multiply` scale the partial accumulators and `add` to main accumulator
+  // (FFMA).
+  template <class EngineScale, class LayoutScale>
+  CUTLASS_DEVICE void scale_core(
+      const cute::Tensor<EngineScale, LayoutScale> &scale) {
     using TensorScale = cute::Tensor<EngineScale, LayoutScale>;
 
-    static_assert(is_static<LayoutScale>::value, "Scale Layout should be static");
-    static_assert(is_rmem<TensorScale>::value , "Scale tensor must be rmem resident.");
+    static_assert(is_static<LayoutScale>::value,
+                  "Scale Layout should be static");
+    static_assert(is_rmem<TensorScale>::value,
+                  "Scale tensor must be rmem resident.");
 
-    static_assert(LayoutAccum{}.shape() == LayoutScale{}.shape(), "Accumulator and scale must have same shape.");
+    static_assert(LayoutAccum{}.shape() == LayoutScale{}.shape(),
+                  "Accumulator and scale must have same shape.");
 
     warpgroup_wait<0>();
     CUTLASS_PRAGMA_UNROLL
@@ -110,18 +116,16 @@ private:
     }
   }
 
-public:
+ public:
   CUTLASS_DEVICE
-  GmmaFP8AccumulationWithScale(
-      TensorAccum &accum,
-      uint32_t accum_promotion_interval,
-      uint32_t mma_count_per_mainloop_iteration)
+  GmmaFP8AccumulationWithScale(TensorAccum &accum,
+                               uint32_t accum_promotion_interval,
+                               uint32_t mma_count_per_mainloop_iteration)
       : accum_(accum),
         accum_promotion_interval_(accum_promotion_interval),
         mma_count_per_mainloop_iteration_(mma_count_per_mainloop_iteration),
         mma_count_(0),
-        reset_accum_flag_(0)
-  {
+        reset_accum_flag_(0) {
     accum_temp_ = cute::make_fragment_like(accum);
   }
 
@@ -130,32 +134,31 @@ public:
   //
 
   CUTLASS_DEVICE
-  TensorAccum& operator()() {
-    return accum_temp_;
-  }
+  TensorAccum &operator()() { return accum_temp_; }
 
   /// prepare the MMA accumulators when initialization or zeroing is required.
   CUTLASS_DEVICE
-  bool prepare_if_needed() {
-    return reset_accum_flag_;
-  }
+  bool prepare_if_needed() { return reset_accum_flag_; }
 
   //
   // Methods (for FADD version)
   //
 
-  /// promote (add) the results from the MMA accumulators to main accumulator if needed.
+  /// promote (add) the results from the MMA accumulators to main accumulator if
+  /// needed.
   CUTLASS_DEVICE
   void promote_if_needed() {
     mma_count_ += mma_count_per_mainloop_iteration_;
-    reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
+    reset_accum_flag_ =
+        __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
     if (reset_accum_flag_) {
       promote_core();
       mma_count_ = 0;
     }
   }
 
-  /// promote (add) the residue results from the MMA accumulators to main accumulator if needed.
+  /// promote (add) the residue results from the MMA accumulators to main
+  /// accumulator if needed.
   CUTLASS_DEVICE
   void promote_residue_if_needed() {
     if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
@@ -167,30 +170,29 @@ public:
   // Methods (for FFMA version)
   //
 
-  /// scale (multiply_add) the results from the MMA accumulators to main accumulator if needed.
-  template <
-    class EngineScale,
-    class LayoutScale>
-  CUTLASS_DEVICE
-  void scale_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+  /// scale (multiply_add) the results from the MMA accumulators to main
+  /// accumulator if needed.
+  template <class EngineScale, class LayoutScale>
+  CUTLASS_DEVICE void scale_if_needed(
+      const cute::Tensor<EngineScale, LayoutScale> &scale) {
     mma_count_ += mma_count_per_mainloop_iteration_;
-    reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
+    reset_accum_flag_ =
+        __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
     if (reset_accum_flag_) {
       scale_core(scale);
       mma_count_ = 0;
     }
   }
 
-  /// scale (multiply_add) the residue results from the MMA accumulators to main accumulator if needed.
-  template <
-    class EngineScale,
-    class LayoutScale>
-  CUTLASS_DEVICE
-  void scale_residue_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+  /// scale (multiply_add) the residue results from the MMA accumulators to main
+  /// accumulator if needed.
+  template <class EngineScale, class LayoutScale>
+  CUTLASS_DEVICE void scale_residue_if_needed(
+      const cute::Tensor<EngineScale, LayoutScale> &scale) {
     if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
       scale_core(scale);
     }
   }
 };
 
-} // namespace cutlass::gemm::collective
+}  // namespace cutlass::gemm::collective
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_gated_tma_gmma_ss_warpspecialized.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_gated_tma_gmma_ss_warpspecialized.hpp
index 8ff14a2a49..f335ec2d39 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_gated_tma_gmma_ss_warpspecialized.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_gated_tma_gmma_ss_warpspecialized.hpp
@@ -53,18 +53,43 @@ using namespace cute;
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 // WarpSpecialized Mainloop
-template <int Stages, class ClusterShape, class KernelSchedule,
-          class TileShape_, class ElementA_, class StrideA_, class ElementB_,
-          class StrideB_, class TiledMma_, class GmemTiledCopyA_,
-          class SmemLayoutAtomA_, class SmemCopyAtomA_, class TransformA_,
-          class GmemTiledCopyB_, class SmemLayoutAtomB_, class SmemCopyAtomB_,
+template <int Stages,
+          class ClusterShape,
+          class KernelSchedule,
+          class TileShape_,
+          class ElementA_,
+          class StrideA_,
+          class ElementB_,
+          class StrideB_,
+          class TiledMma_,
+          class GmemTiledCopyA_,
+          class SmemLayoutAtomA_,
+          class SmemCopyAtomA_,
+          class TransformA_,
+          class GmemTiledCopyB_,
+          class SmemLayoutAtomB_,
+          class SmemCopyAtomB_,
           class TransformB_,
-          template <class /* ElementCompute */> class Activation_, bool SwapAB_>
+          template <class /* ElementCompute */>
+          class Activation_,
+          bool SwapAB_>
 struct CollectiveMmaGated<
     MainloopSm90TmaGmmaWarpSpecialized<Stages, ClusterShape, KernelSchedule>,
-    TileShape_, ElementA_, StrideA_, ElementB_, StrideB_, TiledMma_,
-    GmemTiledCopyA_, SmemLayoutAtomA_, SmemCopyAtomA_, TransformA_,
-    GmemTiledCopyB_, SmemLayoutAtomB_, SmemCopyAtomB_, TransformB_, Activation_,
+    TileShape_,
+    ElementA_,
+    StrideA_,
+    ElementB_,
+    StrideB_,
+    TiledMma_,
+    GmemTiledCopyA_,
+    SmemLayoutAtomA_,
+    SmemCopyAtomA_,
+    TransformA_,
+    GmemTiledCopyB_,
+    SmemLayoutAtomB_,
+    SmemCopyAtomB_,
+    TransformB_,
+    Activation_,
     SwapAB_> {
   static constexpr bool isGated = true;
   static constexpr bool SwapAB = SwapAB_;
@@ -93,7 +118,8 @@ struct CollectiveMmaGated<
   using Activation = Activation_<ElementAccumulator>;
 
   using ElementAux = cute::conditional_t<SwapAB, ElementA_, ElementB_>;
-  using ValTypeAux = cute::conditional_t<SwapAB, typename TiledMma::ValTypeA,
+  using ValTypeAux = cute::conditional_t<SwapAB,
+                                         typename TiledMma::ValTypeA,
                                          typename TiledMma::ValTypeB>;
 
   using MainloopPipeline = cutlass::PipelineTmaAsync<DispatchPolicy::Stages>;
@@ -118,16 +144,20 @@ struct CollectiveMmaGated<
   // Tile along modes in a way that maximizes the TMA box size.
   using SmemLayoutA = decltype(tile_to_shape(
       SmemLayoutAtomA{},
-      make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}),
+      make_shape(shape<0>(TileShape{}),
+                 shape<2>(TileShape{}),
                  Int<DispatchPolicy::Stages>{}),
       conditional_t<::cutlass::gemm::detail::is_major<0, StrideA>(),
-                    Step<_2, _1, _3>, Step<_1, _2, _3>>{}));
+                    Step<_2, _1, _3>,
+                    Step<_1, _2, _3>>{}));
   using SmemLayoutB = decltype(tile_to_shape(
       SmemLayoutAtomB{},
-      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}),
+      make_shape(shape<1>(TileShape{}),
+                 shape<2>(TileShape{}),
                  Int<DispatchPolicy::Stages>{}),
       conditional_t<::cutlass::gemm::detail::is_major<0, StrideB>(),
-                    Step<_2, _1, _3>, Step<_1, _2, _3>>{}));
+                    Step<_2, _1, _3>,
+                    Step<_1, _2, _3>>{}));
   using SmemLayoutAux = cute::conditional_t<SwapAB, SmemLayoutA, SmemLayoutB>;
 
   static_assert(DispatchPolicy::Stages >= 2,
@@ -151,10 +181,12 @@ struct CollectiveMmaGated<
   static constexpr bool ConvertF32toTF32A = cute::is_same_v<float, ElementA>;
   static constexpr bool ConvertF32toTF32B = cute::is_same_v<float, ElementB>;
   using InternalElementA =
-      cute::conditional_t<ConvertF32toTF32A, tfloat32_t,
+      cute::conditional_t<ConvertF32toTF32A,
+                          tfloat32_t,
                           uint_bit_t<sizeof_bits_v<ElementA>>>;
   using InternalElementB =
-      cute::conditional_t<ConvertF32toTF32B, tfloat32_t,
+      cute::conditional_t<ConvertF32toTF32B,
+                          tfloat32_t,
                           uint_bit_t<sizeof_bits_v<ElementB>>>;
   using InternalElementAux =
       cute::conditional_t<SwapAB, InternalElementA, InternalElementB>;
@@ -195,18 +227,22 @@ struct CollectiveMmaGated<
     using TMA_A = decltype(make_tma_copy(
         GmemTiledCopyA{},
         make_tensor(static_cast<InternalElementA const *>(nullptr),
-                    repeat_like(StrideA{}, int32_t(0)), StrideA{}),
+                    repeat_like(StrideA{}, int32_t(0)),
+                    StrideA{}),
         SmemLayoutA{}(_, _, cute::Int<0>{}),
         make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
-        size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any
+        size<1>(
+            ClusterShape{})));  // mcast along N mode for this M load, if any
     // Assumption: StrideB is congruent with Problem_NK
     using TMA_B = decltype(make_tma_copy(
         GmemTiledCopyB{},
         make_tensor(static_cast<InternalElementB const *>(nullptr),
-                    repeat_like(StrideB{}, int32_t(0)), StrideB{}),
+                    repeat_like(StrideB{}, int32_t(0)),
+                    StrideB{}),
         SmemLayoutB{}(_, _, cute::Int<0>{}),
         make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
-        size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any
+        size<0>(
+            ClusterShape{})));  // mcast along M mode for this N load, if any
     using TMA_Aux = cute::conditional_t<SwapAB, TMA_A, TMA_B>;
     TMA_A tma_load_a;
     TMA_B tma_load_b;
@@ -220,9 +256,10 @@ struct CollectiveMmaGated<
   //
 
   template <class ProblemShape>
-  static constexpr Params
-  to_underlying_arguments(ProblemShape const &problem_shape,
-                          Arguments const &args, void *workspace) {
+  static constexpr Params to_underlying_arguments(
+      ProblemShape const &problem_shape,
+      Arguments const &args,
+      void *workspace) {
     (void)workspace;
 
     // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is
@@ -238,36 +275,44 @@ struct CollectiveMmaGated<
     Tensor tensor_b =
         make_tensor(ptr_B0, make_layout(make_shape(N, K, L), args.dB));
     typename Params::TMA_A tma_load_a = make_tma_copy(
-        GmemTiledCopyA{}, tensor_a, SmemLayoutA{}(_, _, cute::Int<0>{}),
+        GmemTiledCopyA{},
+        tensor_a,
+        SmemLayoutA{}(_, _, cute::Int<0>{}),
         make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
-        size<1>(ClusterShape{})); // mcast along N mode for this M load, if any
+        size<1>(ClusterShape{}));  // mcast along N mode for this M load, if any
     typename Params::TMA_B tma_load_b = make_tma_copy(
-        GmemTiledCopyB{}, tensor_b, SmemLayoutB{}(_, _, cute::Int<0>{}),
+        GmemTiledCopyB{},
+        tensor_b,
+        SmemLayoutB{}(_, _, cute::Int<0>{}),
         make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
-        size<0>(ClusterShape{})); // mcast along M mode for this N load, if any
+        size<0>(ClusterShape{}));  // mcast along M mode for this N load, if any
 
     if constexpr (SwapAB) {
       auto ptr_Aux = reinterpret_cast<InternalElementA const *>(args.ptr_B1);
       Tensor tensor_aux =
           make_tensor(ptr_Aux, make_layout(make_shape(M, K, L), args.dA));
       typename Params::TMA_Aux tma_load_aux = make_tma_copy(
-          GmemTiledCopyA{}, tensor_aux, SmemLayoutA{}(_, _, cute::Int<0>{}),
+          GmemTiledCopyA{},
+          tensor_aux,
+          SmemLayoutA{}(_, _, cute::Int<0>{}),
           make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
           size<1>(
-              ClusterShape{})); // mcast along N mode for this M load, if any
-      return {tma_load_a, tma_load_b, tma_load_aux, args.scale_d0,
-              args.scale_d1};
+              ClusterShape{}));  // mcast along N mode for this M load, if any
+      return {
+          tma_load_a, tma_load_b, tma_load_aux, args.scale_d0, args.scale_d1};
     } else {
       auto ptr_Aux = reinterpret_cast<InternalElementB const *>(args.ptr_B1);
       Tensor tensor_aux =
           make_tensor(ptr_Aux, make_layout(make_shape(N, K, L), args.dB));
       typename Params::TMA_Aux tma_load_aux = make_tma_copy(
-          GmemTiledCopyB{}, tensor_aux, SmemLayoutB{}(_, _, cute::Int<0>{}),
+          GmemTiledCopyB{},
+          tensor_aux,
+          SmemLayoutB{}(_, _, cute::Int<0>{}),
           make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
           size<0>(
-              ClusterShape{})); // mcast along M mode for this N load, if any
-      return {tma_load_a, tma_load_b, tma_load_aux, args.scale_d0,
-              args.scale_d1};
+              ClusterShape{}));  // mcast along M mode for this N load, if any
+      return {
+          tma_load_a, tma_load_b, tma_load_aux, args.scale_d0, args.scale_d1};
     }
   }
 
@@ -293,8 +338,9 @@ struct CollectiveMmaGated<
             cute::make_shape(N, K, L), StrideB{});
 
     if (!implementable) {
-      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the "
-                         "minimum alignment requirements for TMA.\n");
+      CUTLASS_TRACE_HOST(
+          "  CAN IMPLEMENT: Problem Size doesn't meet the "
+          "minimum alignment requirements for TMA.\n");
     }
     return implementable;
   }
@@ -342,49 +388,64 @@ struct CollectiveMmaGated<
     // TMA requires special handling of strides to deal with coord codomain
     // mapping Represent the full tensors -- get these from TMA
     Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(
-        make_shape(M, K, L)); // (m,k,l)
+        make_shape(M, K, L));  // (m,k,l)
     Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(
-        make_shape(N, K, L)); // (n,k,l)
+        make_shape(N, K, L));  // (n,k,l)
 
     // Make tiled views, defer the slice
-    Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_, _, _),
-                               Step<_1, X, _1>{}); // (BLK_M,BLK_K,m,k,l)
-    Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_, _, _),
-                               Step<X, _1, _1>{}); // (BLK_N,BLK_K,n,k,l)
+    Tensor gA_mkl = local_tile(mA_mkl,
+                               TileShape{},
+                               make_coord(_, _, _),
+                               Step<_1, X, _1>{});  // (BLK_M,BLK_K,m,k,l)
+    Tensor gB_nkl = local_tile(mB_nkl,
+                               TileShape{},
+                               make_coord(_, _, _),
+                               Step<X, _1, _1>{});  // (BLK_N,BLK_K,n,k,l)
 
     if constexpr (SwapAB) {
       Tensor mAux_xkl = mainloop_params.tma_load_aux.get_tma_tensor(
-          make_shape(M, K, L)); // (m,k,l)
-      Tensor gAux_xkl = local_tile(mAux_xkl, TileShape{}, make_coord(_, _, _),
-                                   Step<_1, X, _1>{}); // (BLK_M,BLK_K,m,k,l)
+          make_shape(M, K, L));  // (m,k,l)
+      Tensor gAux_xkl = local_tile(mAux_xkl,
+                                   TileShape{},
+                                   make_coord(_, _, _),
+                                   Step<_1, X, _1>{});  // (BLK_M,BLK_K,m,k,l)
       return cute::make_tuple(gA_mkl, gB_nkl, gAux_xkl);
     } else {
       Tensor mAux_xkl = mainloop_params.tma_load_aux.get_tma_tensor(
-          make_shape(N, K, L)); // (n,k,l)
-      Tensor gAux_xkl = local_tile(mAux_xkl, TileShape{}, make_coord(_, _, _),
-                                   Step<X, _1, _1>{}); // (BLK_N,BLK_K,n,k,l)
+          make_shape(N, K, L));  // (n,k,l)
+      Tensor gAux_xkl = local_tile(mAux_xkl,
+                                   TileShape{},
+                                   make_coord(_, _, _),
+                                   Step<X, _1, _1>{});  // (BLK_N,BLK_K,n,k,l)
       return cute::make_tuple(gA_mkl, gB_nkl, gAux_xkl);
     }
   }
 
   /// Perform a collective-scoped matrix multiply-accumulate
   /// Producer Perspective
-  template <class TensorA, class TensorB, class TensorAux, class KTileIterator,
+  template <class TensorA,
+            class TensorB,
+            class TensorAux,
+            class KTileIterator,
             class BlockCoord>
-  CUTLASS_DEVICE void
-  load(Params const &mainloop_params, MainloopPipeline pipeline,
-       PipelineState smem_pipe_write,
-       cute::tuple<TensorA, TensorB, TensorAux> const &load_inputs,
-       BlockCoord const &blk_coord, KTileIterator k_tile_iter, int k_tile_count,
-       int thread_idx, uint32_t block_rank_in_cluster,
-       TensorStorage &shared_tensors) {
+  CUTLASS_DEVICE void load(
+      Params const &mainloop_params,
+      MainloopPipeline pipeline,
+      PipelineState smem_pipe_write,
+      cute::tuple<TensorA, TensorB, TensorAux> const &load_inputs,
+      BlockCoord const &blk_coord,
+      KTileIterator k_tile_iter,
+      int k_tile_count,
+      int thread_idx,
+      uint32_t block_rank_in_cluster,
+      TensorStorage &shared_tensors) {
     int lane_predicate = cute::elect_one_sync();
 
     if (lane_predicate) {
       Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()),
-                              SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
+                              SmemLayoutA{});  // (BLK_M,BLK_K,PIPE)
       Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
-                              SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
+                              SmemLayoutB{});  // (BLK_N,BLK_K,PIPE)
       Tensor sAux = make_tensor(make_smem_ptr(shared_tensors.smem_Aux.data()),
                                 SmemLayoutAux{});
 
@@ -412,17 +473,17 @@ struct CollectiveMmaGated<
                     cluster_local_block_id.x);
       // Partition the inputs based on the current block coordinates.
       auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
-      Tensor gA = gA_mkl(_, _, m_coord, _, l_coord); // (BLK_M,BLK_K,k)
-      Tensor gB = gB_nkl(_, _, n_coord, _, l_coord); // (BLK_N,BLK_K,k)
+      Tensor gA = gA_mkl(_, _, m_coord, _, l_coord);  // (BLK_M,BLK_K,k)
+      Tensor gB = gB_nkl(_, _, n_coord, _, l_coord);  // (BLK_N,BLK_K,k)
       Tensor gAux = SwapAB ? gAux_xkl(_, _, m_coord, _, l_coord)
                            : gAux_xkl(_, _, n_coord, _, l_coord);
 
       // Applies the mapping from block_tma_a
-      Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
-      Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE)
+      Tensor tAgA = block_tma_a.partition_S(gA);  // (TMA,TMA_M,TMA_K,k)
+      Tensor tAsA = block_tma_a.partition_D(sA);  // (TMA,TMA_M,TMA_K,PIPE)
 
-      Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k)
-      Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE)
+      Tensor tBgB = block_tma_b.partition_S(gB);  // (TMA,TMA_N,TMA_K,k)
+      Tensor tBsB = block_tma_b.partition_D(sB);  // (TMA,TMA_N,TMA_K,PIPE)
 
       Tensor tAuxgAux = block_tma_aux.partition_S(gAux);
       Tensor tAuxsAux = block_tma_aux.partition_D(sAux);
@@ -435,18 +496,18 @@ struct CollectiveMmaGated<
       // Maps the tile -> block, value
       if constexpr (cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>) {
         auto block_layout =
-            Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) ->
-                                                             // block_id
+            Layout<typename DispatchPolicy::ClusterShape>{};  // (m,n) ->
+                                                              // block_id
         for (int n = 0; n < size<1>(block_layout); ++n) {
-          mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,
-                                                       n, Int<0>{}));
+          mcast_mask_a |= (uint16_t(1) << block_layout(
+                               cluster_local_block_id.x, n, Int<0>{}));
         }
       }
 
       if constexpr (cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>) {
         auto block_layout =
-            Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) ->
-                                                             // block_id
+            Layout<typename DispatchPolicy::ClusterShape>{};  // (m,n) ->
+                                                              // block_id
         for (int m = 0; m < size<0>(block_layout); ++m) {
           mcast_mask_b |= (uint16_t(1) << block_layout(
                                m, cluster_local_block_id.y, Int<0>{}));
@@ -475,11 +536,14 @@ struct CollectiveMmaGated<
 
         int write_stage = smem_pipe_write.index();
         copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a),
-             tAgA(_, _, _, *k_tile_iter), tAsA(_, _, _, write_stage));
+             tAgA(_, _, _, *k_tile_iter),
+             tAsA(_, _, _, write_stage));
         copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b),
-             tBgB(_, _, _, *k_tile_iter), tBsB(_, _, _, write_stage));
+             tBgB(_, _, _, *k_tile_iter),
+             tBsB(_, _, _, write_stage));
         copy(mainloop_params.tma_load_aux.with(*tma_barrier, mcast_mask_aux),
-             tAuxgAux(_, _, _, *k_tile_iter), tAuxsAux(_, _, _, write_stage));
+             tAuxgAux(_, _, _, *k_tile_iter),
+             tAuxsAux(_, _, _, write_stage));
         ++k_tile_iter;
 
         // Advance smem_pipe_write
@@ -508,10 +572,14 @@ struct CollectiveMmaGated<
   /// Perform a collective-scoped matrix multiply-accumulate
   /// Consumer Perspective
   template <class FrgTensorC>
-  CUTLASS_DEVICE void
-  mma(MainloopPipeline pipeline, PipelineState smem_pipe_read,
-      FrgTensorC &accum0, FrgTensorC &accum1, int k_tile_count, int thread_idx,
-      TensorStorage &shared_tensors, Params const &mainloop_params) {
+  CUTLASS_DEVICE void mma(MainloopPipeline pipeline,
+                          PipelineState smem_pipe_read,
+                          FrgTensorC &accum0,
+                          FrgTensorC &accum1,
+                          int k_tile_count,
+                          int thread_idx,
+                          TensorStorage &shared_tensors,
+                          Params const &mainloop_params) {
     static_assert(is_rmem<FrgTensorC>::value,
                   "C tensor must be rmem resident.");
     static_assert(cute::rank(SmemLayoutA{}) == 3,
@@ -528,9 +596,9 @@ struct CollectiveMmaGated<
                   "smem sourced instructions.");
 
     Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()),
-                            SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
+                            SmemLayoutA{});  // (BLK_M,BLK_K,PIPE)
     Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
-                            SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
+                            SmemLayoutB{});  // (BLK_N,BLK_K,PIPE)
     Tensor sAux = make_tensor(make_smem_ptr(shared_tensors.smem_Aux.data()),
                               SmemLayoutAux{});
 
@@ -541,12 +609,12 @@ struct CollectiveMmaGated<
     TiledMma tiled_mma;
     auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
 
-    Tensor tCsA = thread_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
-    Tensor tCsB = thread_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
+    Tensor tCsA = thread_mma.partition_A(sA);  // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCsB = thread_mma.partition_B(sB);  // (MMA,MMA_N,MMA_K,PIPE)
 
     // Allocate "fragments/descriptors"
-    Tensor tCrA = thread_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)
-    Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
+    Tensor tCrA = thread_mma.make_fragment_A(tCsA);  // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCrB = thread_mma.make_fragment_B(tCsB);  // (MMA,MMA_N,MMA_K,PIPE)
 
     auto tCsAux = [&]() -> auto {
       if constexpr (SwapAB) {
@@ -554,34 +622,36 @@ struct CollectiveMmaGated<
       } else {
         return thread_mma.partition_B(sAux);
       }
-    }();
+    }
+    ();
     auto tCrAux = [&]() -> auto {
       if constexpr (SwapAB) {
         return thread_mma.make_fragment_A(tCsAux);
       } else {
         return thread_mma.make_fragment_B(tCsAux);
       }
-    }();
-
-    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum0)); // M
-    CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum0)); // N
-    CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));   // K
-    CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));   // PIPE
-    if constexpr (SwapAB) {
-      CUTE_STATIC_ASSERT_V(size<1>(tCsAux) == size<1>(accum1)); // M
-      CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum1));   // N
-      CUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCsAux));   // K
-      CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tCsAux));   // PIPE
-    } else {
-      CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum1));   // M
-      CUTE_STATIC_ASSERT_V(size<1>(tCsAux) == size<2>(accum1)); // N
-      CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsAux));   // K
-      CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsAux));   // PIPE
     }
-    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
-    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
+    ();
+
+    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum0));  // M
+    CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum0));  // N
+    CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));    // K
+    CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));    // PIPE
+    if constexpr (SwapAB) {
+      CUTE_STATIC_ASSERT_V(size<1>(tCsAux) == size<1>(accum1));  // M
+      CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum1));    // N
+      CUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCsAux));    // K
+      CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tCsAux));    // PIPE
+    } else {
+      CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum1));    // M
+      CUTE_STATIC_ASSERT_V(size<1>(tCsAux) == size<2>(accum1));  // N
+      CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsAux));    // K
+      CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsAux));    // PIPE
+    }
+    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA));  // PIPE
+    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB));  // PIPE
     CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} ==
-                         size<2>(sAux)); // PIPE
+                         size<2>(sAux));  // PIPE
 
     //
     // PIPELINED MAIN LOOP
@@ -613,14 +683,20 @@ struct CollectiveMmaGated<
       CUTLASS_PRAGMA_UNROLL
       for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
         // (V,M,K) x (V,N,K) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA(_, _, k_block, read_stage),
-                   tCrB(_, _, k_block, read_stage), accum0);
+        cute::gemm(tiled_mma,
+                   tCrA(_, _, k_block, read_stage),
+                   tCrB(_, _, k_block, read_stage),
+                   accum0);
         if constexpr (SwapAB) {
-          cute::gemm(tiled_mma, tCrAux(_, _, k_block, read_stage),
-                     tCrB(_, _, k_block, read_stage), accum1);
+          cute::gemm(tiled_mma,
+                     tCrAux(_, _, k_block, read_stage),
+                     tCrB(_, _, k_block, read_stage),
+                     accum1);
         } else {
-          cute::gemm(tiled_mma, tCrA(_, _, k_block, read_stage),
-                     tCrAux(_, _, k_block, read_stage), accum1);
+          cute::gemm(tiled_mma,
+                     tCrA(_, _, k_block, read_stage),
+                     tCrAux(_, _, k_block, read_stage),
+                     accum1);
         }
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
       }
@@ -654,14 +730,20 @@ struct CollectiveMmaGated<
       CUTLASS_PRAGMA_UNROLL
       for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
         // (V,M,K) x (V,N,K) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA(_, _, k_block, read_stage),
-                   tCrB(_, _, k_block, read_stage), accum0);
+        cute::gemm(tiled_mma,
+                   tCrA(_, _, k_block, read_stage),
+                   tCrB(_, _, k_block, read_stage),
+                   accum0);
         if constexpr (SwapAB) {
-          cute::gemm(tiled_mma, tCrAux(_, _, k_block, read_stage),
-                     tCrB(_, _, k_block, read_stage), accum1);
+          cute::gemm(tiled_mma,
+                     tCrAux(_, _, k_block, read_stage),
+                     tCrB(_, _, k_block, read_stage),
+                     accum1);
         } else {
-          cute::gemm(tiled_mma, tCrA(_, _, k_block, read_stage),
-                     tCrAux(_, _, k_block, read_stage), accum1);
+          cute::gemm(tiled_mma,
+                     tCrA(_, _, k_block, read_stage),
+                     tCrAux(_, _, k_block, read_stage),
+                     accum1);
         }
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
       }
@@ -699,8 +781,9 @@ struct CollectiveMmaGated<
     warpgroup_wait<0>();
 
     for (int count = 0; count < prologue_mma_count; ++count) {
-      pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release,
-                                                    // done _computing_ on it
+      pipeline.consumer_release(
+          smem_pipe_release);  // UNLOCK smem_pipe_release,
+                               // done _computing_ on it
       ++smem_pipe_release;
     }
   }
@@ -708,6 +791,6 @@ struct CollectiveMmaGated<
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm::collective
+}  // namespace cutlass::gemm::collective
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_gated_tma_gmma_ss_warpspecialized_fp8.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_gated_tma_gmma_ss_warpspecialized_fp8.hpp
index 76ffbdb2e6..c34ad242e2 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_gated_tma_gmma_ss_warpspecialized_fp8.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_gated_tma_gmma_ss_warpspecialized_fp8.hpp
@@ -55,18 +55,43 @@ using namespace cute;
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 // WarpSpecialized Mainloop
-template <int Stages, class ClusterShape, class KernelSchedule,
-          class TileShape_, class ElementA_, class StrideA_, class ElementB_,
-          class StrideB_, class TiledMma_, class GmemTiledCopyA_,
-          class SmemLayoutAtomA_, class SmemCopyAtomA_, class TransformA_,
-          class GmemTiledCopyB_, class SmemLayoutAtomB_, class SmemCopyAtomB_,
+template <int Stages,
+          class ClusterShape,
+          class KernelSchedule,
+          class TileShape_,
+          class ElementA_,
+          class StrideA_,
+          class ElementB_,
+          class StrideB_,
+          class TiledMma_,
+          class GmemTiledCopyA_,
+          class SmemLayoutAtomA_,
+          class SmemCopyAtomA_,
+          class TransformA_,
+          class GmemTiledCopyB_,
+          class SmemLayoutAtomB_,
+          class SmemCopyAtomB_,
           class TransformB_,
-          template <class /* ElementCompute */> class Activation_, bool SwapAB_>
+          template <class /* ElementCompute */>
+          class Activation_,
+          bool SwapAB_>
 struct CollectiveMmaGated<
     MainloopSm90TmaGmmaWarpSpecializedFP8<Stages, ClusterShape, KernelSchedule>,
-    TileShape_, ElementA_, StrideA_, ElementB_, StrideB_, TiledMma_,
-    GmemTiledCopyA_, SmemLayoutAtomA_, SmemCopyAtomA_, TransformA_,
-    GmemTiledCopyB_, SmemLayoutAtomB_, SmemCopyAtomB_, TransformB_, Activation_,
+    TileShape_,
+    ElementA_,
+    StrideA_,
+    ElementB_,
+    StrideB_,
+    TiledMma_,
+    GmemTiledCopyA_,
+    SmemLayoutAtomA_,
+    SmemCopyAtomA_,
+    TransformA_,
+    GmemTiledCopyB_,
+    SmemLayoutAtomB_,
+    SmemCopyAtomB_,
+    TransformB_,
+    Activation_,
     SwapAB_> {
   static constexpr bool isGated = true;
   static constexpr bool SwapAB = SwapAB_;
@@ -74,9 +99,9 @@ struct CollectiveMmaGated<
   //
   // Type Aliases
   //
-  using DispatchPolicy =
-      MainloopSm90TmaGmmaWarpSpecializedFP8<Stages, ClusterShape,
-                                            KernelSchedule>;
+  using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedFP8<Stages,
+                                                               ClusterShape,
+                                                               KernelSchedule>;
   using TileShape = TileShape_;
   using ElementA = ElementA_;
   using StrideA = StrideA_;
@@ -96,7 +121,8 @@ struct CollectiveMmaGated<
   using Activation = Activation_<ElementAccumulator>;
 
   using ElementAux = cute::conditional_t<SwapAB, ElementA_, ElementB_>;
-  using ValTypeAux = cute::conditional_t<SwapAB, typename TiledMma::ValTypeA,
+  using ValTypeAux = cute::conditional_t<SwapAB,
+                                         typename TiledMma::ValTypeA,
                                          typename TiledMma::ValTypeB>;
 
   using MainloopPipeline = cutlass::PipelineTmaAsync<DispatchPolicy::Stages>;
@@ -121,16 +147,20 @@ struct CollectiveMmaGated<
   // Tile along modes in a way that maximizes the TMA box size.
   using SmemLayoutA = decltype(tile_to_shape(
       SmemLayoutAtomA{},
-      make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}),
+      make_shape(shape<0>(TileShape{}),
+                 shape<2>(TileShape{}),
                  Int<DispatchPolicy::Stages>{}),
       conditional_t<::cutlass::gemm::detail::is_major<0, StrideA>(),
-                    Step<_2, _1, _3>, Step<_1, _2, _3>>{}));
+                    Step<_2, _1, _3>,
+                    Step<_1, _2, _3>>{}));
   using SmemLayoutB = decltype(tile_to_shape(
       SmemLayoutAtomB{},
-      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}),
+      make_shape(shape<1>(TileShape{}),
+                 shape<2>(TileShape{}),
                  Int<DispatchPolicy::Stages>{}),
       conditional_t<::cutlass::gemm::detail::is_major<0, StrideB>(),
-                    Step<_2, _1, _3>, Step<_1, _2, _3>>{}));
+                    Step<_2, _1, _3>,
+                    Step<_1, _2, _3>>{}));
   using SmemLayoutAux = cute::conditional_t<SwapAB, SmemLayoutA, SmemLayoutB>;
 
   static_assert(DispatchPolicy::Stages >= 2,
@@ -184,18 +214,22 @@ struct CollectiveMmaGated<
     using TMA_A = decltype(make_tma_copy(
         GmemTiledCopyA{},
         make_tensor(static_cast<ElementA const *>(nullptr),
-                    repeat_like(StrideA{}, int32_t(0)), StrideA{}),
+                    repeat_like(StrideA{}, int32_t(0)),
+                    StrideA{}),
         SmemLayoutA{}(_, _, 0),
         make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
-        size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any
+        size<1>(
+            ClusterShape{})));  // mcast along N mode for this M load, if any
     // Assumption: StrideB is congruent with Problem_NK
     using TMA_B = decltype(make_tma_copy(
         GmemTiledCopyB{},
         make_tensor(static_cast<ElementB const *>(nullptr),
-                    repeat_like(StrideB{}, int32_t(0)), StrideB{}),
+                    repeat_like(StrideB{}, int32_t(0)),
+                    StrideB{}),
         SmemLayoutB{}(_, _, 0),
         make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
-        size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any
+        size<0>(
+            ClusterShape{})));  // mcast along M mode for this N load, if any
     using TMA_Aux = cute::conditional_t<SwapAB, TMA_A, TMA_B>;
     TMA_A tma_load_a;
     TMA_B tma_load_b;
@@ -210,9 +244,10 @@ struct CollectiveMmaGated<
   //
 
   template <class ProblemShape>
-  static constexpr Params
-  to_underlying_arguments(ProblemShape const &problem_shape,
-                          Arguments const &args, void *workspace) {
+  static constexpr Params to_underlying_arguments(
+      ProblemShape const &problem_shape,
+      Arguments const &args,
+      void *workspace) {
     (void)workspace;
 
     // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is
@@ -228,35 +263,51 @@ struct CollectiveMmaGated<
     Tensor tensor_b =
         make_tensor(ptr_B0, make_layout(make_shape(N, K, L), args.dB));
     typename Params::TMA_A tma_load_a = make_tma_copy(
-        GmemTiledCopyA{}, tensor_a, SmemLayoutA{}(_, _, cute::Int<0>{}),
+        GmemTiledCopyA{},
+        tensor_a,
+        SmemLayoutA{}(_, _, cute::Int<0>{}),
         make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
-        size<1>(ClusterShape{})); // mcast along N mode for this M load, if any
+        size<1>(ClusterShape{}));  // mcast along N mode for this M load, if any
     typename Params::TMA_B tma_load_b = make_tma_copy(
-        GmemTiledCopyB{}, tensor_b, SmemLayoutB{}(_, _, cute::Int<0>{}),
+        GmemTiledCopyB{},
+        tensor_b,
+        SmemLayoutB{}(_, _, cute::Int<0>{}),
         make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
-        size<0>(ClusterShape{})); // mcast along M mode for this N load, if any
+        size<0>(ClusterShape{}));  // mcast along M mode for this N load, if any
     if constexpr (SwapAB) {
       auto ptr_Aux = reinterpret_cast<ElementA const *>(args.ptr_B1);
       Tensor tensor_aux =
           make_tensor(ptr_Aux, make_layout(make_shape(M, K, L), args.dA));
       typename Params::TMA_Aux tma_load_aux = make_tma_copy(
-          GmemTiledCopyA{}, tensor_aux, SmemLayoutA{}(_, _, cute::Int<0>{}),
+          GmemTiledCopyA{},
+          tensor_aux,
+          SmemLayoutA{}(_, _, cute::Int<0>{}),
           make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
           size<1>(
-              ClusterShape{})); // mcast along N mode for this M load, if any
-      return {tma_load_a,    tma_load_b,    tma_load_aux,
-              args.scale_d0, args.scale_d1, args.mma_promotion_interval};
+              ClusterShape{}));  // mcast along N mode for this M load, if any
+      return {tma_load_a,
+              tma_load_b,
+              tma_load_aux,
+              args.scale_d0,
+              args.scale_d1,
+              args.mma_promotion_interval};
     } else {
       auto ptr_Aux = reinterpret_cast<ElementB const *>(args.ptr_B1);
       Tensor tensor_aux =
           make_tensor(ptr_Aux, make_layout(make_shape(N, K, L), args.dB));
       typename Params::TMA_Aux tma_load_aux = make_tma_copy(
-          GmemTiledCopyB{}, tensor_aux, SmemLayoutB{}(_, _, cute::Int<0>{}),
+          GmemTiledCopyB{},
+          tensor_aux,
+          SmemLayoutB{}(_, _, cute::Int<0>{}),
           make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
           size<0>(
-              ClusterShape{})); // mcast along M mode for this N load, if any
-      return {tma_load_a,    tma_load_b,    tma_load_aux,
-              args.scale_d0, args.scale_d1, args.mma_promotion_interval};
+              ClusterShape{}));  // mcast along M mode for this N load, if any
+      return {tma_load_a,
+              tma_load_b,
+              tma_load_aux,
+              args.scale_d0,
+              args.scale_d1,
+              args.mma_promotion_interval};
     }
   }
 
@@ -285,8 +336,9 @@ struct CollectiveMmaGated<
     implementable = implementable && (args.mma_promotion_interval % 4 == 0);
 
     if (!implementable) {
-      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the "
-                         "minimum alignment requirements for TMA.\n");
+      CUTLASS_TRACE_HOST(
+          "  CAN IMPLEMENT: Problem Size doesn't meet the "
+          "minimum alignment requirements for TMA.\n");
     }
     return implementable;
   }
@@ -333,49 +385,64 @@ struct CollectiveMmaGated<
     // TMA requires special handling of strides to deal with coord codomain
     // mapping Represent the full tensors -- get these from TMA
     Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(
-        make_shape(M, K, L)); // (m,k,l)
+        make_shape(M, K, L));  // (m,k,l)
     Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(
-        make_shape(N, K, L)); // (n,k,l)
+        make_shape(N, K, L));  // (n,k,l)
 
     // Make tiled views, defer the slice
-    Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_, _, _),
-                               Step<_1, X, _1>{}); // (BLK_M,BLK_K,m,k,l)
-    Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_, _, _),
-                               Step<X, _1, _1>{}); // (BLK_N,BLK_K,n,k,l)
+    Tensor gA_mkl = local_tile(mA_mkl,
+                               TileShape{},
+                               make_coord(_, _, _),
+                               Step<_1, X, _1>{});  // (BLK_M,BLK_K,m,k,l)
+    Tensor gB_nkl = local_tile(mB_nkl,
+                               TileShape{},
+                               make_coord(_, _, _),
+                               Step<X, _1, _1>{});  // (BLK_N,BLK_K,n,k,l)
 
     if constexpr (SwapAB) {
       Tensor mAux_xkl = mainloop_params.tma_load_aux.get_tma_tensor(
-          make_shape(M, K, L)); // (m,k,l)
-      Tensor gAux_xkl = local_tile(mAux_xkl, TileShape{}, make_coord(_, _, _),
-                                   Step<_1, X, _1>{}); // (BLK_M,BLK_K,m,k,l)
+          make_shape(M, K, L));  // (m,k,l)
+      Tensor gAux_xkl = local_tile(mAux_xkl,
+                                   TileShape{},
+                                   make_coord(_, _, _),
+                                   Step<_1, X, _1>{});  // (BLK_M,BLK_K,m,k,l)
       return cute::make_tuple(gA_mkl, gB_nkl, gAux_xkl);
     } else {
       Tensor mAux_xkl = mainloop_params.tma_load_aux.get_tma_tensor(
-          make_shape(N, K, L)); // (n,k,l)
-      Tensor gAux_xkl = local_tile(mAux_xkl, TileShape{}, make_coord(_, _, _),
-                                   Step<X, _1, _1>{}); // (BLK_N,BLK_K,n,k,l)
+          make_shape(N, K, L));  // (n,k,l)
+      Tensor gAux_xkl = local_tile(mAux_xkl,
+                                   TileShape{},
+                                   make_coord(_, _, _),
+                                   Step<X, _1, _1>{});  // (BLK_N,BLK_K,n,k,l)
       return cute::make_tuple(gA_mkl, gB_nkl, gAux_xkl);
     }
   }
 
   /// Perform a collective-scoped matrix multiply-accumulate
   /// Producer Perspective
-  template <class TensorA, class TensorB, class TensorAux, class KTileIterator,
+  template <class TensorA,
+            class TensorB,
+            class TensorAux,
+            class KTileIterator,
             class BlockCoord>
-  CUTLASS_DEVICE void
-  load(Params const &mainloop_params, MainloopPipeline pipeline,
-       PipelineState smem_pipe_write,
-       cute::tuple<TensorA, TensorB, TensorAux> const &load_inputs,
-       BlockCoord const &blk_coord, KTileIterator k_tile_iter, int k_tile_count,
-       int thread_idx, uint32_t block_rank_in_cluster,
-       TensorStorage &shared_tensors) {
+  CUTLASS_DEVICE void load(
+      Params const &mainloop_params,
+      MainloopPipeline pipeline,
+      PipelineState smem_pipe_write,
+      cute::tuple<TensorA, TensorB, TensorAux> const &load_inputs,
+      BlockCoord const &blk_coord,
+      KTileIterator k_tile_iter,
+      int k_tile_count,
+      int thread_idx,
+      uint32_t block_rank_in_cluster,
+      TensorStorage &shared_tensors) {
     int lane_predicate = cute::elect_one_sync();
 
     if (lane_predicate) {
       Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()),
-                              SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
+                              SmemLayoutA{});  // (BLK_M,BLK_K,PIPE)
       Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
-                              SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
+                              SmemLayoutB{});  // (BLK_N,BLK_K,PIPE)
       Tensor sAux = make_tensor(make_smem_ptr(shared_tensors.smem_Aux.data()),
                                 SmemLayoutAux{});
 
@@ -403,17 +470,17 @@ struct CollectiveMmaGated<
 
       // Partition the inputs based on the current block coordinates.
       auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
-      Tensor gA = gA_mkl(_, _, m_coord, _, l_coord); // (BLK_M,BLK_K,k)
-      Tensor gB = gB_nkl(_, _, n_coord, _, l_coord); // (BLK_N,BLK_K,k)
+      Tensor gA = gA_mkl(_, _, m_coord, _, l_coord);  // (BLK_M,BLK_K,k)
+      Tensor gB = gB_nkl(_, _, n_coord, _, l_coord);  // (BLK_N,BLK_K,k)
       Tensor gAux = SwapAB ? gAux_xkl(_, _, m_coord, _, l_coord)
                            : gAux_xkl(_, _, n_coord, _, l_coord);
 
       // Applies the mapping from block_tma_a
-      Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
-      Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE)
+      Tensor tAgA = block_tma_a.partition_S(gA);  // (TMA,TMA_M,TMA_K,k)
+      Tensor tAsA = block_tma_a.partition_D(sA);  // (TMA,TMA_M,TMA_K,PIPE)
 
-      Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k)
-      Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE)
+      Tensor tBgB = block_tma_b.partition_S(gB);  // (TMA,TMA_N,TMA_K,k)
+      Tensor tBsB = block_tma_b.partition_D(sB);  // (TMA,TMA_N,TMA_K,PIPE)
 
       Tensor tAuxgAux = block_tma_aux.partition_S(gAux);
       Tensor tAuxsAux = block_tma_aux.partition_D(sAux);
@@ -426,18 +493,18 @@ struct CollectiveMmaGated<
       // Maps the tile -> block, value
       if constexpr (cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>) {
         auto block_layout =
-            Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) ->
-                                                             // block_id
+            Layout<typename DispatchPolicy::ClusterShape>{};  // (m,n) ->
+                                                              // block_id
         for (int n = 0; n < size<1>(block_layout); ++n) {
-          mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,
-                                                       n, Int<0>{}));
+          mcast_mask_a |= (uint16_t(1) << block_layout(
+                               cluster_local_block_id.x, n, Int<0>{}));
         }
       }
 
       if constexpr (cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>) {
         auto block_layout =
-            Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) ->
-                                                             // block_id
+            Layout<typename DispatchPolicy::ClusterShape>{};  // (m,n) ->
+                                                              // block_id
         for (int m = 0; m < size<0>(block_layout); ++m) {
           mcast_mask_b |= (uint16_t(1) << block_layout(
                                m, cluster_local_block_id.y, Int<0>{}));
@@ -466,11 +533,14 @@ struct CollectiveMmaGated<
 
         int write_stage = smem_pipe_write.index();
         copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a),
-             tAgA(_, _, _, *k_tile_iter), tAsA(_, _, _, write_stage));
+             tAgA(_, _, _, *k_tile_iter),
+             tAsA(_, _, _, write_stage));
         copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b),
-             tBgB(_, _, _, *k_tile_iter), tBsB(_, _, _, write_stage));
+             tBgB(_, _, _, *k_tile_iter),
+             tBsB(_, _, _, write_stage));
         copy(mainloop_params.tma_load_aux.with(*tma_barrier, mcast_mask_aux),
-             tAuxgAux(_, _, _, *k_tile_iter), tAuxsAux(_, _, _, write_stage));
+             tAuxgAux(_, _, _, *k_tile_iter),
+             tAuxsAux(_, _, _, write_stage));
         ++k_tile_iter;
 
         // Advance smem_pipe_write
@@ -499,11 +569,14 @@ struct CollectiveMmaGated<
   /// Perform a collective-scoped matrix multiply-accumulate
   /// Consumer Perspective
   template <class FrgTensorC>
-  CUTLASS_DEVICE void
-  mma(MainloopPipeline pipeline, PipelineState smem_pipe_read,
-      FrgTensorC &accum0, FrgTensorC &accum1, int k_tile_count, int thread_idx,
-      TensorStorage &shared_tensors, Params const &mainloop_params) {
-
+  CUTLASS_DEVICE void mma(MainloopPipeline pipeline,
+                          PipelineState smem_pipe_read,
+                          FrgTensorC &accum0,
+                          FrgTensorC &accum1,
+                          int k_tile_count,
+                          int thread_idx,
+                          TensorStorage &shared_tensors,
+                          Params const &mainloop_params) {
     static_assert(is_rmem<FrgTensorC>::value,
                   "C tensor must be rmem resident.");
     static_assert(cute::rank(SmemLayoutA{}) == 3,
@@ -518,9 +591,9 @@ struct CollectiveMmaGated<
                   "smem sourced instructions.");
 
     Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()),
-                            SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
+                            SmemLayoutA{});  // (BLK_M,BLK_K,PIPE)
     Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
-                            SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
+                            SmemLayoutB{});  // (BLK_N,BLK_K,PIPE)
     Tensor sAux = make_tensor(make_smem_ptr(shared_tensors.smem_Aux.data()),
                               SmemLayoutAux{});
 
@@ -531,12 +604,12 @@ struct CollectiveMmaGated<
     TiledMma tiled_mma;
     auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
 
-    Tensor tCsA = thread_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
-    Tensor tCsB = thread_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
+    Tensor tCsA = thread_mma.partition_A(sA);  // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCsB = thread_mma.partition_B(sB);  // (MMA,MMA_N,MMA_K,PIPE)
 
     // Allocate "fragments/descriptors"
-    Tensor tCrA = thread_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)
-    Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
+    Tensor tCrA = thread_mma.make_fragment_A(tCsA);  // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCrB = thread_mma.make_fragment_B(tCsB);  // (MMA,MMA_N,MMA_K,PIPE)
 
     auto tCsAux = [&]() -> auto {
       if constexpr (SwapAB) {
@@ -544,34 +617,36 @@ struct CollectiveMmaGated<
       } else {
         return thread_mma.partition_B(sAux);
       }
-    }();
+    }
+    ();
     auto tCrAux = [&]() -> auto {
       if constexpr (SwapAB) {
         return thread_mma.make_fragment_A(tCsAux);
       } else {
         return thread_mma.make_fragment_B(tCsAux);
       }
-    }();
-
-    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum0)); // M
-    CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum0)); // N
-    CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));   // K
-    CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));   // PIPE
-    if constexpr (SwapAB) {
-      CUTE_STATIC_ASSERT_V(size<1>(tCsAux) == size<1>(accum1)); // M
-      CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum1));   // N
-      CUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCsAux));   // K
-      CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tCsAux));   // PIPE
-    } else {
-      CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum1));   // M
-      CUTE_STATIC_ASSERT_V(size<1>(tCsAux) == size<2>(accum1)); // N
-      CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsAux));   // K
-      CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsAux));   // PIPE
     }
-    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
-    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
+    ();
+
+    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum0));  // M
+    CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum0));  // N
+    CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));    // K
+    CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));    // PIPE
+    if constexpr (SwapAB) {
+      CUTE_STATIC_ASSERT_V(size<1>(tCsAux) == size<1>(accum1));  // M
+      CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum1));    // N
+      CUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCsAux));    // K
+      CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tCsAux));    // PIPE
+    } else {
+      CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum1));    // M
+      CUTE_STATIC_ASSERT_V(size<1>(tCsAux) == size<2>(accum1));  // N
+      CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsAux));    // K
+      CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsAux));    // PIPE
+    }
+    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA));  // PIPE
+    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB));  // PIPE
     CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} ==
-                         size<2>(sAux)); // PIPE
+                         size<2>(sAux));  // PIPE
 
     //
     // PIPELINED MAIN LOOP
@@ -611,14 +686,20 @@ struct CollectiveMmaGated<
       CUTLASS_PRAGMA_UNROLL
       for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
         // (V,M,K) x (V,N,K) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA(_, _, k_block, read_stage),
-                   tCrB(_, _, k_block, read_stage), accumulation0());
+        cute::gemm(tiled_mma,
+                   tCrA(_, _, k_block, read_stage),
+                   tCrB(_, _, k_block, read_stage),
+                   accumulation0());
         if constexpr (SwapAB) {
-          cute::gemm(tiled_mma, tCrAux(_, _, k_block, read_stage),
-                     tCrB(_, _, k_block, read_stage), accumulation1());
+          cute::gemm(tiled_mma,
+                     tCrAux(_, _, k_block, read_stage),
+                     tCrB(_, _, k_block, read_stage),
+                     accumulation1());
         } else {
-          cute::gemm(tiled_mma, tCrA(_, _, k_block, read_stage),
-                     tCrAux(_, _, k_block, read_stage), accumulation1());
+          cute::gemm(tiled_mma,
+                     tCrA(_, _, k_block, read_stage),
+                     tCrAux(_, _, k_block, read_stage),
+                     accumulation1());
         }
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
       }
@@ -659,14 +740,20 @@ struct CollectiveMmaGated<
       CUTLASS_PRAGMA_UNROLL
       for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
         // (V,M,K) x (V,N,K) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA(_, _, k_block, read_stage),
-                   tCrB(_, _, k_block, read_stage), accumulation0());
+        cute::gemm(tiled_mma,
+                   tCrA(_, _, k_block, read_stage),
+                   tCrB(_, _, k_block, read_stage),
+                   accumulation0());
         if constexpr (SwapAB) {
-          cute::gemm(tiled_mma, tCrAux(_, _, k_block, read_stage),
-                     tCrB(_, _, k_block, read_stage), accumulation1());
+          cute::gemm(tiled_mma,
+                     tCrAux(_, _, k_block, read_stage),
+                     tCrB(_, _, k_block, read_stage),
+                     accumulation1());
         } else {
-          cute::gemm(tiled_mma, tCrA(_, _, k_block, read_stage),
-                     tCrAux(_, _, k_block, read_stage), accumulation1());
+          cute::gemm(tiled_mma,
+                     tCrA(_, _, k_block, read_stage),
+                     tCrAux(_, _, k_block, read_stage),
+                     accumulation1());
         }
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
       }
@@ -681,8 +768,9 @@ struct CollectiveMmaGated<
       accumulation0.promote_if_needed();
       accumulation1.promote_if_needed();
 
-      pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release,
-                                                    // done _computing_ on it
+      pipeline.consumer_release(
+          smem_pipe_release);  // UNLOCK smem_pipe_release,
+                               // done _computing_ on it
 
       // Advance smem_pipe_read and smem_pipe_release
       ++smem_pipe_read;
@@ -710,8 +798,9 @@ struct CollectiveMmaGated<
     warpgroup_wait<0>();
 
     for (int count = 0; count < prologue_mma_count; ++count) {
-      pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release,
-                                                    // done _computing_ on it
+      pipeline.consumer_release(
+          smem_pipe_release);  // UNLOCK smem_pipe_release,
+                               // done _computing_ on it
       ++smem_pipe_release;
     }
   }
@@ -719,6 +808,6 @@ struct CollectiveMmaGated<
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm::collective
+}  // namespace cutlass::gemm::collective
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
index be1f9747e7..837d65ae54 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
@@ -12,18 +12,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Adapt from https://github.com/vllm-project/vllm/blob/v0.7.2/csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
-// Adapted (Heavily) from: https://github.com/soundOfDestiny/cutlass/blob/9d997ce0dea4c5fa1a617db6b7ff29aa9235822c/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+// Adapt from
+// https://github.com/vllm-project/vllm/blob/v0.7.2/csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+// Adapted (Heavily) from:
+// https://github.com/soundOfDestiny/cutlass/blob/9d997ce0dea4c5fa1a617db6b7ff29aa9235822c/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
 
 /***************************************************************************************************
- * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -35,14 +37,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
@@ -73,46 +76,52 @@ using namespace cute;
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 // WarpSpecialized Mainloop
-template <
-  int Stages,
-  class ClusterShape,
-  class KernelSchedule,
-  int ScaleGranularityM_,
-  class TileShape_,
-  class ElementA_,
-  class StrideA_,
-  class ElementB_,
-  class StrideB_,
-  class TiledMma_,
-  class GmemTiledCopyA_,
-  class SmemLayoutAtomA_,
-  class SmemCopyAtomA_,
-  class TransformA_,
-  class GmemTiledCopyB_,
-  class SmemLayoutAtomB_,
-  class SmemCopyAtomB_,
-  class TransformB_>
-struct CollectiveMma<
-    MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_>,
-    TileShape_,
-    ElementA_,
-    StrideA_,
-    ElementB_,
-    StrideB_,
-    TiledMma_,
-    GmemTiledCopyA_,
-    SmemLayoutAtomA_,
-    SmemCopyAtomA_,
-    TransformA_,
-    GmemTiledCopyB_,
-    SmemLayoutAtomB_,
-    SmemCopyAtomB_,
-    TransformB_>
-{
+template <int Stages,
+          class ClusterShape,
+          class KernelSchedule,
+          int ScaleGranularityM_,
+          class TileShape_,
+          class ElementA_,
+          class StrideA_,
+          class ElementB_,
+          class StrideB_,
+          class TiledMma_,
+          class GmemTiledCopyA_,
+          class SmemLayoutAtomA_,
+          class SmemCopyAtomA_,
+          class TransformA_,
+          class GmemTiledCopyB_,
+          class SmemLayoutAtomB_,
+          class SmemCopyAtomB_,
+          class TransformB_>
+struct CollectiveMma<MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<
+                         Stages,
+                         ClusterShape,
+                         KernelSchedule,
+                         ScaleGranularityM_>,
+                     TileShape_,
+                     ElementA_,
+                     StrideA_,
+                     ElementB_,
+                     StrideB_,
+                     TiledMma_,
+                     GmemTiledCopyA_,
+                     SmemLayoutAtomA_,
+                     SmemCopyAtomA_,
+                     TransformA_,
+                     GmemTiledCopyB_,
+                     SmemLayoutAtomB_,
+                     SmemCopyAtomB_,
+                     TransformB_> {
   //
   // Type Aliases
   //
-  using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_>;
+  using DispatchPolicy =
+      MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<
+          Stages,
+          ClusterShape,
+          KernelSchedule,
+          ScaleGranularityM_>;
   using TileShape = TileShape_;
   using ElementA = ElementA_;
   using StrideA = StrideA_;
@@ -139,55 +148,91 @@ struct CollectiveMma<
   // Two threads per CTA are producers (1 for operand tile and 32 for scales)
   static constexpr int NumProducerThreadEvents = 33;
 
-  static constexpr int ScaleGranularityM = ScaleGranularityM_ == 0 ? size<0>(TileShape{}) : ScaleGranularityM_;
-  static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
+  static constexpr int ScaleGranularityM =
+      ScaleGranularityM_ == 0 ? size<0>(TileShape{}) : ScaleGranularityM_;
+  static constexpr int ScaleMsPerTile =
+      size<0>(TileShape{}) / ScaleGranularityM;
 
-  static_assert(cute::rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
-  static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
-  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert(cute::rank(SmemLayoutAtomA{}) == 2,
+                "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0,
+                "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0,
+                "SmemLayoutAtom must evenly divide tile shape.");
 
-  static_assert(cute::rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
-  static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
-  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert(cute::rank(SmemLayoutAtomB{}) == 2,
+                "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0,
+                "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0,
+                "SmemLayoutAtom must evenly divide tile shape.");
 
-  static_assert((size<0>(TileShape{}) % ScaleGranularityM) == 0, "FP8 scaling granularity must evenly divide tile shape along M.");
+  static_assert(
+      (size<0>(TileShape{}) % ScaleGranularityM) == 0,
+      "FP8 scaling granularity must evenly divide tile shape along M.");
 
   // Tile along modes in a way that maximizes the TMA box size.
   using SmemLayoutA = decltype(tile_to_shape(
       SmemLayoutAtomA{},
-      make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
-      cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
+      make_shape(shape<0>(TileShape{}),
+                 shape<2>(TileShape{}),
+                 Int<DispatchPolicy::Stages>{}),
+      cute::conditional_t<::cutlass::gemm::detail::is_major<0, StrideA>(),
+                          Step<_2, _1, _3>,
+                          Step<_1, _2, _3>>{}));
   using SmemLayoutB = decltype(tile_to_shape(
       SmemLayoutAtomB{},
-      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
-      cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
+      make_shape(shape<1>(TileShape{}),
+                 shape<2>(TileShape{}),
+                 Int<DispatchPolicy::Stages>{}),
+      cute::conditional_t<::cutlass::gemm::detail::is_major<0, StrideB>(),
+                          Step<_2, _1, _3>,
+                          Step<_1, _2, _3>>{}));
 
   // Block scaling gmem-to-smem copy atom
-  using SmemBlockScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
-  using SmemBlockScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
+  using SmemBlockScalingCopyAtomA =
+      Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>,
+                ElementBlockScale>;
+  using SmemBlockScalingCopyAtomB =
+      Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>,
+                ElementBlockScale>;
 
   // Block scaling smem layout
-  using SmemLayoutScaleA = Layout<Shape<Int<ScaleMsPerTile>, Int<DispatchPolicy::Stages>>>;
-  using SmemLayoutScaleB = Layout<Shape<Int<DispatchPolicy::Stages>>, Stride<_1>>; // `ScaleNsPerTile` is always 1.
+  using SmemLayoutScaleA =
+      Layout<Shape<Int<ScaleMsPerTile>, Int<DispatchPolicy::Stages>>>;
+  using SmemLayoutScaleB = Layout<Shape<Int<DispatchPolicy::Stages>>,
+                                  Stride<_1>>;  // `ScaleNsPerTile` is always 1.
 
-  static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
-  static_assert(cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
-                cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
-                "MMA atom must source both A and B operand from smem_desc for this mainloop.");
-  static_assert(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>,
-      "GmemTiledCopy - invalid SM90 TMA copy atom specified.");
-  static_assert(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>,
-      "GmemTiledCopy - invalid SM90 TMA copy atom specified.");
-  static_assert(cute::is_same_v<ElementAccumulator, ElementBlockScale>,
-             "ElementAccumulator and ElementBlockScale should be same datatype");
+  static_assert(DispatchPolicy::Stages >= 2,
+                "Specialization requires Stages set to value 1 or more.");
+  static_assert(cute::is_base_of<cute::GMMA::DescriptorIterator,
+                                 typename TiledMma::FrgTypeA>::value &&
+                    cute::is_base_of<cute::GMMA::DescriptorIterator,
+                                     typename TiledMma::FrgTypeB>::value,
+                "MMA atom must source both A and B operand from smem_desc for "
+                "this mainloop.");
+  static_assert(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> ||
+                    cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>,
+                "GmemTiledCopy - invalid SM90 TMA copy atom specified.");
+  static_assert(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> ||
+                    cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>,
+                "GmemTiledCopy - invalid SM90 TMA copy atom specified.");
+  static_assert(
+      cute::is_same_v<ElementAccumulator, ElementBlockScale>,
+      "ElementAccumulator and ElementBlockScale should be same datatype");
 
-  struct SharedStorage
-  {
+  struct SharedStorage {
     struct TensorStorage : cute::aligned_struct<128> {
-      cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>> smem_A;  // mxk
-      cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>> smem_B;  // nxk
-      cute::array_aligned<ElementBlockScale, cute::cosize_v<SmemLayoutScaleA>> smem_scale_A; // ScaleMsPerTile x k
-      cute::array_aligned<ElementBlockScale, cute::cosize_v<SmemLayoutScaleB>> smem_scale_B; // 1xk
+      cute::array_aligned<typename TiledMma::ValTypeA,
+                          cute::cosize_v<SmemLayoutA>>
+          smem_A;  // mxk
+      cute::array_aligned<typename TiledMma::ValTypeB,
+                          cute::cosize_v<SmemLayoutB>>
+          smem_B;  // nxk
+      cute::array_aligned<ElementBlockScale, cute::cosize_v<SmemLayoutScaleA>>
+          smem_scale_A;  // ScaleMsPerTile x k
+      cute::array_aligned<ElementBlockScale, cute::cosize_v<SmemLayoutScaleB>>
+          smem_scale_B;  // 1xk
     } tensors;
 
     using PipelineStorage = typename MainloopPipeline::SharedStorage;
@@ -211,15 +256,19 @@ struct CollectiveMma<
     // Assumption: StrideA is congruent with Problem_MK
     using TMA_A = decltype(make_tma_copy_A_sm90(
         GmemTiledCopyA{},
-        make_tensor(static_cast<ElementA const*>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
-        SmemLayoutA{}(_,_,0),
+        make_tensor(static_cast<ElementA const*>(nullptr),
+                    repeat_like(StrideA{}, int32_t(0)),
+                    StrideA{}),
+        SmemLayoutA{}(_, _, 0),
         TileShape{},
         ClusterShape{}));
     // Assumption: StrideB is congruent with Problem_NK
     using TMA_B = decltype(make_tma_copy_B_sm90(
         GmemTiledCopyB{},
-        make_tensor(static_cast<ElementB const*>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
-        SmemLayoutB{}(_,_,0),
+        make_tensor(static_cast<ElementB const*>(nullptr),
+                    repeat_like(StrideB{}, int32_t(0)),
+                    StrideB{}),
+        SmemLayoutB{}(_, _, 0),
         TileShape{},
         ClusterShape{}));
     TMA_A tma_load_a;
@@ -237,103 +286,128 @@ struct CollectiveMma<
   //
 
   template <class ProblemShape>
-  static constexpr Params
-  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
-    (void) workspace;
+  static constexpr Params to_underlying_arguments(
+      ProblemShape const& problem_shape,
+      Arguments const& args,
+      void* workspace) {
+    (void)workspace;
 
-    // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
+    // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is
+    // only rank-3 (MNK)
     auto problem_shape_MNKL = append<4>(problem_shape, 1);
-    auto [M,N,K,L] = problem_shape_MNKL;
+    auto [M, N, K, L] = problem_shape_MNKL;
 
     auto ptr_A = reinterpret_cast<ElementA const*>(args.ptr_A);
     auto ptr_B = reinterpret_cast<ElementB const*>(args.ptr_B);
 
-    Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA));
-    Tensor tensor_b = make_tensor(ptr_B, make_layout(make_shape(N,K,L), args.dB));
-    typename Params::TMA_A tma_load_a = make_tma_copy_A_sm90(
-        GmemTiledCopyA{},
-        tensor_a,
-        SmemLayoutA{}(_,_,cute::Int<0>{}),
-        TileShape{},
-        ClusterShape{});
-    typename Params::TMA_B tma_load_b = make_tma_copy_B_sm90(
-        GmemTiledCopyB{},
-        tensor_b,
-        SmemLayoutB{}(_,_,cute::Int<0>{}),
-        TileShape{},
-        ClusterShape{});
+    Tensor tensor_a =
+        make_tensor(ptr_A, make_layout(make_shape(M, K, L), args.dA));
+    Tensor tensor_b =
+        make_tensor(ptr_B, make_layout(make_shape(N, K, L), args.dB));
+    typename Params::TMA_A tma_load_a =
+        make_tma_copy_A_sm90(GmemTiledCopyA{},
+                             tensor_a,
+                             SmemLayoutA{}(_, _, cute::Int<0>{}),
+                             TileShape{},
+                             ClusterShape{});
+    typename Params::TMA_B tma_load_b =
+        make_tma_copy_B_sm90(GmemTiledCopyB{},
+                             tensor_b,
+                             SmemLayoutB{}(_, _, cute::Int<0>{}),
+                             TileShape{},
+                             ClusterShape{});
     uint32_t transaction_bytes_mk = TmaTransactionBytesMK;
     uint32_t transaction_bytes_nk = TmaTransactionBytesNK;
     uint32_t transaction_bytes = transaction_bytes_mk + transaction_bytes_nk;
 
-    return {
-      tma_load_a,
-      tma_load_b,
-      transaction_bytes,
-      transaction_bytes_mk,
-      transaction_bytes_nk,
-      args.ptr_scale_A,
-      args.ptr_scale_B
-    };
+    return {tma_load_a,
+            tma_load_b,
+            transaction_bytes,
+            transaction_bytes_mk,
+            transaction_bytes_nk,
+            args.ptr_scale_A,
+            args.ptr_scale_B};
   }
 
-  template<class ProblemShape>
-  static bool
-  can_implement(
-      ProblemShape const& problem_shape,
-      [[maybe_unused]] Arguments const& args) {
+  template <class ProblemShape>
+  static bool can_implement(ProblemShape const& problem_shape,
+                            [[maybe_unused]] Arguments const& args) {
     constexpr int tma_alignment_bits = 128;
     auto problem_shape_MNKL = append<4>(problem_shape, 1);
-    auto [M,N,K,L] = problem_shape_MNKL;
+    auto [M, N, K, L] = problem_shape_MNKL;
 
     bool implementable = true;
-    constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
-    implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
-    constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
-    implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
+    constexpr int min_tma_aligned_elements_A =
+        tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
+    implementable =
+        implementable &&
+        cutlass::detail::check_alignment<min_tma_aligned_elements_A>(
+            cute::make_shape(M, K, L), StrideA{});
+    constexpr int min_tma_aligned_elements_B =
+        tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
+    implementable =
+        implementable &&
+        cutlass::detail::check_alignment<min_tma_aligned_elements_B>(
+            cute::make_shape(N, K, L), StrideB{});
 
     if (!implementable) {
-      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
+      CUTLASS_TRACE_HOST(
+          "  CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment "
+          "requirements for TMA.\n");
     }
     return implementable;
   }
 
   static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
   static constexpr int K_PIPE_MMAS = 1;
-  static constexpr uint32_t TmaTransactionBytesMK =
-        cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof_bits<ElementA>::value));
-  static constexpr uint32_t TmaTransactionBytesNK =
-        cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof_bits<ElementB>::value));
-  static constexpr uint32_t TmaTransactionBytes = TmaTransactionBytesMK + TmaTransactionBytesNK;
+  static constexpr uint32_t TmaTransactionBytesMK = cutlass::bits_to_bytes(
+      size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) *
+      static_cast<uint32_t>(sizeof_bits<ElementA>::value));
+  static constexpr uint32_t TmaTransactionBytesNK = cutlass::bits_to_bytes(
+      size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) *
+      static_cast<uint32_t>(sizeof_bits<ElementB>::value));
+  static constexpr uint32_t TmaTransactionBytes =
+      TmaTransactionBytesMK + TmaTransactionBytesNK;
 
-  /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
+  /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best
+  /// performance
   CUTLASS_DEVICE
-  static void prefetch_tma_descriptors(Params const& mainloop_params)
-  {
-    cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor());
-    cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor());
+  static void prefetch_tma_descriptors(Params const& mainloop_params) {
+    cute::prefetch_tma_descriptor(
+        mainloop_params.tma_load_a.get_tma_descriptor());
+    cute::prefetch_tma_descriptor(
+        mainloop_params.tma_load_b.get_tma_descriptor());
   }
 
   /// Set up the data needed by this collective for load and mma.
-  /// Returns a tuple of tensors. The collective and the kernel layer have the contract
-  /// Returned tuple must contain at least two elements, with the first two elements being:
-  /// gA_mkl - The tma tensor, A after a local tile so it has shape  (BLK_M,BLK_K,m,k,l)
-  /// gB_nkl - The tma tensor, B after a local tile so it has shape  (BLK_N,BLK_K,n,k,l)
+  /// Returns a tuple of tensors. The collective and the kernel layer have the
+  /// contract Returned tuple must contain at least two elements, with the first
+  /// two elements being: gA_mkl - The tma tensor, A after a local tile so it
+  /// has shape  (BLK_M,BLK_K,m,k,l) gB_nkl - The tma tensor, B after a local
+  /// tile so it has shape  (BLK_N,BLK_K,n,k,l)
   template <class ProblemShape_MNKL>
-  CUTLASS_DEVICE auto
-  load_init(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params) const {
+  CUTLASS_DEVICE auto load_init(ProblemShape_MNKL const& problem_shape_MNKL,
+                                Params const& mainloop_params) const {
     using X = Underscore;
     // Separate out problem shape for convenience
-    auto [M,N,K,L] = problem_shape_MNKL;
+    auto [M, N, K, L] = problem_shape_MNKL;
 
-    // TMA requires special handling of strides to deal with coord codomain mapping
-    // Represent the full tensors -- get these from TMA
-    Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(make_shape(M,K,L));                            // (m,k,l)
-    Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(make_shape(N,K,L));                            // (n,k,l)
+    // TMA requires special handling of strides to deal with coord codomain
+    // mapping Represent the full tensors -- get these from TMA
+    Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(
+        make_shape(M, K, L));  // (m,k,l)
+    Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(
+        make_shape(N, K, L));  // (n,k,l)
 
     // Make tiled views, defer the slice
-    Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{});        // (BLK_M,BLK_K,m,k,l)
-    Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{});        // (BLK_N,BLK_K,n,k,l)
+    Tensor gA_mkl = local_tile(mA_mkl,
+                               TileShape{},
+                               make_coord(_, _, _),
+                               Step<_1, X, _1>{});  // (BLK_M,BLK_K,m,k,l)
+    Tensor gB_nkl = local_tile(mB_nkl,
+                               TileShape{},
+                               make_coord(_, _, _),
+                               Step<X, _1, _1>{});  // (BLK_N,BLK_K,n,k,l)
 
     constexpr auto scales_m = Int<ScaleMsPerTile>{};
     auto tM = get<2>(gA_mkl.shape());
@@ -341,84 +415,103 @@ struct CollectiveMma<
     auto tK = get<3>(gA_mkl.shape());
 
     // Make the tiled views of scale tensors
-    auto scaleA_shape = make_shape(M / ScaleGranularityM, tK, L); // (scale_m,k,l)
-    auto scaleA_layout = make_ordered_layout(scaleA_shape,  Step<_0, _1, _2>{});
-    auto scaleB_shape = make_shape(tN, tK, L); // (n,k,l)
+    auto scaleA_shape =
+        make_shape(M / ScaleGranularityM, tK, L);  // (scale_m,k,l)
+    auto scaleA_layout = make_ordered_layout(scaleA_shape, Step<_0, _1, _2>{});
+    auto scaleB_shape = make_shape(tN, tK, L);  // (n,k,l)
     auto scaleB_layout = make_ordered_layout(scaleB_shape, Step<_1, _0, _2>{});
 
-    // Note that mScaleA_mkl and mScaleB_nkl are already blocked tiled in the `m` host and
-    // gScaleA_mkl and gScaleB_nkl in `g` global memory are same as mScaleA_mkl and mScaleB_nkl.
-    Tensor mScaleA_mkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_A), scaleA_layout); // (scale_m,k,l)
-    Tensor mScaleB_nkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_B), scaleB_layout); // (n,k,l)
+    // Note that mScaleA_mkl and mScaleB_nkl are already blocked tiled in the
+    // `m` host and gScaleA_mkl and gScaleB_nkl in `g` global memory are same as
+    // mScaleA_mkl and mScaleB_nkl.
+    Tensor mScaleA_mkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_A),
+                                     scaleA_layout);  // (scale_m,k,l)
+    Tensor mScaleB_nkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_B),
+                                     scaleB_layout);  // (n,k,l)
 
     return cute::make_tuple(gA_mkl, gB_nkl, mScaleA_mkl, mScaleB_nkl);
   }
 
   /// Perform a collective-scoped matrix multiply-accumulate
   /// Producer Perspective
-  template <
-    class TensorA, class TensorB,
-    class TensorScaleA, class TensorScaleB,
-    class KTileIterator, class BlockCoord
-  >
-  CUTLASS_DEVICE void
-  load(
+  template <class TensorA,
+            class TensorB,
+            class TensorScaleA,
+            class TensorScaleB,
+            class KTileIterator,
+            class BlockCoord>
+  CUTLASS_DEVICE void load(
       Params const& mainloop_params,
       MainloopPipeline pipeline,
       PipelineState smem_pipe_write,
-      cute::tuple<TensorA, TensorB, TensorScaleA, TensorScaleB> const& load_inputs,
+      cute::tuple<TensorA, TensorB, TensorScaleA, TensorScaleB> const&
+          load_inputs,
       BlockCoord const& blk_coord,
-      KTileIterator k_tile_iter, int k_tile_count,
+      KTileIterator k_tile_iter,
+      int k_tile_count,
       int thread_idx,
       uint32_t block_rank_in_cluster,
       TensorStorage& shared_tensors) {
     int lane_predicate = cute::elect_one_sync();
 
     // Blockscaling: Tma loads for load_input and CpAsync for load_scale
-    Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{});        // (BLK_M,BLK_K,PIPE)
-    Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{});        // (BLK_N,BLK_K,PIPE)
-    Tensor sScaleA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_A.data()), SmemLayoutScaleA{}); // (ScaleMsPerTile,k)
-    Tensor sScaleB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()), SmemLayoutScaleB{}); // (k)
+    Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()),
+                            SmemLayoutA{});  // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
+                            SmemLayoutB{});  // (BLK_N,BLK_K,PIPE)
+    Tensor sScaleA =
+        make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_A.data()),
+                    SmemLayoutScaleA{});  // (ScaleMsPerTile,k)
+    Tensor sScaleB =
+        make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()),
+                    SmemLayoutScaleB{});  // (k)
 
     //
     // Prepare the TMA loads for A and B
     //
 
     constexpr uint32_t cluster_shape_x = get<0>(ClusterShape());
-    uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
+    uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x,
+                                    block_rank_in_cluster / cluster_shape_x};
 
     Tensor gA_mkl = get<0>(load_inputs);
     Tensor gB_nkl = get<1>(load_inputs);
 
-    auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
-    auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
+    auto block_tma_a =
+        mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
+    auto block_tma_b =
+        mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
 
     // Partition the inputs based on the current block coordinates.
     auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
-    Tensor gA = gA_mkl(_,_,m_coord,_,l_coord);                                                     // (BLK_M,BLK_K,k)
-    Tensor gB = gB_nkl(_,_,n_coord,_,l_coord);                                                     // (BLK_N,BLK_K,k)
+    Tensor gA = gA_mkl(_, _, m_coord, _, l_coord);  // (BLK_M,BLK_K,k)
+    Tensor gB = gB_nkl(_, _, n_coord, _, l_coord);  // (BLK_N,BLK_K,k)
 
-
-    // Block scaling: load_scale has scaling tensors in global memory which are not tiled
+    // Block scaling: load_scale has scaling tensors in global memory which are
+    // not tiled
     Tensor mScaleA_mkl = get<2>(load_inputs);
     Tensor mScaleB_nkl = get<3>(load_inputs);
     auto scales_m = get<0>(mScaleA_mkl.shape());
 
     Tensor cScaleA_mkl = make_identity_tensor(mScaleA_mkl.shape());
 
-    Tensor gScaleA = local_tile(
-      mScaleA_mkl, make_tile(Int<ScaleMsPerTile>{}),
-      make_coord(m_coord,_,l_coord));                   // (ScaleMsPerTile,k,1)
-    Tensor cScaleA = local_tile(
-      cScaleA_mkl, make_tile(Int<ScaleMsPerTile>{}),
-      make_coord(m_coord,_,l_coord));
-    Tensor gScaleB = mScaleB_nkl(n_coord,_,l_coord);                                           // (1,k,1)
+    Tensor gScaleA =
+        local_tile(mScaleA_mkl,
+                   make_tile(Int<ScaleMsPerTile>{}),
+                   make_coord(m_coord, _, l_coord));  // (ScaleMsPerTile,k,1)
+    Tensor cScaleA = local_tile(cScaleA_mkl,
+                                make_tile(Int<ScaleMsPerTile>{}),
+                                make_coord(m_coord, _, l_coord));
+    Tensor gScaleB = mScaleB_nkl(n_coord, _, l_coord);  // (1,k,1)
 
     // TODO: test `scale_copy_a` with `ScaleMsPerTile` < 128
-    TiledCopy scale_copy_a = make_tiled_copy(SmemBlockScalingCopyAtomA{},
-      Layout<Shape<_32, _1>>{}, Layout<Shape<_4, _1>>{}); // (1,1,1)
+    TiledCopy scale_copy_a =
+        make_tiled_copy(SmemBlockScalingCopyAtomA{},
+                        Layout<Shape<_32, _1>>{},
+                        Layout<Shape<_4, _1>>{});  // (1,1,1)
     TiledCopy scale_copy_b = make_tiled_copy(SmemBlockScalingCopyAtomB{},
-      Layout<Shape<_1>>{}, Layout<Shape<_1>>{}); // (1,1,1)
+                                             Layout<Shape<_1>>{},
+                                             Layout<Shape<_1>>{});  // (1,1,1)
     ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(threadIdx.x);
     ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(threadIdx.x);
 
@@ -430,11 +523,11 @@ struct CollectiveMma<
     Tensor tBsB_ScaleB = thr_scale_copy_b.partition_D(sScaleB);
 
     // Applies the mapping from block_tma_a
-    Tensor tAgA = block_tma_a.partition_S(gA);                                              // (TMA,TMA_M,TMA_K,k)
-    Tensor tAsA = block_tma_a.partition_D(sA);                                              // (TMA,TMA_M,TMA_K,PIPE)
+    Tensor tAgA = block_tma_a.partition_S(gA);  // (TMA,TMA_M,TMA_K,k)
+    Tensor tAsA = block_tma_a.partition_D(sA);  // (TMA,TMA_M,TMA_K,PIPE)
 
-    Tensor tBgB = block_tma_b.partition_S(gB);                                              // (TMA,TMA_N,TMA_K,k)
-    Tensor tBsB = block_tma_b.partition_D(sB);                                              // (TMA,TMA_N,TMA_K,PIPE)
+    Tensor tBgB = block_tma_b.partition_S(gB);  // (TMA,TMA_N,TMA_K,k)
+    Tensor tBsB = block_tma_b.partition_D(sB);  // (TMA,TMA_N,TMA_K,PIPE)
 
     uint16_t mcast_mask_a = 0;
     uint16_t mcast_mask_b = 0;
@@ -442,30 +535,34 @@ struct CollectiveMma<
     // Issue TmaLoads for GEMM operands A/B and CpAsync for scale tensors
     // Maps the tile -> block, value
     if constexpr (cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>) {
-      auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{};                       // (m,n) -> block_id
+      auto block_layout =
+          Layout<typename DispatchPolicy::ClusterShape>{};  // (m,n) -> block_id
       for (int n = 0; n < size<1>(block_layout); ++n) {
-        mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,n,Int<0>{}));
+        mcast_mask_a |= (uint16_t(1) << block_layout(
+                             cluster_local_block_id.x, n, Int<0>{}));
       }
     }
 
     if constexpr (cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>) {
-      auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{};                       // (m,n) -> block_id
+      auto block_layout =
+          Layout<typename DispatchPolicy::ClusterShape>{};  // (m,n) -> block_id
       for (int m = 0; m < size<0>(block_layout); ++m) {
-        mcast_mask_b |= (uint16_t(1) << block_layout(m,cluster_local_block_id.y,Int<0>{}));
+        mcast_mask_b |= (uint16_t(1) << block_layout(
+                             m, cluster_local_block_id.y, Int<0>{}));
       }
     }
 
     // Allocate predicate tensors for a_scales (since we can't guarantee that
     // all scales are valid, since we could have a partial tiles along M)
-    Tensor tApA_ScaleA = make_tensor<bool>(shape(tAsA_ScaleA(_,_,0)));
-    #pragma unroll
+    Tensor tApA_ScaleA = make_tensor<bool>(shape(tAsA_ScaleA(_, _, 0)));
+#pragma unroll
     for (int i = 0; i < size(tApA_ScaleA); ++i) {
       tApA_ScaleA(i) = get<0>(tAcA_ScaleA(i)) < scales_m;
     }
 
     // Mainloop
     CUTLASS_PRAGMA_NO_UNROLL
-    for ( ; k_tile_count > 0; --k_tile_count) {
+    for (; k_tile_count > 0; --k_tile_count) {
       // LOCK smem_pipe_write for _writing_
       pipeline.producer_acquire(smem_pipe_write);
 
@@ -477,13 +574,25 @@ struct CollectiveMma<
       BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
 
       // Copy operands A and B from global memory to shared memory
-      if (lane_predicate) copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
-      if (lane_predicate) copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
+      if (lane_predicate)
+        copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a),
+             tAgA(_, _, _, *k_tile_iter),
+             tAsA(_, _, _, write_stage));
+      if (lane_predicate)
+        copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b),
+             tBgB(_, _, _, *k_tile_iter),
+             tBsB(_, _, _, write_stage));
 
       // Copy scale tensors from global memory to shared memory
-      copy_if(scale_copy_a, tApA_ScaleA, tAgA_ScaleA(_,_,*k_tile_iter), tAsA_ScaleA(_,_,write_stage));
-      copy(scale_copy_b, tBgB_ScaleB(_,*k_tile_iter), tBsB_ScaleB(_,write_stage));
-      pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive_noinc);
+      copy_if(scale_copy_a,
+              tApA_ScaleA,
+              tAgA_ScaleA(_, _, *k_tile_iter),
+              tAsA_ScaleA(_, _, write_stage));
+      copy(scale_copy_b,
+           tBgB_ScaleB(_, *k_tile_iter),
+           tBsB_ScaleB(_, write_stage));
+      pipeline.producer_commit(smem_pipe_write,
+                               cutlass::arch::cpasync_barrier_arrive_noinc);
 
       ++k_tile_iter;
 
@@ -493,10 +602,8 @@ struct CollectiveMma<
   }
 
   /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
-  CUTLASS_DEVICE void
-  load_tail(
-      MainloopPipeline pipeline,
-      PipelineState smem_pipe_write) {
+  CUTLASS_DEVICE void load_tail(MainloopPipeline pipeline,
+                                PipelineState smem_pipe_write) {
     int lane_predicate = cute::elect_one_sync();
 
     // Issue the epilogue waits
@@ -513,37 +620,46 @@ struct CollectiveMma<
 
   /// Perform a collective-scoped matrix multiply-accumulate
   /// Consumer Perspective
-  template <
-    class FrgTensorC
-  >
-  CUTLASS_DEVICE void
-  mma(MainloopPipeline pipeline,
-      PipelineState smem_pipe_read,
-      FrgTensorC& accum,
-      int k_tile_count,
-      int thread_idx,
-      TensorStorage& shared_tensors,
-      Params const& mainloop_params) {
-
-
-    static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
-    static_assert(cute::rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
-    static_assert(cute::rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
+  template <class FrgTensorC>
+  CUTLASS_DEVICE void mma(MainloopPipeline pipeline,
+                          PipelineState smem_pipe_read,
+                          FrgTensorC& accum,
+                          int k_tile_count,
+                          int thread_idx,
+                          TensorStorage& shared_tensors,
+                          Params const& mainloop_params) {
+    static_assert(is_rmem<FrgTensorC>::value,
+                  "C tensor must be rmem resident.");
+    static_assert(cute::rank(SmemLayoutA{}) == 3,
+                  "Smem layout must be rank 3.");
+    static_assert(cute::rank(SmemLayoutB{}) == 3,
+                  "Smem layout must be rank 3.");
     static_assert(cute::is_void_v<SmemCopyAtomA>,
-      "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
+                  "SM90 GMMA mainloops cannot have a non-void copy atom for "
+                  "smem sourced instructions.");
     static_assert(cute::is_void_v<SmemCopyAtomB>,
-      "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
+                  "SM90 GMMA mainloops cannot have a non-void copy atom for "
+                  "smem sourced instructions.");
 
-    Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{});          // (BLK_M,BLK_K,PIPE)
-    Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{});          // (BLK_N,BLK_K,PIPE)
+    Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()),
+                            SmemLayoutA{});  // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()),
+                            SmemLayoutB{});  // (BLK_N,BLK_K,PIPE)
 
     // Block scaling
-    Tensor sScaleAViewAsC = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_A.data()),
-      Layout<
-        Shape<Shape<Int<ScaleGranularityM>, Int<ScaleMsPerTile>>, cute::tuple_element_t<1, TileShape>, Int<DispatchPolicy::Stages>>,
-        Stride<Stride<_0, _1>, _0, Int<ScaleMsPerTile>>
-      >{}); // ((ScaleGranularityM,ScaleMsPerTile),n,k)
-    Tensor sScaleB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()), SmemLayoutScaleB{}); // (k)
+    Tensor sScaleAViewAsC = make_tensor(
+        cute::make_smem_ptr(shared_tensors.smem_scale_A.data()),
+        Layout<
+            Shape<Shape<Int<ScaleGranularityM>, Int<ScaleMsPerTile>>,
+                  cute::tuple_element_t<1, TileShape>,
+                  Int<DispatchPolicy::Stages>>,
+            Stride<
+                Stride<_0, _1>,
+                _0,
+                Int<ScaleMsPerTile>>>{});  // ((ScaleGranularityM,ScaleMsPerTile),n,k)
+    Tensor sScaleB =
+        make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()),
+                    SmemLayoutScaleB{});  // (k)
 
     //
     // Define C accumulators and A/B partitioning
@@ -551,52 +667,68 @@ struct CollectiveMma<
 
     // Layout of warp group to thread mapping
 
-    static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and
-                  stride<0>(typename TiledMma::BLayout{}) == 0 and
-                  size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
-                  size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
-                  "Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
+    static_assert(
+        stride<0>(typename TiledMma::ALayout{}) == 0 and
+            stride<0>(typename TiledMma::BLayout{}) == 0 and
+            size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
+            size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup,
+        "Stride of the first mode must be 0 and the size of the mode must be "
+        "NumThreadsPerWarpGroup");
 
     constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
-    Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{},
-                                                  Int<NumThreadsPerWarpGroup>{});
+    Layout warp_group_thread_layout =
+        make_layout(Int<MmaWarpGroups>{}, Int<NumThreadsPerWarpGroup>{});
 
-    int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
+    int warp_group_idx =
+        __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
 
     TiledMma tiled_mma;
-    auto thread_mma = tiled_mma.get_slice(warp_group_thread_layout(warp_group_idx));
+    auto thread_mma =
+        tiled_mma.get_slice(warp_group_thread_layout(warp_group_idx));
 
-    Tensor tCsScaleAViewAsC = tiled_mma.get_slice(thread_idx).partition_C(sScaleAViewAsC);    // (MMA,MMA_M,MMA_N,PIPE), `thread_mma` above is correct when partitioning A and B, but it is not correct when partitioning C.
+    Tensor tCsScaleAViewAsC =
+        tiled_mma.get_slice(thread_idx)
+            .partition_C(
+                sScaleAViewAsC);  // (MMA,MMA_M,MMA_N,PIPE), `thread_mma` above
+                                  // is correct when partitioning A and B, but
+                                  // it is not correct when partitioning C.
 
-    Tensor tCsA = thread_mma.partition_A(sA);                                                 // (MMA,MMA_M,MMA_K,PIPE)
-    Tensor tCsB = thread_mma.partition_B(sB);                                                 // (MMA,MMA_N,MMA_K,PIPE)
+    Tensor tCsA = thread_mma.partition_A(sA);  // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCsB = thread_mma.partition_B(sB);  // (MMA,MMA_N,MMA_K,PIPE)
 
     // Allocate "fragments/descriptors"
-    Tensor tCrA = thread_mma.make_fragment_A(tCsA);                                           // (MMA,MMA_M,MMA_K,PIPE)
-    Tensor tCrB = thread_mma.make_fragment_B(tCsB);                                           // (MMA,MMA_N,MMA_K,PIPE)
+    Tensor tCrA = thread_mma.make_fragment_A(tCsA);  // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCrB = thread_mma.make_fragment_B(tCsB);  // (MMA,MMA_N,MMA_K,PIPE)
 
-    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum));                                                         // M
-    CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum));                                                         // N
-    CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));                                                          // K
-    CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));                                                       // PIPE
-    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA));                                         // PIPE
-    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB));                                         // PIPE
+    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum));               // M
+    CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum));               // N
+    CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));                // K
+    CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));                // PIPE
+    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA));  // PIPE
+    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB));  // PIPE
 
     //
     // PIPELINED MAIN LOOP
     //
-    static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS <  K_PIPE_MAX),
-        "ERROR : Incorrect number of MMAs in flight");
+    static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS < K_PIPE_MAX),
+                  "ERROR : Incorrect number of MMAs in flight");
 
     // We release buffers to producer warps(dma load) with some mmas in flight
     PipelineState smem_pipe_release = smem_pipe_read;
 
     // Per block scale values for operand A and B
 
-    using RegLayoutScaleAViewAsC = decltype(make_layout_like(tCsScaleAViewAsC(_, _, _, 0).layout())); // `make_layout_like` makes a compact layout.
-    using RegLayoutScaleAEssential = decltype(filter_zeros(RegLayoutScaleAViewAsC{}.stride(), RegLayoutScaleAViewAsC{}.shape())); // an interface to traverse the underlying storage for the compact layout mentioned above
+    using RegLayoutScaleAViewAsC = decltype(make_layout_like(
+        tCsScaleAViewAsC(_, _, _, 0)
+            .layout()));  // `make_layout_like` makes a compact layout.
+    using RegLayoutScaleAEssential = decltype(filter_zeros(
+        RegLayoutScaleAViewAsC{}.stride(),
+        RegLayoutScaleAViewAsC{}
+            .shape()));  // an interface to traverse the underlying storage for
+                         // the compact layout mentioned above
 
-    Tensor tCrScaleAViewAsC = make_tensor<ElementBlockScale>(RegLayoutScaleAViewAsC{});              // (MMA,MMA_M,MMA_N)
+    Tensor tCrScaleAViewAsC = make_tensor<ElementBlockScale>(
+        RegLayoutScaleAViewAsC{});  // (MMA,MMA_M,MMA_N)
     ElementBlockScale scale_b;
 
     // Prologue GMMAs
@@ -604,12 +736,16 @@ struct CollectiveMma<
 
     tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
 
-    GmmaFP8AccumulationWithScale accumulation(accum, size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{}), size<2>(tCrA));
+    GmmaFP8AccumulationWithScale accumulation(
+        accum,
+        size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{}),
+        size<2>(tCrA));
     warpgroup_fence_operand(accumulation());
     CUTLASS_PRAGMA_UNROLL
-    for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue)
-    {
-      // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
+    for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0;
+         --k_tile_prologue) {
+      // WAIT on smem_pipe_read until its data are available (phase bit flips
+      // from rdPhaseBit value)
       auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
       pipeline.consumer_wait(smem_pipe_read, barrier_token);
 
@@ -623,11 +759,16 @@ struct CollectiveMma<
       scale_b = sScaleB[read_stage];
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
-        tCrScaleAViewAsC.data()[i] = tCsScaleAViewAsC(_, _, _, read_stage)(idx2crd(i, RegLayoutScaleAEssential{}));
+        tCrScaleAViewAsC.data()[i] = tCsScaleAViewAsC(
+            _, _, _, read_stage)(idx2crd(i, RegLayoutScaleAEssential{}));
       }
       if constexpr (ScaleMsPerTile == 1) {
         static_assert(size(RegLayoutScaleAEssential{}) == 1);
-        tCrScaleAViewAsC.data()[0] = __shfl_sync(0xffffffff, tCrScaleAViewAsC.data()[0] * scale_b, 0); // `tCrScaleAViewAsC.data()[0]` are all same in a warp group when `ScaleMsPerTile == 1`.
+        tCrScaleAViewAsC.data()[0] =
+            __shfl_sync(0xffffffff,
+                        tCrScaleAViewAsC.data()[0] * scale_b,
+                        0);  // `tCrScaleAViewAsC.data()[0]` are all same in a
+                             // warp group when `ScaleMsPerTile == 1`.
       } else {
         CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
@@ -640,7 +781,10 @@ struct CollectiveMma<
       CUTLASS_PRAGMA_UNROLL
       for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
         // (V,M,K) x (V,N,K) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
+        cute::gemm(tiled_mma,
+                   tCrA(_, _, k_block, read_stage),
+                   tCrB(_, _, k_block, read_stage),
+                   accumulation());
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
       }
       warpgroup_commit_batch();
@@ -656,9 +800,9 @@ struct CollectiveMma<
     k_tile_count -= prologue_mma_count;
 
     CUTLASS_PRAGMA_NO_UNROLL
-    for ( ; k_tile_count > 0; --k_tile_count)
-    {
-      // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
+    for (; k_tile_count > 0; --k_tile_count) {
+      // WAIT on smem_pipe_read until its data are available (phase bit flips
+      // from rdPhaseBit value)
       auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
       pipeline.consumer_wait(smem_pipe_read, barrier_token);
 
@@ -668,15 +812,21 @@ struct CollectiveMma<
 
       int read_stage = smem_pipe_read.index();
 
-      // Load per block scale values from shared memory to registers (at most twice per block along M and exactly once per block along N)
+      // Load per block scale values from shared memory to registers (at most
+      // twice per block along M and exactly once per block along N)
       scale_b = sScaleB[read_stage];
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
-        tCrScaleAViewAsC.data()[i] = tCsScaleAViewAsC(_, _, _, read_stage)(idx2crd(i, RegLayoutScaleAEssential{}));
+        tCrScaleAViewAsC.data()[i] = tCsScaleAViewAsC(
+            _, _, _, read_stage)(idx2crd(i, RegLayoutScaleAEssential{}));
       }
       if constexpr (ScaleMsPerTile == 1) {
         static_assert(size(RegLayoutScaleAEssential{}) == 1);
-        tCrScaleAViewAsC.data()[0] = __shfl_sync(0xffffffff, tCrScaleAViewAsC.data()[0] * scale_b, 0); // `tCrScaleAViewAsC.data()[0]` are all same in a warp group when `ScaleMsPerTile == 1`.
+        tCrScaleAViewAsC.data()[0] =
+            __shfl_sync(0xffffffff,
+                        tCrScaleAViewAsC.data()[0] * scale_b,
+                        0);  // `tCrScaleAViewAsC.data()[0]` are all same in a
+                             // warp group when `ScaleMsPerTile == 1`.
       } else {
         CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
@@ -694,19 +844,25 @@ struct CollectiveMma<
       CUTLASS_PRAGMA_UNROLL
       for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
         // (V,M,K) x (V,N,K) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
+        cute::gemm(tiled_mma,
+                   tCrA(_, _, k_block, read_stage),
+                   tCrB(_, _, k_block, read_stage),
+                   accumulation());
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
       }
       warpgroup_commit_batch();
 
-      /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
+      /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to
+      /// ensure smem_pipe_write is consumed
       warpgroup_wait<K_PIPE_MMAS>();
       warpgroup_fence_operand(accumulation());
 
       // Block scale the accumulators with reg tensor `tCrScaleAViewAsC`
       accumulation.scale_if_needed(tCrScaleAViewAsC);
 
-      pipeline.consumer_release(smem_pipe_release);                 // UNLOCK smem_pipe_release, done _computing_ on it
+      pipeline.consumer_release(
+          smem_pipe_release);  // UNLOCK smem_pipe_release, done _computing_ on
+                               // it
 
       // Advance smem_pipe_read and smem_pipe_release
       ++smem_pipe_read;
@@ -719,8 +875,9 @@ struct CollectiveMma<
   }
 
   /// Perform a Consumer Epilogue to release all buffers
-  CUTLASS_DEVICE void
-  mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
+  CUTLASS_DEVICE void mma_tail(MainloopPipeline pipeline,
+                               PipelineState smem_pipe_release,
+                               int k_tile_count) {
     // Prologue GMMAs
     int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
     k_tile_count -= prologue_mma_count;
@@ -731,7 +888,9 @@ struct CollectiveMma<
     warpgroup_wait<0>();
 
     for (int count = 0; count < prologue_mma_count; ++count) {
-      pipeline.consumer_release(smem_pipe_release);                 // UNLOCK smem_pipe_release, done _computing_ on it
+      pipeline.consumer_release(
+          smem_pipe_release);  // UNLOCK smem_pipe_release, done _computing_ on
+                               // it
       ++smem_pipe_release;
     }
   }
@@ -739,6 +898,6 @@ struct CollectiveMma<
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm::collective
+}  // namespace cutlass::gemm::collective
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/device/gemm_universal_base_compat.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/device/gemm_universal_base_compat.h
index 2edd5a228b..ae5ff6decc 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/device/gemm_universal_base_compat.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/device/gemm_universal_base_compat.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,20 +18,21 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*!
   \file
-  \brief The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and
-    batched array variants.
+  \brief The universal GEMM accommodates serial reductions, parallel reductions,
+  batched strided, and batched array variants.
 */
 
 #pragma once
@@ -54,385 +55,363 @@
 
 ////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace device
-{
+namespace cutlass {
+namespace gemm {
+namespace device {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 /*
-    This is the device layer from CUTLASS 2.10 (SHA - cc85b64cf676c45f98a17e3a47c0aafcf817f088)
-    It is replicated here since we needed to duplicate kernel level APIs for mixed dtype GEMMs
-    and SmoothQuant. The newer device layer is not compatible with these older kernel level APIs.
+    This is the device layer from CUTLASS 2.10 (SHA -
+   cc85b64cf676c45f98a17e3a47c0aafcf817f088) It is replicated here since we
+   needed to duplicate kernel level APIs for mixed dtype GEMMs and SmoothQuant.
+   The newer device layer is not compatible with these older kernel level APIs.
 
-    Note: While CUTLASS 3.x supports stream-k, none of the kernels in the extensions folder support
-          that feature at the moment.
+    Note: While CUTLASS 3.x supports stream-k, none of the kernels in the
+   extensions folder support that feature at the moment.
   */
 
 template <typename GemmKernel_>
-class GemmUniversalBaseCompat
-{
-public:
-    using GemmKernel = GemmKernel_;
-    using ThreadblockShape = typename GemmKernel::Mma::Shape;
+class GemmUniversalBaseCompat {
+ public:
+  using GemmKernel = GemmKernel_;
+  using ThreadblockShape = typename GemmKernel::Mma::Shape;
 
-    using ElementA = typename GemmKernel::ElementA;
-    using LayoutA = typename GemmKernel::LayoutA;
-    using TensorRefA = TensorRef<ElementA const, LayoutA>;
-    static ComplexTransform const kTransformA = GemmKernel::kTransformA;
+  using ElementA = typename GemmKernel::ElementA;
+  using LayoutA = typename GemmKernel::LayoutA;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  static ComplexTransform const kTransformA = GemmKernel::kTransformA;
 
-    using ElementB = typename GemmKernel::ElementB;
-    using LayoutB = typename GemmKernel::LayoutB;
-    using TensorRefB = TensorRef<ElementB const, LayoutB>;
-    static ComplexTransform const kTransformB = GemmKernel::kTransformB;
+  using ElementB = typename GemmKernel::ElementB;
+  using LayoutB = typename GemmKernel::LayoutB;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  static ComplexTransform const kTransformB = GemmKernel::kTransformB;
 
-    using ElementC = typename GemmKernel::ElementC;
-    using LayoutC = typename GemmKernel::LayoutC;
-    using TensorRefC = TensorRef<ElementC const, LayoutC>;
-    using TensorRefD = TensorRef<ElementC, LayoutC>;
+  using ElementC = typename GemmKernel::ElementC;
+  using LayoutC = typename GemmKernel::LayoutC;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
 
-    using ElementAccumulator = typename GemmKernel::Mma::Policy::Operator::ElementC;
+  using ElementAccumulator =
+      typename GemmKernel::Mma::Policy::Operator::ElementC;
 
-    using EpilogueOutputOp = typename GemmKernel::EpilogueOutputOp;
-    using ThreadblockSwizzle = typename GemmKernel::ThreadblockSwizzle;
-    using Operator = typename GemmKernel::Operator;
+  using EpilogueOutputOp = typename GemmKernel::EpilogueOutputOp;
+  using ThreadblockSwizzle = typename GemmKernel::ThreadblockSwizzle;
+  using Operator = typename GemmKernel::Operator;
 
-    /// Argument structure
-    using Arguments = typename GemmKernel::Arguments;
+  /// Argument structure
+  using Arguments = typename GemmKernel::Arguments;
 
-protected:
-    /// Kernel parameters object
-    typename GemmKernel::Params params_;
+ protected:
+  /// Kernel parameters object
+  typename GemmKernel::Params params_;
 
-protected:
-    /// Private helper to obtain the grid dimensions with fix-up for split-K
-    static void get_grid_shape_(gemm::GemmCoord& grid_tiled_shape, int& gemm_k_size, Arguments const& args)
-    {
+ protected:
+  /// Private helper to obtain the grid dimensions with fix-up for split-K
+  static void get_grid_shape_(gemm::GemmCoord& grid_tiled_shape,
+                              int& gemm_k_size,
+                              Arguments const& args) {
+    // Determine grid shape
+    ThreadblockSwizzle threadblock_swizzle;
 
-        // Determine grid shape
-        ThreadblockSwizzle threadblock_swizzle;
+    grid_tiled_shape = threadblock_swizzle.get_tiled_shape(
+        args.problem_size,
+        {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+        args.batch_count);
 
-        grid_tiled_shape = threadblock_swizzle.get_tiled_shape(
-            args.problem_size, {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, args.batch_count);
+    gemm_k_size = args.problem_size.k();
 
-        gemm_k_size = args.problem_size.k();
+    if (args.mode == GemmUniversalMode::kGemm ||
+        args.mode == GemmUniversalMode::kGemmSplitKParallel) {
+      int const kAlignK =
+          const_max(const_max(128 / sizeof_bits<ElementA>::value,
+                              128 / sizeof_bits<ElementB>::value),
+                    1);
 
-        if (args.mode == GemmUniversalMode::kGemm || args.mode == GemmUniversalMode::kGemmSplitKParallel)
-        {
+      gemm_k_size =
+          round_up(ceil_div(args.problem_size.k(), args.batch_count), kAlignK);
 
-            int const kAlignK
-                = const_max(const_max(128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value), 1);
+      if (gemm_k_size) {
+        grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size);
+      }
+    }
+  }
 
-            gemm_k_size = round_up(ceil_div(args.problem_size.k(), args.batch_count), kAlignK);
+ public:
+  /// Constructs the GEMM.
+  GemmUniversalBaseCompat() {}
 
-            if (gemm_k_size)
-            {
-                grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size);
-            }
-        }
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const& args) {
+    // Determine grid shape
+    cutlass::gemm::GemmCoord grid_tiled_shape;
+    int gemm_k_size = 0;
+
+    get_grid_shape_(grid_tiled_shape, gemm_k_size, args);
+
+    ThreadblockSwizzle threadblock_swizzle;
+    dim3 grid = threadblock_swizzle.get_grid_shape(grid_tiled_shape);
+
+    uint32_t const kGridYZMax = ((1 << (sizeof(uint16_t) * 8)) - 1);
+
+    if (!(grid.y <= kGridYZMax && grid.z <= kGridYZMax)) {
+      return Status::kErrorInvalidProblem;
     }
 
-public:
-    /// Constructs the GEMM.
-    GemmUniversalBaseCompat() {}
+    return GemmKernel::can_implement(args);
+  }
 
-    /// Determines whether the GEMM can execute the given problem.
-    static Status can_implement(Arguments const& args)
-    {
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const& args) {
+    CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::get_workspace_size()");
 
-        // Determine grid shape
-        cutlass::gemm::GemmCoord grid_tiled_shape;
-        int gemm_k_size = 0;
+    size_t workspace_bytes = 0;
 
-        get_grid_shape_(grid_tiled_shape, gemm_k_size, args);
+    // Determine grid shape
+    cutlass::gemm::GemmCoord grid_tiled_shape;
+    int gemm_k_size = 0;
 
-        ThreadblockSwizzle threadblock_swizzle;
-        dim3 grid = threadblock_swizzle.get_grid_shape(grid_tiled_shape);
+    get_grid_shape_(grid_tiled_shape, gemm_k_size, args);
 
-        uint32_t const kGridYZMax = ((1 << (sizeof(uint16_t) * 8)) - 1);
-
-        if (!(grid.y <= kGridYZMax && grid.z <= kGridYZMax))
-        {
-
-            return Status::kErrorInvalidProblem;
-        }
-
-        return GemmKernel::can_implement(args);
+    if (args.mode == GemmUniversalMode::kGemmSplitKParallel) {
+      // Split-K parallel always requires a temporary workspace
+      workspace_bytes = sizeof(ElementC) * size_t(args.batch_stride_D) *
+                        size_t(grid_tiled_shape.k());
+    } else if (args.mode == GemmUniversalMode::kGemm &&
+               grid_tiled_shape.k() > 1) {
+      // Serial split-K only requires a temporary workspace if the number of
+      // partitions along the GEMM K dimension is greater than one.
+      workspace_bytes = sizeof(int) * size_t(grid_tiled_shape.m()) *
+                        size_t(grid_tiled_shape.n());
     }
 
-    /// Gets the workspace size
-    static size_t get_workspace_size(Arguments const& args)
-    {
+    CUTLASS_TRACE_HOST("  workspace_bytes: " << workspace_bytes);
 
-        CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::get_workspace_size()");
+    workspace_bytes +=
+        GemmKernel::get_extra_workspace_size(args, grid_tiled_shape);
 
-        size_t workspace_bytes = 0;
+    return workspace_bytes;
+  }
 
-        // Determine grid shape
-        cutlass::gemm::GemmCoord grid_tiled_shape;
-        int gemm_k_size = 0;
+  /// Computes the grid shape
+  static dim3 get_grid_shape(Arguments const& args) {
+    CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::get_grid_shape()");
 
-        get_grid_shape_(grid_tiled_shape, gemm_k_size, args);
+    ThreadblockSwizzle threadblock_swizzle;
 
-        if (args.mode == GemmUniversalMode::kGemmSplitKParallel)
-        {
+    cutlass::gemm::GemmCoord grid_tiled_shape;
+    int gemm_k_size = 0;
 
-            // Split-K parallel always requires a temporary workspace
-            workspace_bytes = sizeof(ElementC) * size_t(args.batch_stride_D) * size_t(grid_tiled_shape.k());
-        }
-        else if (args.mode == GemmUniversalMode::kGemm && grid_tiled_shape.k() > 1)
-        {
+    get_grid_shape_(grid_tiled_shape, gemm_k_size, args);
+    dim3 result = threadblock_swizzle.get_grid_shape(grid_tiled_shape);
 
-            // Serial split-K only requires a temporary workspace if the number of partitions along the
-            // GEMM K dimension is greater than one.
-            workspace_bytes = sizeof(int) * size_t(grid_tiled_shape.m()) * size_t(grid_tiled_shape.n());
-        }
+    CUTLASS_TRACE_HOST("  grid_tiled_shape: " << grid_tiled_shape << "\n"
+                                              << "  result = {" << result
+                                              << "}");
 
-        CUTLASS_TRACE_HOST("  workspace_bytes: " << workspace_bytes);
+    return result;
+  }
 
-        workspace_bytes += GemmKernel::get_extra_workspace_size(args, grid_tiled_shape);
+  /// Computes the maximum number of active blocks per multiprocessor
+  static int maximum_active_blocks(int smem_capacity = -1) {
+    CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::maximum_active_blocks()");
 
-        return workspace_bytes;
-    }
+    int max_active_blocks = -1;
+    int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
 
-    /// Computes the grid shape
-    static dim3 get_grid_shape(Arguments const& args)
-    {
+    CUTLASS_TRACE_HOST("  smem_size: " << smem_size << " bytes");
 
-        CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::get_grid_shape()");
+    if (smem_size <= (48 << 10)) {
+      cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+          &max_active_blocks,
+          Kernel<GemmKernel>,
+          GemmKernel::kThreadCount,
+          smem_size);
 
-        ThreadblockSwizzle threadblock_swizzle;
+      if (result == cudaSuccess) {
+        CUTLASS_TRACE_HOST("  max_active_blocks: " << max_active_blocks);
+        return max_active_blocks;
+      }
+    } else {
+      // Query assuming zero shared memory then compute occupancy limit based on
+      // SMEM
+      cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+          &max_active_blocks, Kernel<GemmKernel>, GemmKernel::kThreadCount, 0);
 
-        cutlass::gemm::GemmCoord grid_tiled_shape;
-        int gemm_k_size = 0;
-
-        get_grid_shape_(grid_tiled_shape, gemm_k_size, args);
-        dim3 result = threadblock_swizzle.get_grid_shape(grid_tiled_shape);
-
-        CUTLASS_TRACE_HOST("  grid_tiled_shape: " << grid_tiled_shape << "\n"
-                                                  << "  result = {" << result << "}");
-
-        return result;
-    }
-
-    /// Computes the maximum number of active blocks per multiprocessor
-    static int maximum_active_blocks(int smem_capacity = -1)
-    {
-
-        CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::maximum_active_blocks()");
-
-        int max_active_blocks = -1;
-        int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
-
-        CUTLASS_TRACE_HOST("  smem_size: " << smem_size << " bytes");
-
-        if (smem_size <= (48 << 10))
-        {
-
-            cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-                &max_active_blocks, Kernel<GemmKernel>, GemmKernel::kThreadCount, smem_size);
-
-            if (result == cudaSuccess)
-            {
-                CUTLASS_TRACE_HOST("  max_active_blocks: " << max_active_blocks);
-                return max_active_blocks;
-            }
-        }
-        else
-        {
-
-            // Query assuming zero shared memory then compute occupancy limit based on SMEM
-            cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-                &max_active_blocks, Kernel<GemmKernel>, GemmKernel::kThreadCount, 0);
-
-            if (result != cudaSuccess)
-            {
-
-                CUTLASS_TRACE_HOST(
-                    "  cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error " << cudaGetErrorString(result));
-
-                return -1;
-            }
-
-            if (smem_capacity < 0)
-            {
-                int device_idx = 0;
-                result = cudaGetDevice(&device_idx);
-
-                if (result != cudaSuccess)
-                {
-                    return -1;
-                }
-
-                cudaDeviceProp properties;
-                result = cudaGetDeviceProperties(&properties, device_idx);
-
-                if (result != cudaSuccess)
-                {
-                    return -1;
-                }
-
-                smem_capacity = static_cast<int>(properties.sharedMemPerMultiprocessor);
-            }
-
-            int occupancy = std::min(max_active_blocks, smem_capacity / smem_size);
-
-            CUTLASS_TRACE_HOST("  occupancy: " << occupancy);
-
-            return occupancy;
-        }
-
-        CUTLASS_TRACE_HOST("  returning internal error");
+      if (result != cudaSuccess) {
+        CUTLASS_TRACE_HOST(
+            "  cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error "
+            << cudaGetErrorString(result));
 
         return -1;
-    }
+      }
 
-    /// Initializes GEMM state from arguments.
-    Status initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr)
-    {
+      if (smem_capacity < 0) {
+        int device_idx = 0;
+        result = cudaGetDevice(&device_idx);
 
-        CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::initialize() - workspace "
-            << workspace << ", stream: " << (stream ? "non-null" : "null"));
-
-        size_t workspace_bytes = get_workspace_size(args);
-
-        CUTLASS_TRACE_HOST("  workspace_bytes: " << workspace_bytes);
-
-        if (workspace_bytes)
-        {
-
-            if (!workspace)
-            {
-                CUTLASS_TRACE_HOST("  error: device workspace must not be null");
-
-                return Status::kErrorWorkspaceNull;
-            }
-
-            if (args.mode == GemmUniversalMode::kGemm)
-            {
-                CUTLASS_TRACE_HOST("  clearing device workspace");
-                cudaError_t result = cudaMemsetAsync(workspace, 0, workspace_bytes, stream);
-
-                if (result != cudaSuccess)
-                {
-                    CUTLASS_TRACE_HOST("  cudaMemsetAsync() returned error " << cudaGetErrorString(result));
-
-                    return Status::kErrorInternal;
-                }
-            }
+        if (result != cudaSuccess) {
+          return -1;
         }
 
-        // Get CUDA grid shape
-        cutlass::gemm::GemmCoord grid_tiled_shape;
-        int gemm_k_size = 0;
+        cudaDeviceProp properties;
+        result = cudaGetDeviceProperties(&properties, device_idx);
 
-        get_grid_shape_(grid_tiled_shape, gemm_k_size, args);
-
-        // Initialize the Params structure
-        params_ = typename GemmKernel::Params(args, grid_tiled_shape, gemm_k_size, static_cast<int*>(workspace));
-
-        // Specify shared memory capacity for kernel.
-        int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
-
-        if (smem_size >= (48 << 10))
-        {
-            cudaError_t result
-                = cudaFuncSetAttribute(Kernel<GemmKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-
-            if (result != cudaSuccess)
-            {
-                return Status::kErrorInternal;
-            }
+        if (result != cudaSuccess) {
+          return -1;
         }
 
-        return Status::kSuccess;
+        smem_capacity = static_cast<int>(properties.sharedMemPerMultiprocessor);
+      }
+
+      int occupancy = std::min(max_active_blocks, smem_capacity / smem_size);
+
+      CUTLASS_TRACE_HOST("  occupancy: " << occupancy);
+
+      return occupancy;
     }
 
-    /// Lightweight update given a subset of arguments
-    Status update(Arguments const& args, void* workspace = nullptr)
-    {
+    CUTLASS_TRACE_HOST("  returning internal error");
 
-        CUTLASS_TRACE_HOST("GemmUniversalBaseCompat()::update() - workspace: " << workspace);
+    return -1;
+  }
 
-        size_t workspace_bytes = get_workspace_size(args);
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const& args,
+                    void* workspace = nullptr,
+                    cudaStream_t stream = nullptr) {
+    CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::initialize() - workspace "
+                       << workspace
+                       << ", stream: " << (stream ? "non-null" : "null"));
 
-        if (workspace_bytes && !workspace)
-        {
-            return Status::kErrorWorkspaceNull;
+    size_t workspace_bytes = get_workspace_size(args);
+
+    CUTLASS_TRACE_HOST("  workspace_bytes: " << workspace_bytes);
+
+    if (workspace_bytes) {
+      if (!workspace) {
+        CUTLASS_TRACE_HOST("  error: device workspace must not be null");
+
+        return Status::kErrorWorkspaceNull;
+      }
+
+      if (args.mode == GemmUniversalMode::kGemm) {
+        CUTLASS_TRACE_HOST("  clearing device workspace");
+        cudaError_t result =
+            cudaMemsetAsync(workspace, 0, workspace_bytes, stream);
+
+        if (result != cudaSuccess) {
+          CUTLASS_TRACE_HOST("  cudaMemsetAsync() returned error "
+                             << cudaGetErrorString(result));
+
+          return Status::kErrorInternal;
         }
-
-        params_.update(args, workspace);
-
-        return Status::kSuccess;
+      }
     }
 
-    /// Runs the kernel using initialized state.
-    Status run(cudaStream_t stream = nullptr)
-    {
-        CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::run()");
+    // Get CUDA grid shape
+    cutlass::gemm::GemmCoord grid_tiled_shape;
+    int gemm_k_size = 0;
 
-        //
-        // Configure grid and block dimensions
-        //
+    get_grid_shape_(grid_tiled_shape, gemm_k_size, args);
 
-        ThreadblockSwizzle threadblock_swizzle;
+    // Initialize the Params structure
+    params_ = typename GemmKernel::Params(
+        args, grid_tiled_shape, gemm_k_size, static_cast<int*>(workspace));
 
-        dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);
-        dim3 block(GemmKernel::kThreadCount, 1, 1);
+    // Specify shared memory capacity for kernel.
+    int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
 
-        int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
+    if (smem_size >= (48 << 10)) {
+      cudaError_t result =
+          cudaFuncSetAttribute(Kernel<GemmKernel>,
+                               cudaFuncAttributeMaxDynamicSharedMemorySize,
+                               smem_size);
 
-        //
-        // Launch kernel
-        //
-
-        CUTLASS_TRACE_HOST("  grid: (" << grid << "),  block: (" << block << "),  SMEM: " << smem_size << " bytes");
-
-        // Launch
-        cutlass::Kernel<GemmKernel><<<grid, block, smem_size, stream>>>(params_);
-
-        //
-        // Query for errors
-        //
-        cudaError_t result = cudaGetLastError();
-
-        if (result != cudaSuccess)
-        {
-            CUTLASS_TRACE_HOST("  grid launch failed with error " << cudaGetErrorString(result));
-            return Status::kErrorInternal;
-        }
-
-        return Status::kSuccess;
+      if (result != cudaSuccess) {
+        return Status::kErrorInternal;
+      }
     }
 
-    /// Runs the kernel using initialized state.
-    Status operator()(cudaStream_t stream = nullptr)
-    {
-        return run(stream);
+    return Status::kSuccess;
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const& args, void* workspace = nullptr) {
+    CUTLASS_TRACE_HOST(
+        "GemmUniversalBaseCompat()::update() - workspace: " << workspace);
+
+    size_t workspace_bytes = get_workspace_size(args);
+
+    if (workspace_bytes && !workspace) {
+      return Status::kErrorWorkspaceNull;
     }
 
-    /// Runs the kernel using initialized state.
-    Status operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr)
-    {
+    params_.update(args, workspace);
 
-        Status status = initialize(args, workspace, stream);
+    return Status::kSuccess;
+  }
 
-        if (status == Status::kSuccess)
-        {
-            status = run(stream);
-        }
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+    CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::run()");
 
-        return status;
+    //
+    // Configure grid and block dimensions
+    //
+
+    ThreadblockSwizzle threadblock_swizzle;
+
+    dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);
+    dim3 block(GemmKernel::kThreadCount, 1, 1);
+
+    int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
+
+    //
+    // Launch kernel
+    //
+
+    CUTLASS_TRACE_HOST("  grid: (" << grid << "),  block: (" << block
+                                   << "),  SMEM: " << smem_size << " bytes");
+
+    // Launch
+    cutlass::Kernel<GemmKernel><<<grid, block, smem_size, stream>>>(params_);
+
+    //
+    // Query for errors
+    //
+    cudaError_t result = cudaGetLastError();
+
+    if (result != cudaSuccess) {
+      CUTLASS_TRACE_HOST("  grid launch failed with error "
+                         << cudaGetErrorString(result));
+      return Status::kErrorInternal;
     }
+
+    return Status::kSuccess;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(Arguments const& args,
+                    void* workspace = nullptr,
+                    cudaStream_t stream = nullptr) {
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace device
-} // namespace gemm
-} // namespace cutlass
+}  // namespace device
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/device/splitk_gemm_grouped.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/device/splitk_gemm_grouped.h
index bfd3666b9c..02859ee0d1 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/device/splitk_gemm_grouped.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/device/splitk_gemm_grouped.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*!
@@ -55,488 +56,479 @@
 
 ////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace device
-{
+namespace cutlass {
+namespace gemm {
+namespace device {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <typename T_IN, typename T_OUT>
-__global__ void splitkReduction(T_OUT** out_tensor, const T_IN* in_tensor, GemmCoord const* problem_sizes, int splitk,
-    int64_t* splitk_buffer_offsets)
-{
-    // in_tensor: [problem_idx, k_partition, hidden_size]
-    //      Note that different requests of in_tensor might have different hidden_size (=m*n)
-    //      so, we need to use splitk_buffer_offsets.
-    // out_tensor: problem_idx * [hidden_size]
+__global__ void splitkReduction(T_OUT** out_tensor,
+                                const T_IN* in_tensor,
+                                GemmCoord const* problem_sizes,
+                                int splitk,
+                                int64_t* splitk_buffer_offsets) {
+  // in_tensor: [problem_idx, k_partition, hidden_size]
+  //      Note that different requests of in_tensor might have different
+  //      hidden_size (=m*n) so, we need to use splitk_buffer_offsets.
+  // out_tensor: problem_idx * [hidden_size]
 
-    int const problem_idx = blockIdx.y;
-    GemmCoord problem = problem_sizes[problem_idx];
-    int const hidden_size = problem.m() * problem.n();
-    const T_IN* in_tensor_ = in_tensor + splitk_buffer_offsets[problem_idx] * splitk;
-    T_OUT* out_tensor_ = out_tensor[problem_idx];
+  int const problem_idx = blockIdx.y;
+  GemmCoord problem = problem_sizes[problem_idx];
+  int const hidden_size = problem.m() * problem.n();
+  const T_IN* in_tensor_ =
+      in_tensor + splitk_buffer_offsets[problem_idx] * splitk;
+  T_OUT* out_tensor_ = out_tensor[problem_idx];
 
-    for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < hidden_size; i += blockDim.x * gridDim.x)
-    {
-        float sum = 0.0f;
-        for (int k_idx = 0; k_idx < splitk; k_idx++)
-        {
-            sum += (float) in_tensor_[k_idx * hidden_size + i];
-        }
-        out_tensor_[i] = (T_OUT) (sum);
+  for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < hidden_size;
+       i += blockDim.x * gridDim.x) {
+    float sum = 0.0f;
+    for (int k_idx = 0; k_idx < splitk; k_idx++) {
+      sum += (float)in_tensor_[k_idx * hidden_size + i];
     }
+    out_tensor_[i] = (T_OUT)(sum);
+  }
 }
 
 /// GEMM Grouped
 template <typename BaseKernel_>
-class BaseSplitkGrouped
-{
-public:
-    using BaseKernel = BaseKernel_;
+class BaseSplitkGrouped {
+ public:
+  using BaseKernel = BaseKernel_;
 
-    using ElementA = typename BaseKernel::ElementA;
-    using LayoutA = typename BaseKernel::LayoutA;
-    using TensorRefA = TensorRef<ElementA const, LayoutA>;
-    static ComplexTransform const kTransformA = BaseKernel::kTransformA;
-    static int const kAlignmentA = BaseKernel::kAlignmentA;
+  using ElementA = typename BaseKernel::ElementA;
+  using LayoutA = typename BaseKernel::LayoutA;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  static ComplexTransform const kTransformA = BaseKernel::kTransformA;
+  static int const kAlignmentA = BaseKernel::kAlignmentA;
 
-    using ElementB = typename BaseKernel::ElementB;
-    using LayoutB = typename BaseKernel::LayoutB;
-    using TensorRefB = TensorRef<ElementB const, LayoutB>;
-    static ComplexTransform const kTransformB = BaseKernel::kTransformB;
-    static int const kAlignmentB = BaseKernel::kAlignmentB;
+  using ElementB = typename BaseKernel::ElementB;
+  using LayoutB = typename BaseKernel::LayoutB;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  static ComplexTransform const kTransformB = BaseKernel::kTransformB;
+  static int const kAlignmentB = BaseKernel::kAlignmentB;
 
-    using ElementC = typename BaseKernel::ElementC;
-    using LayoutC = typename BaseKernel::LayoutC;
-    using TensorRefC = TensorRef<ElementC const, LayoutC>;
-    using TensorRefD = TensorRef<ElementC, LayoutC>;
-    static int const kAlignmentC = BaseKernel::kAlignmentC;
+  using ElementC = typename BaseKernel::ElementC;
+  using LayoutC = typename BaseKernel::LayoutC;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
+  static int const kAlignmentC = BaseKernel::kAlignmentC;
 
-    using ElementAccumulator = typename BaseKernel::Mma::Policy::Operator::ElementC;
+  using ElementAccumulator =
+      typename BaseKernel::Mma::Policy::Operator::ElementC;
 
-    using EpilogueOutputOp = typename BaseKernel::EpilogueOutputOp;
-    using ThreadblockSwizzle = typename threadblock::GemmSplitKHorizontalThreadblockSwizzle;
+  using EpilogueOutputOp = typename BaseKernel::EpilogueOutputOp;
+  using ThreadblockSwizzle =
+      typename threadblock::GemmSplitKHorizontalThreadblockSwizzle;
 
-    using Operator = typename BaseKernel::Operator;
-    using WarpMmaOperator = typename BaseKernel::Mma::Policy::Operator;
+  using Operator = typename BaseKernel::Operator;
+  using WarpMmaOperator = typename BaseKernel::Mma::Policy::Operator;
 
-    using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator;
-    using MathOperator = typename WarpMmaOperator::MathOperator;
-    using OperatorClass = typename WarpMmaOperator::OperatorClass;
-    using ArchTag = typename WarpMmaOperator::ArchTag;
-    using ThreadblockShape = typename BaseKernel::Mma::Shape;
-    using WarpShape = typename BaseKernel::WarpShape;
-    using InstructionShape = typename BaseKernel::InstructionShape;
-    static int const kStages = BaseKernel::Mma::kStages;
+  using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator;
+  using MathOperator = typename WarpMmaOperator::MathOperator;
+  using OperatorClass = typename WarpMmaOperator::OperatorClass;
+  using ArchTag = typename WarpMmaOperator::ArchTag;
+  using ThreadblockShape = typename BaseKernel::Mma::Shape;
+  using WarpShape = typename BaseKernel::WarpShape;
+  using InstructionShape = typename BaseKernel::InstructionShape;
+  static int const kStages = BaseKernel::Mma::kStages;
 
-    /// Argument structure
-    using Arguments = typename BaseKernel::Arguments;
+  /// Argument structure
+  using Arguments = typename BaseKernel::Arguments;
 
-    using ProblemInfo = typename BaseKernel::ProblemVisitor::ProblemInfo;
+  using ProblemInfo = typename BaseKernel::ProblemVisitor::ProblemInfo;
 
-protected:
-    /// Kernel parameters object
-    typename BaseKernel::Params gemm_params_;
+ protected:
+  /// Kernel parameters object
+  typename BaseKernel::Params gemm_params_;
 
-private:
-    /// Get the number of tiles across all problems in a group
-    static int32_t group_tile_count(cutlass::gemm::GemmCoord const* problem_sizes_ptr, int problem_count)
-    {
-        int32_t tiles = 0;
-        for (int32_t i = 0; i < problem_count; ++i)
-        {
-            cutlass::gemm::GemmCoord problem = problem_sizes_ptr[i];
-            BaseKernel::ProblemVisitor::possibly_transpose_problem(problem);
-            tiles += problem_tile_count(problem);
-        }
-        return tiles;
+ private:
+  /// Get the number of tiles across all problems in a group
+  static int32_t group_tile_count(
+      cutlass::gemm::GemmCoord const* problem_sizes_ptr, int problem_count) {
+    int32_t tiles = 0;
+    for (int32_t i = 0; i < problem_count; ++i) {
+      cutlass::gemm::GemmCoord problem = problem_sizes_ptr[i];
+      BaseKernel::ProblemVisitor::possibly_transpose_problem(problem);
+      tiles += problem_tile_count(problem);
+    }
+    return tiles;
+  }
+
+  /// Copy from `data` to `workspace`
+  Status copy_to_workspace(void* workspace, void* data, size_t bytes) {
+    cudaError_t cuda_error =
+        cudaMemcpy(workspace, data, bytes, cudaMemcpyHostToDevice);
+    if (cuda_error != cudaSuccess) {
+      // Call cudaGetLastError() to clear the error bit
+      cuda_error = cudaGetLastError();
+      CUTLASS_TRACE_HOST("  cudaMemcpy() returned error "
+                         << cudaGetErrorString(cuda_error));
+      return Status::kErrorInternal;
     }
 
-    /// Copy from `data` to `workspace`
-    Status copy_to_workspace(void* workspace, void* data, size_t bytes)
-    {
-        cudaError_t cuda_error = cudaMemcpy(workspace, data, bytes, cudaMemcpyHostToDevice);
-        if (cuda_error != cudaSuccess)
-        {
-            // Call cudaGetLastError() to clear the error bit
-            cuda_error = cudaGetLastError();
-            CUTLASS_TRACE_HOST("  cudaMemcpy() returned error " << cudaGetErrorString(cuda_error));
-            return Status::kErrorInternal;
-        }
+    return Status::kSuccess;
+  }
 
-        return Status::kSuccess;
+  /// Precomputes scheduling information for the grouped GEMM
+  Status precompute(Arguments const& args,
+                    int32_t tile_count,
+                    void* workspace) {
+    size_t workspace_bytes = get_workspace_size(args);
+    std::vector<uint8_t> host_workspace(workspace_bytes);
+    BaseKernel::ProblemVisitor::host_precompute(args.host_problem_sizes,
+                                                args.problem_count,
+                                                args.threadblock_count,
+                                                (void*)host_workspace.data());
+    return copy_to_workspace(workspace, host_workspace.data(), workspace_bytes);
+  }
+
+  /// Reorder `data` according to `indices`
+  template <typename T>
+  static void reorder_array(T* data, std::vector<size_t> const& indices) {
+    // For now, simply create a copy of the data and then copy over to the
+    // original.
+    std::vector<T> copy(indices.size());
+    for (size_t i = 0; i < indices.size(); ++i) {
+      copy.at(i) = data[indices[i]];
     }
 
-    /// Precomputes scheduling information for the grouped GEMM
-    Status precompute(Arguments const& args, int32_t tile_count, void* workspace)
-    {
-        size_t workspace_bytes = get_workspace_size(args);
-        std::vector<uint8_t> host_workspace(workspace_bytes);
-        BaseKernel::ProblemVisitor::host_precompute(
-            args.host_problem_sizes, args.problem_count, args.threadblock_count, (void*) host_workspace.data());
-        return copy_to_workspace(workspace, host_workspace.data(), workspace_bytes);
+    memcpy(data, copy.data(), indices.size() * sizeof(T));
+  }
+
+ public:
+  /// Constructs the GEMM.
+  BaseSplitkGrouped() {}
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const& args) {
+    return BaseKernel::can_implement(args);
+  }
+
+  /// Get the number of tiles in a problem
+  static int32_t problem_tile_count(cutlass::gemm::GemmCoord const& problem) {
+    auto grid = BaseKernel::ProblemVisitor::grid_shape(problem);
+    return BaseKernel::ProblemVisitor::tile_count(grid);
+  }
+
+  /// Get the number of tiles across all problems in a group
+  static int32_t group_tile_count(Arguments const& args) {
+    if (args.host_problem_sizes == nullptr) {
+      CUTLASS_TRACE_HOST("Received nullptr for `args.host_problem_sizes");
+      return -1;
     }
 
-    /// Reorder `data` according to `indices`
-    template <typename T>
-    static void reorder_array(T* data, std::vector<size_t> const& indices)
-    {
-        // For now, simply create a copy of the data and then copy over to the original.
-        std::vector<T> copy(indices.size());
-        for (size_t i = 0; i < indices.size(); ++i)
-        {
-            copy.at(i) = data[indices[i]];
-        }
+    return group_tile_count(args.host_problem_sizes, args.problem_count);
+  }
 
-        memcpy(data, copy.data(), indices.size() * sizeof(T));
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const& args) {
+    size_t total_mn = 0;
+    for (int i = 0; i < args.problem_count; i++) {
+      total_mn +=
+          args.host_problem_sizes[i].m() * args.host_problem_sizes[i].n();
+    }
+    size_t workSpaceSize =
+        total_mn * sizeof(ElementAccumulator) * args.split_k_slices;
+
+    if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) {
+      workSpaceSize += BaseKernel::ProblemVisitor::get_workspace_size(
+          args.host_problem_sizes, args.problem_count, args.threadblock_count);
+    }
+    return workSpaceSize;
+  }
+
+  /// Computes the grid shape
+  static dim3 get_grid_shape(Arguments const& args) {
+    return dim3(args.threadblock_count, 1, 1);
+  }
+
+  /// Computes the maximum number of active blocks per multiprocessor
+  static int maximum_active_blocks(int smem_capacity = -1) {
+    CUTLASS_TRACE_HOST("BaseSplitkGrouped::maximum_active_blocks()");
+
+    int smem_size = int(sizeof(typename BaseKernel::SharedStorage));
+
+    CUTLASS_TRACE_HOST("  smem_size: " << smem_size << " bytes");
+
+    cudaError_t result;
+    if (smem_size > (48 << 10)) {
+      result = cudaFuncSetAttribute(Kernel<BaseKernel>,
+                                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                    smem_size);
+
+      if (result != cudaSuccess) {
+        // Call cudaGetLastError() to clear the error bit
+        result = cudaGetLastError();
+        CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error "
+                           << cudaGetErrorString(result));
+        return -1;
+      }
     }
 
-public:
-    /// Constructs the GEMM.
-    BaseSplitkGrouped() {}
+    int max_active_blocks = -1;
+    result =
+        cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_active_blocks,
+                                                      Kernel<BaseKernel>,
+                                                      BaseKernel::kThreadCount,
+                                                      smem_size);
 
-    /// Determines whether the GEMM can execute the given problem.
-    static Status can_implement(Arguments const& args)
-    {
-
-        return BaseKernel::can_implement(args);
+    if (result != cudaSuccess) {
+      // Call cudaGetLastError() to clear the error bit
+      result = cudaGetLastError();
+      CUTLASS_TRACE_HOST(
+          "  cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error "
+          << cudaGetErrorString(result));
+      return -1;
     }
 
-    /// Get the number of tiles in a problem
-    static int32_t problem_tile_count(cutlass::gemm::GemmCoord const& problem)
-    {
-        auto grid = BaseKernel::ProblemVisitor::grid_shape(problem);
-        return BaseKernel::ProblemVisitor::tile_count(grid);
+    CUTLASS_TRACE_HOST("  max_active_blocks: " << max_active_blocks);
+    return max_active_blocks;
+  }
+
+  /// Sorts each pointer passed in according to the indices that sort
+  /// `problem_sizes_ptr` in descending order of problem-K dimension.
+  static void sort_problems(int problem_count,
+                            cutlass::gemm::GemmCoord* problem_sizes_ptr,
+                            int64_t* lda_host_ptr,
+                            int64_t* ldb_host_ptr,
+                            int64_t* ldc_host_ptr,
+                            int64_t* ldd_host_ptr,
+                            int64_t* offset_A_ptr,
+                            int64_t* offset_B_ptr,
+                            int64_t* offset_C_ptr,
+                            int64_t* offset_D_ptr) {
+    std::vector<size_t> indices(problem_count);
+    std::iota(indices.begin(), indices.end(), 0);
+    std::stable_sort(indices.begin(),
+                     indices.end(),
+                     [&problem_sizes_ptr](size_t i, size_t j) {
+                       return problem_sizes_ptr[i].k() >
+                              problem_sizes_ptr[j].k();
+                     });
+
+    reorder_array(problem_sizes_ptr, indices);
+    reorder_array(lda_host_ptr, indices);
+    reorder_array(ldb_host_ptr, indices);
+    reorder_array(ldc_host_ptr, indices);
+    reorder_array(ldd_host_ptr, indices);
+    reorder_array(offset_A_ptr, indices);
+    reorder_array(offset_B_ptr, indices);
+    reorder_array(offset_C_ptr, indices);
+    reorder_array(offset_D_ptr, indices);
+  }
+
+  /// Computes the number of threadblocks to launch for the grouped kernel
+  static int sufficient(
+      cutlass::gemm::GemmCoord const* problem_sizes_ptr = nullptr,
+      int problem_count = 0,
+      int available_sm_count = -1) {
+    // Determine the number of blocks that would be launched to fill up a single
+    // wave on the GPU with each SM having maximum occupancy.
+    int device_idx;
+    cudaError_t result = cudaGetDevice(&device_idx);
+    if (result != cudaSuccess) {
+      // Call cudaGetLastError() to clear the error bit
+      result = cudaGetLastError();
+      CUTLASS_TRACE_HOST("  cudaGetDevice() returned error "
+                         << cudaGetErrorString(result));
+      return 0;
     }
 
-    /// Get the number of tiles across all problems in a group
-    static int32_t group_tile_count(Arguments const& args)
-    {
-        if (args.host_problem_sizes == nullptr)
-        {
-            CUTLASS_TRACE_HOST("Received nullptr for `args.host_problem_sizes");
-            return -1;
-        }
-
-        return group_tile_count(args.host_problem_sizes, args.problem_count);
+    int multiprocessor_count;
+    result = cudaDeviceGetAttribute(
+        &multiprocessor_count, cudaDevAttrMultiProcessorCount, device_idx);
+    if (result != cudaSuccess) {
+      CUTLASS_TRACE_HOST("  cudaDeviceGetAttribute() returned error "
+                         << cudaGetErrorString(result));
+      return 0;
     }
 
-    /// Gets the workspace size
-    static size_t get_workspace_size(Arguments const& args)
-    {
-        size_t total_mn = 0;
-        for (int i = 0; i < args.problem_count; i++)
-        {
-            total_mn += args.host_problem_sizes[i].m() * args.host_problem_sizes[i].n();
-        }
-        size_t workSpaceSize = total_mn * sizeof(ElementAccumulator) * args.split_k_slices;
-
-        if (BaseKernel::ProblemVisitor::kRequiresPrecomputation)
-        {
-            workSpaceSize += BaseKernel::ProblemVisitor::get_workspace_size(
-                args.host_problem_sizes, args.problem_count, args.threadblock_count);
-        }
-        return workSpaceSize;
+    bool override_sm_count =
+        (available_sm_count < 0 || available_sm_count > multiprocessor_count);
+    if (override_sm_count) {
+      available_sm_count = multiprocessor_count;
     }
 
-    /// Computes the grid shape
-    static dim3 get_grid_shape(Arguments const& args)
-    {
-
-        return dim3(args.threadblock_count, 1, 1);
+    int max_active_blocks = maximum_active_blocks();
+    if (max_active_blocks <= 0) {
+      return 0;
     }
 
-    /// Computes the maximum number of active blocks per multiprocessor
-    static int maximum_active_blocks(int smem_capacity = -1)
-    {
+    int occupancy_based_block_count = available_sm_count * max_active_blocks;
 
-        CUTLASS_TRACE_HOST("BaseSplitkGrouped::maximum_active_blocks()");
-
-        int smem_size = int(sizeof(typename BaseKernel::SharedStorage));
-
-        CUTLASS_TRACE_HOST("  smem_size: " << smem_size << " bytes");
-
-        cudaError_t result;
-        if (smem_size > (48 << 10))
-        {
-            result = cudaFuncSetAttribute(Kernel<BaseKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-
-            if (result != cudaSuccess)
-            {
-                // Call cudaGetLastError() to clear the error bit
-                result = cudaGetLastError();
-                CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error " << cudaGetErrorString(result));
-                return -1;
-            }
-        }
-
-        int max_active_blocks = -1;
-        result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-            &max_active_blocks, Kernel<BaseKernel>, BaseKernel::kThreadCount, smem_size);
-
-        if (result != cudaSuccess)
-        {
-            // Call cudaGetLastError() to clear the error bit
-            result = cudaGetLastError();
-            CUTLASS_TRACE_HOST(
-                "  cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error " << cudaGetErrorString(result));
-            return -1;
-        }
-
-        CUTLASS_TRACE_HOST("  max_active_blocks: " << max_active_blocks);
-        return max_active_blocks;
+    if (problem_sizes_ptr == nullptr || problem_count == 0) {
+      return occupancy_based_block_count;
     }
 
-    /// Sorts each pointer passed in according to the indices that sort
-    /// `problem_sizes_ptr` in descending order of problem-K dimension.
-    static void sort_problems(int problem_count, cutlass::gemm::GemmCoord* problem_sizes_ptr, int64_t* lda_host_ptr,
-        int64_t* ldb_host_ptr, int64_t* ldc_host_ptr, int64_t* ldd_host_ptr, int64_t* offset_A_ptr,
-        int64_t* offset_B_ptr, int64_t* offset_C_ptr, int64_t* offset_D_ptr)
-    {
-        std::vector<size_t> indices(problem_count);
-        std::iota(indices.begin(), indices.end(), 0);
-        std::stable_sort(indices.begin(), indices.end(),
-            [&problem_sizes_ptr](size_t i, size_t j) { return problem_sizes_ptr[i].k() > problem_sizes_ptr[j].k(); });
+    int total_tiles = group_tile_count(problem_sizes_ptr, problem_count);
 
-        reorder_array(problem_sizes_ptr, indices);
-        reorder_array(lda_host_ptr, indices);
-        reorder_array(ldb_host_ptr, indices);
-        reorder_array(ldc_host_ptr, indices);
-        reorder_array(ldd_host_ptr, indices);
-        reorder_array(offset_A_ptr, indices);
-        reorder_array(offset_B_ptr, indices);
-        reorder_array(offset_C_ptr, indices);
-        reorder_array(offset_D_ptr, indices);
+    // If the group contains a single problem, launching the exact number of
+    // threadblocks needed to cover the problem minimizes the work performed
+    // per threadblock in finding the next tile to compute. We return
+    // total_tiles unless the user has provided the SM count.
+    if (problem_count == 1 && override_sm_count) {
+      return total_tiles;
     }
 
-    /// Computes the number of threadblocks to launch for the grouped kernel
-    static int sufficient(
-        cutlass::gemm::GemmCoord const* problem_sizes_ptr = nullptr, int problem_count = 0, int available_sm_count = -1)
-    {
-        // Determine the number of blocks that would be launched to fill up a single
-        // wave on the GPU with each SM having maximum occupancy.
-        int device_idx;
-        cudaError_t result = cudaGetDevice(&device_idx);
-        if (result != cudaSuccess)
-        {
-            // Call cudaGetLastError() to clear the error bit
-            result = cudaGetLastError();
-            CUTLASS_TRACE_HOST("  cudaGetDevice() returned error " << cudaGetErrorString(result));
-            return 0;
-        }
+    // Choose between the full wave of threadblocks and the tile count. If there
+    // are fewer tiles in the group than threadblocks in the full wave, only
+    // some threadblocks will be assigned tiles. Those threadblocks
+    // which are not assigned tiles still need to perform the work of iterating
+    // through problem sizes to determine that they have no work to do. This
+    // competes for cycles with those threadblocks that are assigned tiles to
+    // compute.
+    return std::min(total_tiles, occupancy_based_block_count);
+  }
 
-        int multiprocessor_count;
-        result = cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device_idx);
-        if (result != cudaSuccess)
-        {
-            CUTLASS_TRACE_HOST("  cudaDeviceGetAttribute() returned error " << cudaGetErrorString(result));
-            return 0;
-        }
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const& args,
+                    void* workspace = nullptr,
+                    cudaStream_t stream = nullptr) {
+    CUTLASS_TRACE_HOST("BaseSplitkGrouped::initialize() - workspace "
+                       << workspace
+                       << ", stream: " << (stream ? "non-null" : "null"));
 
-        bool override_sm_count = (available_sm_count < 0 || available_sm_count > multiprocessor_count);
-        if (override_sm_count)
-        {
-            available_sm_count = multiprocessor_count;
-        }
+    // Workspace
+    size_t workspace_bytes = get_workspace_size(args);
 
-        int max_active_blocks = maximum_active_blocks();
-        if (max_active_blocks <= 0)
-        {
-            return 0;
-        }
-
-        int occupancy_based_block_count = available_sm_count * max_active_blocks;
-
-        if (problem_sizes_ptr == nullptr || problem_count == 0)
-        {
-            return occupancy_based_block_count;
-        }
-
-        int total_tiles = group_tile_count(problem_sizes_ptr, problem_count);
-
-        // If the group contains a single problem, launching the exact number of
-        // threadblocks needed to cover the problem minimizes the work performed
-        // per threadblock in finding the next tile to compute. We return total_tiles
-        // unless the user has provided the SM count.
-        if (problem_count == 1 && override_sm_count)
-        {
-            return total_tiles;
-        }
-
-        // Choose between the full wave of threadblocks and the tile count. If there
-        // are fewer tiles in the group than threadblocks in the full wave, only
-        // some threadblocks will be assigned tiles. Those threadblocks
-        // which are not assigned tiles still need to perform the work of iterating through
-        // problem sizes to determine that they have no work to do. This competes for cycles
-        // with those threadblocks that are assigned tiles to compute.
-        return std::min(total_tiles, occupancy_based_block_count);
+    if (workspace_bytes && !workspace) {
+      return Status::kErrorWorkspaceNull;
     }
 
-    /// Initializes GEMM state from arguments.
-    Status initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr)
-    {
-
-        CUTLASS_TRACE_HOST("BaseSplitkGrouped::initialize() - workspace "
-            << workspace << ", stream: " << (stream ? "non-null" : "null"));
-
-        // Workspace
-        size_t workspace_bytes = get_workspace_size(args);
-
-        if (workspace_bytes && !workspace)
-        {
-            return Status::kErrorWorkspaceNull;
-        }
-
-        if (BaseKernel::ProblemVisitor::kRequiresPrecomputation)
-        {
-            int32_t tile_count = group_tile_count(args);
-            Status status = precompute(args, tile_count, workspace);
-            if (status != Status::kSuccess)
-            {
-                return status;
-            }
-
-            gemm_params_ = typename BaseKernel::Params(args, workspace, tile_count);
-        }
-        else
-        {
-            gemm_params_ = typename BaseKernel::Params(args, workspace);
-        }
-
-        // Specify shared memory capacity for kernel.
-        int smem_size = int(sizeof(typename BaseKernel::SharedStorage));
-
-        if (smem_size >= (48 << 10))
-        {
-            cudaError_t result
-                = cudaFuncSetAttribute(Kernel<BaseKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-
-            if (result != cudaSuccess)
-            {
-                return Status::kErrorInternal;
-            }
-        }
-
-        return Status::kSuccess;
-    }
-
-    /// Lightweight update given a subset of arguments
-    Status update(Arguments const& args, void* workspace = nullptr)
-    {
-
-        size_t workspace_bytes = get_workspace_size(args);
-
-        if (workspace_bytes && !workspace)
-        {
-            return Status::kErrorWorkspaceNull;
-        }
-
-        if (BaseKernel::ProblemVisitor::kRequiresPrecomputation)
-        {
-            int32_t tile_count = group_tile_count(args);
-            Status status = precompute(args, tile_count, workspace);
-            if (status != Status::kSuccess)
-            {
-                return status;
-            }
-
-            gemm_params_.update(args, workspace, tile_count);
-        }
-        else
-        {
-            gemm_params_.update(args, workspace);
-        }
-
-        return Status::kSuccess;
-    }
-
-    /// Runs the kernel using initialized state.
-    Status run(cudaStream_t stream = nullptr)
-    {
-        if (!gemm_params_.problem_visitor.problem_count)
-        {
-            return Status::kSuccess;
-        }
-
-        //
-        // Launch kernel
-        //
-
-        // Launch splitk grouped gemm
-        {
-            dim3 grid(gemm_params_.threadblock_count, 1, gemm_params_.split_k_slices);
-            dim3 block(BaseKernel::kThreadCount, 1, 1);
-
-            int smem_size = int(sizeof(typename BaseKernel::SharedStorage));
-            cutlass::Kernel<BaseKernel><<<grid, block, smem_size, stream>>>(gemm_params_);
-
-            cudaError_t result = cudaGetLastError();
-
-            if (result != cudaSuccess)
-            {
-                CUTLASS_TRACE_HOST("  grid launch failed with error " << cudaGetErrorString(result));
-                return Status::kErrorInternal;
-            }
-        }
-
-        // Launch splitkReduction
-        {
-            dim3 grid(32, gemm_params_.problem_visitor.problem_count);
-            dim3 block(256);
-            splitkReduction<<<grid, block, 0, stream>>>(gemm_params_.ptr_D, gemm_params_.ptr_D_split,
-                gemm_params_.problem_visitor.problem_sizes, gemm_params_.split_k_slices,
-                gemm_params_.splitk_buffer_offsets);
-
-            cudaError_t result = cudaGetLastError();
-
-            if (result != cudaSuccess)
-            {
-                CUTLASS_TRACE_HOST("  grid launch failed with error " << cudaGetErrorString(result));
-                return Status::kErrorInternal;
-            }
-        }
-
-        return Status::kSuccess;
-    }
-
-    /// Runs the kernel using initialized state.
-    Status operator()(cudaStream_t stream = nullptr)
-    {
-        return run(stream);
-    }
-
-    /// Initializes and runs the kernel.
-    Status operator()(Arguments const& args, void* workspace, cudaStream_t stream = nullptr)
-    {
-
-        Status status = initialize(args, workspace, stream);
-
-        if (status == Status::kSuccess)
-        {
-            status = run(stream);
-        }
-
+    if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) {
+      int32_t tile_count = group_tile_count(args);
+      Status status = precompute(args, tile_count, workspace);
+      if (status != Status::kSuccess) {
         return status;
+      }
+
+      gemm_params_ = typename BaseKernel::Params(args, workspace, tile_count);
+    } else {
+      gemm_params_ = typename BaseKernel::Params(args, workspace);
     }
+
+    // Specify shared memory capacity for kernel.
+    int smem_size = int(sizeof(typename BaseKernel::SharedStorage));
+
+    if (smem_size >= (48 << 10)) {
+      cudaError_t result =
+          cudaFuncSetAttribute(Kernel<BaseKernel>,
+                               cudaFuncAttributeMaxDynamicSharedMemorySize,
+                               smem_size);
+
+      if (result != cudaSuccess) {
+        return Status::kErrorInternal;
+      }
+    }
+
+    return Status::kSuccess;
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const& args, void* workspace = nullptr) {
+    size_t workspace_bytes = get_workspace_size(args);
+
+    if (workspace_bytes && !workspace) {
+      return Status::kErrorWorkspaceNull;
+    }
+
+    if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) {
+      int32_t tile_count = group_tile_count(args);
+      Status status = precompute(args, tile_count, workspace);
+      if (status != Status::kSuccess) {
+        return status;
+      }
+
+      gemm_params_.update(args, workspace, tile_count);
+    } else {
+      gemm_params_.update(args, workspace);
+    }
+
+    return Status::kSuccess;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+    if (!gemm_params_.problem_visitor.problem_count) {
+      return Status::kSuccess;
+    }
+
+    //
+    // Launch kernel
+    //
+
+    // Launch splitk grouped gemm
+    {
+      dim3 grid(gemm_params_.threadblock_count, 1, gemm_params_.split_k_slices);
+      dim3 block(BaseKernel::kThreadCount, 1, 1);
+
+      int smem_size = int(sizeof(typename BaseKernel::SharedStorage));
+      cutlass::Kernel<BaseKernel>
+          <<<grid, block, smem_size, stream>>>(gemm_params_);
+
+      cudaError_t result = cudaGetLastError();
+
+      if (result != cudaSuccess) {
+        CUTLASS_TRACE_HOST("  grid launch failed with error "
+                           << cudaGetErrorString(result));
+        return Status::kErrorInternal;
+      }
+    }
+
+    // Launch splitkReduction
+    {
+      dim3 grid(32, gemm_params_.problem_visitor.problem_count);
+      dim3 block(256);
+      splitkReduction<<<grid, block, 0, stream>>>(
+          gemm_params_.ptr_D,
+          gemm_params_.ptr_D_split,
+          gemm_params_.problem_visitor.problem_sizes,
+          gemm_params_.split_k_slices,
+          gemm_params_.splitk_buffer_offsets);
+
+      cudaError_t result = cudaGetLastError();
+
+      if (result != cudaSuccess) {
+        CUTLASS_TRACE_HOST("  grid launch failed with error "
+                           << cudaGetErrorString(result));
+        return Status::kErrorInternal;
+      }
+    }
+
+    return Status::kSuccess;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
+
+  /// Initializes and runs the kernel.
+  Status operator()(Arguments const& args,
+                    void* workspace,
+                    cudaStream_t stream = nullptr) {
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 /// GEMM Grouped
 template <typename GemmKernel_>
-class SplitkGemmGrouped : public BaseSplitkGrouped<GemmKernel_>
-{
-public:
-    using GemmKernel = GemmKernel_;
+class SplitkGemmGrouped : public BaseSplitkGrouped<GemmKernel_> {
+ public:
+  using GemmKernel = GemmKernel_;
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace device
-} // namespace gemm
-} // namespace cutlass
+}  // namespace device
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/dispatch_policy.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/dispatch_policy.hpp
index f4cf0bf420..c276c4e9c0 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/dispatch_policy.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/dispatch_policy.hpp
@@ -30,7 +30,8 @@ struct KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum
 
 // n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp
 // specialized dynamic schedule For FP8 kernels with Block Scaling
-template <int Stages_, class ClusterShape_ = Shape<_1, _1, _1>,
+template <int Stages_,
+          class ClusterShape_ = Shape<_1, _1, _1>,
           class KernelSchedule = KernelTmaWarpSpecialized,
           int ScaleGranularityM =
               0  // `ScaleGranularityM` specifies scaling granularity along M,
@@ -38,7 +39,8 @@ template <int Stages_, class ClusterShape_ = Shape<_1, _1, _1>,
                  // granularity is `size<0>(TileShape_MNK{})` along M.
           >
 struct MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8
-    : MainloopSm90TmaGmmaWarpSpecialized<Stages_, ClusterShape_,
+    : MainloopSm90TmaGmmaWarpSpecialized<Stages_,
+                                         ClusterShape_,
                                          KernelSchedule> {
   static_assert(
       cute::is_same_v<
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h
index 3f83470270..0d25460e66 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,164 +27,190 @@
 #include "cutlass_extensions/arch/mma.h"
 #include "cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h"
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace kernel
-{
+namespace cutlass {
+namespace gemm {
+namespace kernel {
 
 template <typename TypeA, typename TypeB, typename arch, typename Enable = void>
-struct MixedGemmArchTraits
-{
-    static_assert(dependent_false<arch>, "Unrecognised parameterization");
+struct MixedGemmArchTraits {
+  static_assert(dependent_false<arch>, "Unrecognised parameterization");
 };
 
 template <typename Arch>
-struct MixedGemmArchTraits<float, float, Arch>
-{
-    static constexpr int Stages = 2;
-    using OperatorClass = cutlass::arch::OpClassSimt;
-    using AccType = float;
-    using LayoutB = cutlass::layout::ColumnMajor;
+struct MixedGemmArchTraits<float, float, Arch> {
+  static constexpr int Stages = 2;
+  using OperatorClass = cutlass::arch::OpClassSimt;
+  using AccType = float;
+  using LayoutB = cutlass::layout::ColumnMajor;
 
-    static constexpr int ElementsPerAccessA = 1;
-    static constexpr int ElementsPerAccessB = 1;
-    static constexpr int ElementsPerAccessC = 1;
-    static constexpr int ThreadblockK = 8;
-    using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>;
+  static constexpr int ElementsPerAccessA = 1;
+  static constexpr int ElementsPerAccessB = 1;
+  static constexpr int ElementsPerAccessC = 1;
+  static constexpr int ThreadblockK = 8;
+  using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>;
 
-    using Operator = cutlass::arch::OpMultiplyAdd;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
 // ========================= Volta Traits ===========================
 // Volta will always dequantize after the global memory load.
 // This will instantiate any HMMA tensorcore kernels for Volta.
-// Note that volta does not have native bfloat support so weights and activations will be casted to fp16
-// and compute will happen in fp16 then will be converted for bf16 output.
+// Note that volta does not have native bfloat support so weights and
+// activations will be casted to fp16 and compute will happen in fp16 then will
+// be converted for bf16 output.
 template <typename TypeA, typename TypeB>
-struct MixedGemmArchTraits<TypeA, TypeB, cutlass::arch::Sm70,
-    typename cutlass::platform::enable_if<cutlass::platform::is_same<TypeA, cutlass::half_t>::value
-        || cutlass::platform::is_same<TypeA, cutlass::bfloat16_t>::value>::type>
-{
-private:
-    using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm70>;
+struct MixedGemmArchTraits<
+    TypeA,
+    TypeB,
+    cutlass::arch::Sm70,
+    typename cutlass::platform::enable_if<
+        cutlass::platform::is_same<TypeA, cutlass::half_t>::value ||
+        cutlass::platform::is_same<TypeA, cutlass::bfloat16_t>::value>::type> {
+ private:
+  using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm70>;
 
-public:
-    static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
+ public:
+  static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
 
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using AccType = float;
-    using LayoutB = typename LayoutDetails::Layout;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using AccType = float;
+  using LayoutB = typename LayoutDetails::Layout;
 
-    static constexpr int ElementsPerAccessA = 128 / cutlass::sizeof_bits<TypeA>::value;
-    static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
-    static constexpr int ElementsPerAccessC = 128 / cutlass::sizeof_bits<TypeA>::value;
-    using InstructionShape = cutlass::gemm::GemmShape<8, 8, 4>;
+  static constexpr int ElementsPerAccessA =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
+  static constexpr int ElementsPerAccessC =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  using InstructionShape = cutlass::gemm::GemmShape<8, 8, 4>;
 
-    using Operator = typename LayoutDetails::Operator;
+  using Operator = typename LayoutDetails::Operator;
 };
 
 // ======================= Turing Traits ==============================
-// Note that turing does not have native bfloat support so weights and activations will be casted to fp16
-// and compute will happen in fp16 then will be converted for bf16 output.
+// Note that turing does not have native bfloat support so weights and
+// activations will be casted to fp16 and compute will happen in fp16 then will
+// be converted for bf16 output.
 template <typename TypeA, typename TypeB>
-struct MixedGemmArchTraits<TypeA, TypeB, cutlass::arch::Sm75,
-    typename cutlass::platform::enable_if<cutlass::platform::is_same<TypeA, cutlass::half_t>::value
-        || cutlass::platform::is_same<TypeA, cutlass::bfloat16_t>::value>::type>
-{
-private:
-    using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm75>;
+struct MixedGemmArchTraits<
+    TypeA,
+    TypeB,
+    cutlass::arch::Sm75,
+    typename cutlass::platform::enable_if<
+        cutlass::platform::is_same<TypeA, cutlass::half_t>::value ||
+        cutlass::platform::is_same<TypeA, cutlass::bfloat16_t>::value>::type> {
+ private:
+  using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm75>;
 
-public:
-    static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
+ public:
+  static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
 
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using AccType = float;
-    using LayoutB = typename LayoutDetails::Layout;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using AccType = float;
+  using LayoutB = typename LayoutDetails::Layout;
 
-    static constexpr int ElementsPerAccessA = 128 / cutlass::sizeof_bits<TypeA>::value;
-    static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
-    static constexpr int ElementsPerAccessC = 128 / cutlass::sizeof_bits<TypeA>::value;
-    using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int ElementsPerAccessA =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
+  static constexpr int ElementsPerAccessC =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>;
 
-    using Operator = typename LayoutDetails::Operator;
+  using Operator = typename LayoutDetails::Operator;
 };
 
 // ======================= Ampere Traits ==============================
 template <typename TypeA, typename TypeB>
-struct MixedGemmArchTraits<TypeA, TypeB, cutlass::arch::Sm80,
-    typename cutlass::platform::enable_if<cutlass::platform::is_same<TypeA, cutlass::half_t>::value
-        || cutlass::platform::is_same<TypeA, cutlass::bfloat16_t>::value>::type>
-{
-private:
-    using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm80>;
+struct MixedGemmArchTraits<
+    TypeA,
+    TypeB,
+    cutlass::arch::Sm80,
+    typename cutlass::platform::enable_if<
+        cutlass::platform::is_same<TypeA, cutlass::half_t>::value ||
+        cutlass::platform::is_same<TypeA, cutlass::bfloat16_t>::value>::type> {
+ private:
+  using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm80>;
 
-public:
-    static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
+ public:
+  static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
 
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using AccType = float;
-    using LayoutB = typename LayoutDetails::Layout;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using AccType = float;
+  using LayoutB = typename LayoutDetails::Layout;
 
-    static constexpr int ElementsPerAccessA = 128 / cutlass::sizeof_bits<TypeA>::value;
-    static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
-    static constexpr int ElementsPerAccessC = 128 / cutlass::sizeof_bits<TypeA>::value;
-    using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int ElementsPerAccessA =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
+  static constexpr int ElementsPerAccessC =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>;
 
-    using Operator = typename LayoutDetails::Operator;
+  using Operator = typename LayoutDetails::Operator;
 };
 
 // ======================= Ada Traits ==============================
 template <typename TypeA, typename TypeB>
-struct MixedGemmArchTraits<TypeA, TypeB, cutlass::arch::Sm89,
-    typename cutlass::platform::enable_if<cutlass::platform::is_same<TypeA, cutlass::half_t>::value
-        || cutlass::platform::is_same<TypeA, cutlass::bfloat16_t>::value>::type>
-{
-private:
-    using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm89>;
+struct MixedGemmArchTraits<
+    TypeA,
+    TypeB,
+    cutlass::arch::Sm89,
+    typename cutlass::platform::enable_if<
+        cutlass::platform::is_same<TypeA, cutlass::half_t>::value ||
+        cutlass::platform::is_same<TypeA, cutlass::bfloat16_t>::value>::type> {
+ private:
+  using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm89>;
 
-public:
-    static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
+ public:
+  static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
 
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using AccType = float;
-    using LayoutB = typename LayoutDetails::Layout;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using AccType = float;
+  using LayoutB = typename LayoutDetails::Layout;
 
-    static constexpr int ElementsPerAccessA = 128 / cutlass::sizeof_bits<TypeA>::value;
-    static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
-    static constexpr int ElementsPerAccessC = 128 / cutlass::sizeof_bits<TypeA>::value;
-    using InstructionShape = cutlass::gemm::GemmShape<16, 8, 256 / cutlass::sizeof_bits<TypeA>::value>;
+  static constexpr int ElementsPerAccessA =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
+  static constexpr int ElementsPerAccessC =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  using InstructionShape =
+      cutlass::gemm::GemmShape<16, 8, 256 / cutlass::sizeof_bits<TypeA>::value>;
 
-    using Operator = typename LayoutDetails::Operator;
+  using Operator = typename LayoutDetails::Operator;
 };
 
 // FP8 A/B = fp8, C/D = fp32
 template <typename TypeA, typename TypeB>
-struct MixedGemmArchTraits<TypeA, TypeB, cutlass::arch::Sm89,
-    typename cutlass::platform::enable_if<cutlass::platform::is_same<TypeA, cutlass::float_e4m3_t>::value
-        || cutlass::platform::is_same<TypeA, cutlass::float_e5m2_t>::value>::type>
-{
-private:
-    using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm89>;
+struct MixedGemmArchTraits<
+    TypeA,
+    TypeB,
+    cutlass::arch::Sm89,
+    typename cutlass::platform::enable_if<
+        cutlass::platform::is_same<TypeA, cutlass::float_e4m3_t>::value ||
+        cutlass::platform::is_same<TypeA,
+                                   cutlass::float_e5m2_t>::value>::type> {
+ private:
+  using LayoutDetails = LayoutDetailsB<TypeA, TypeB, cutlass::arch::Sm89>;
 
-public:
-    static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
+ public:
+  static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
 
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using AccType = float;
-    // be careful, TypeC should align with HopperGroupedGemmInput::OutputTypeAdaptor_t<TypeA>
-    using TypeC = __nv_bfloat16;
-    using LayoutB = typename LayoutDetails::Layout;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using AccType = float;
+  // be careful, TypeC should align with
+  // HopperGroupedGemmInput::OutputTypeAdaptor_t<TypeA>
+  using TypeC = __nv_bfloat16;
+  using LayoutB = typename LayoutDetails::Layout;
 
-    static constexpr int ElementsPerAccessA = 128 / cutlass::sizeof_bits<TypeA>::value;
-    static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
-    static constexpr int ElementsPerAccessC = 128 / cutlass::sizeof_bits<TypeC>::value;
-    using InstructionShape = cutlass::gemm::GemmShape<16, 8, 256 / cutlass::sizeof_bits<TypeA>::value>;
+  static constexpr int ElementsPerAccessA =
+      128 / cutlass::sizeof_bits<TypeA>::value;
+  static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
+  static constexpr int ElementsPerAccessC =
+      128 / cutlass::sizeof_bits<TypeC>::value;
+  using InstructionShape =
+      cutlass::gemm::GemmShape<16, 8, 256 / cutlass::sizeof_bits<TypeA>::value>;
 
-    using Operator = typename LayoutDetails::Operator;
+  using Operator = typename LayoutDetails::Operator;
 };
 
-} // namespace kernel
-} // namespace gemm
-} // namespace cutlass
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/default_int8_traits.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/default_int8_traits.h
index 3fd722994e..beab7fd253 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/default_int8_traits.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/default_int8_traits.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,36 +22,30 @@
 #include "cutlass/gemm/gemm.h"
 #include "cutlass/layout/matrix.h"
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace kernel
-{
+namespace cutlass {
+namespace gemm {
+namespace kernel {
 
 template <typename arch>
-struct Int8GemmArchTraits
-{
-    using OperatorClass = cutlass::arch::OpClassSimt;
-    using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>;
+struct Int8GemmArchTraits {
+  using OperatorClass = cutlass::arch::OpClassSimt;
+  using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>;
 };
 
 // ======================= Turing Traits ==============================
 template <>
-struct Int8GemmArchTraits<cutlass::arch::Sm75>
-{
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>;
+struct Int8GemmArchTraits<cutlass::arch::Sm75> {
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>;
 };
 
 // ======================= Ampere Traits ==============================
 template <>
-struct Int8GemmArchTraits<cutlass::arch::Sm80>
-{
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
+struct Int8GemmArchTraits<cutlass::arch::Sm80> {
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
 };
 
-} // namespace kernel
-} // namespace gemm
-} // namespace cutlass
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h
index b83f6d76e0..f8ee5bdab9 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,19 +18,21 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
 /*! \file
-    \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K.
+    \brief Template for a pipelined GEMM kernel. Does not compute batching or
+   support split-K.
 */
 
 #pragma once
@@ -46,523 +48,560 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace kernel
-{
+namespace cutlass {
+namespace gemm {
+namespace kernel {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace detail
-{
+namespace detail {
 template <typename>
 inline constexpr bool dependent_false_v = false;
 }
 
-template <typename Mma_,          ///! Threadblock-scoped matrix multiply-accumulate
-    typename Epilogue_,           ///! Epilogue
-    typename ThreadblockSwizzle_, ///! Threadblock swizzling function
-    typename KernelArch, ///! The Architecture this kernel is compiled for. Used since SIMT kernels lose top-level
-                         /// arch.
-    bool SplitKSerial    ///! If true, code supporting split-K via serial reduction is enabled.
-    >
-struct GemmFpAIntB
-{
+template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
+          typename Epilogue_,            ///! Epilogue
+          typename ThreadblockSwizzle_,  ///! Threadblock swizzling function
+          typename KernelArch,  ///! The Architecture this kernel is compiled
+                                /// for. Used since SIMT kernels lose top-level
+                                /// arch.
+          bool SplitKSerial     ///! If true, code supporting split-K via serial
+                                /// reduction is enabled.
+          >
+struct GemmFpAIntB {
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using EpilogueOutputOp = typename Epilogue::OutputOp;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static bool const kSplitKSerial = SplitKSerial;
 
-    using Mma = Mma_;
-    using Epilogue = Epilogue_;
-    using EpilogueOutputOp = typename Epilogue::OutputOp;
-    using ThreadblockSwizzle = ThreadblockSwizzle_;
-    static bool const kSplitKSerial = SplitKSerial;
+  using ElementA = typename Mma::IteratorA::Element;
+  using LayoutA = typename Mma::IteratorA::Layout;
+  using ElementB = typename Mma::IteratorB::Element;
+  using LayoutB = typename Mma::IteratorB::Layout;
+  using ElementC = typename Epilogue::OutputTileIterator::Element;
+  using LayoutC = typename Mma::LayoutC;
+  using ElementScale = ElementC;
 
-    using ElementA = typename Mma::IteratorA::Element;
-    using LayoutA = typename Mma::IteratorA::Layout;
-    using ElementB = typename Mma::IteratorB::Element;
-    using LayoutB = typename Mma::IteratorB::Layout;
-    using ElementC = typename Epilogue::OutputTileIterator::Element;
-    using LayoutC = typename Mma::LayoutC;
-    using ElementScale = ElementC;
+  static ComplexTransform const kTransformA = Mma::kTransformA;
+  static ComplexTransform const kTransformB = Mma::kTransformA;
 
-    static ComplexTransform const kTransformA = Mma::kTransformA;
-    static ComplexTransform const kTransformB = Mma::kTransformA;
+  // Type definitions about the mainloop.
+  using Operator = typename Mma::Operator;
+  using OperatorClass = typename Mma::Operator::OperatorClass;
+  using ThreadblockShape = typename Mma::Shape;
+  using WarpShape = typename Mma::Operator::Shape;
+  using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
+  using ArchTag = typename Mma::ArchTag;
 
-    // Type definitions about the mainloop.
-    using Operator = typename Mma::Operator;
-    using OperatorClass = typename Mma::Operator::OperatorClass;
-    using ThreadblockShape = typename Mma::Shape;
-    using WarpShape = typename Mma::Operator::Shape;
-    using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
-    using ArchTag = typename Mma::ArchTag;
+  static int const kStages = Mma::kStages;
+  static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+  static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+  static int const kAlignmentC =
+      Epilogue::OutputTileIterator::kElementsPerAccess;
 
-    static int const kStages = Mma::kStages;
-    static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
-    static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
-    static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess;
+  /// Warp count (concept: GemmShape)
+  using WarpCount = typename Mma::WarpCount;
+  static int const kThreadCount = 32 * WarpCount::kCount;
 
-    /// Warp count (concept: GemmShape)
-    using WarpCount = typename Mma::WarpCount;
-    static int const kThreadCount = 32 * WarpCount::kCount;
+  static constexpr int kInterleave =
+      Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
 
-    static constexpr int kInterleave = Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
+  /// Parameters structure
+  struct Arguments {
+    GemmUniversalMode mode = GemmUniversalMode::kGemm;
 
-    /// Parameters structure
-    struct Arguments
-    {
-        GemmUniversalMode mode = GemmUniversalMode::kGemm;
+    cutlass::gemm::GemmCoord problem_size;
+    int group_size;
+    typename Mma::IteratorA::TensorRef ref_A;
+    typename Mma::IteratorB::TensorRef ref_B;
+    typename Mma::IteratorScale::TensorRef ref_scale;
+    typename Mma::IteratorScale::TensorRef ref_zero;
+    typename Epilogue::OutputTileIterator::TensorRef ref_C;
+    typename Epilogue::OutputTileIterator::TensorRef ref_D;
 
-        cutlass::gemm::GemmCoord problem_size;
-        int group_size;
-        typename Mma::IteratorA::TensorRef ref_A;
-        typename Mma::IteratorB::TensorRef ref_B;
-        typename Mma::IteratorScale::TensorRef ref_scale;
-        typename Mma::IteratorScale::TensorRef ref_zero;
-        typename Epilogue::OutputTileIterator::TensorRef ref_C;
-        typename Epilogue::OutputTileIterator::TensorRef ref_D;
+    // Control serial split-k
+    int batch_count;
 
-        // Control serial split-k
-        int batch_count;
+    typename EpilogueOutputOp::Params output_op;
 
-        typename EpilogueOutputOp::Params output_op;
+    // For gather+scatter operations
+    int const* gather_A_indices;
+    int const* gather_B_indices;
+    int const* scatter_D_indices;
 
-        // For gather+scatter operations
-        int const* gather_A_indices;
-        int const* gather_B_indices;
-        int const* scatter_D_indices;
-
-        // Included so we can use Gemm Universal
-        int batch_stride_D = 0;
-
-        //
-        // Methods
-        //
-
-        CUTLASS_HOST_DEVICE
-        Arguments() {}
-
-        CUTLASS_HOST_DEVICE
-        Arguments(cutlass::gemm::GemmCoord const& problem_size, int const group_size,
-            typename Mma::IteratorA::TensorRef ref_A, typename Mma::IteratorB::TensorRef ref_B,
-            typename Mma::IteratorScale::TensorRef ref_scale, typename Mma::IteratorScale::TensorRef ref_zero,
-            typename Epilogue::OutputTileIterator::TensorRef ref_C,
-            typename Epilogue::OutputTileIterator::TensorRef ref_D, int serial_split_k_factor,
-            typename EpilogueOutputOp::Params output_op = typename EpilogueOutputOp::Params(),
-            int const* gather_A_indices = nullptr, int const* gather_B_indices = nullptr,
-            int const* scatter_D_indices = nullptr)
-            : problem_size(problem_size)
-            , group_size(group_size)
-            , ref_A(ref_A)
-            , ref_B(ref_B)
-            , ref_scale(ref_scale)
-            , ref_zero(ref_zero)
-            , ref_C(ref_C)
-            , ref_D(ref_D)
-            , batch_count(serial_split_k_factor)
-            , output_op(output_op)
-            , gather_A_indices(gather_A_indices)
-            , gather_B_indices(gather_B_indices)
-            , scatter_D_indices(scatter_D_indices)
-        {
-        }
-    };
-
-    /// Parameters structure
-    struct Params
-    {
-        cutlass::gemm::GemmCoord problem_size;
-        int group_size;
-        cutlass::gemm::GemmCoord grid_tiled_shape;
-        int swizzle_log_tile;
-        typename Mma::IteratorA::Params params_A;
-        typename Mma::IteratorA::TensorRef ref_A;
-        typename Mma::IteratorB::Params params_B;
-        typename Mma::IteratorB::TensorRef ref_B;
-        typename Mma::IteratorScale::Params params_scale;
-        typename Mma::IteratorScale::TensorRef ref_scale;
-        typename Mma::IteratorScale::TensorRef ref_zero;
-        typename Epilogue::OutputTileIterator::Params params_C;
-        typename Epilogue::OutputTileIterator::TensorRef ref_C;
-        typename Epilogue::OutputTileIterator::Params params_D;
-        typename Epilogue::OutputTileIterator::TensorRef ref_D;
-        typename EpilogueOutputOp::Params output_op;
-        int* semaphore;
-        int gemm_k_size;
-        // For gather+scatter operations
-        int const* gather_A_indices;
-        int const* gather_B_indices;
-        int const* scatter_D_indices;
-
-        //
-        // Methods
-        //
-
-        CUTLASS_HOST_DEVICE
-        Params()
-            : swizzle_log_tile(0)
-            , semaphore(0)
-            , gemm_k_size(0)
-        {
-        }
-
-        CUTLASS_HOST_DEVICE
-        Params(Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape, int const gemm_k_size,
-            void* workspace = nullptr)
-            : problem_size(args.problem_size)
-            , group_size(args.group_size)
-            , grid_tiled_shape(grid_tiled_shape)
-            , swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape))
-            , params_A(args.ref_A.layout())
-            , ref_A(args.ref_A)
-            , params_B(args.ref_B.layout())
-            , ref_B(args.ref_B)
-            , params_scale(args.ref_scale.layout())
-            , ref_scale(args.ref_scale)
-            , ref_zero(args.ref_zero)
-            , params_C(args.ref_C.layout())
-            , ref_C(args.ref_C)
-            , params_D(args.ref_D.layout())
-            , ref_D(args.ref_D)
-            , output_op(args.output_op)
-            , semaphore(static_cast<int*>(workspace))
-            , gemm_k_size(gemm_k_size)
-            , gather_A_indices(args.gather_A_indices)
-            , gather_B_indices(args.gather_B_indices)
-            , scatter_D_indices(args.scatter_D_indices)
-        {
-        }
-    };
-
-    /// Shared memory storage structure
-    union SharedStorage
-    {
-        typename Mma::SharedStorage main_loop;
-        typename Epilogue::SharedStorage epilogue;
-    };
+    // Included so we can use Gemm Universal
+    int batch_stride_D = 0;
 
     //
     // Methods
     //
 
     CUTLASS_HOST_DEVICE
-    GemmFpAIntB() {}
+    Arguments() {}
 
-    /// Determines whether kernel satisfies alignment
-    static Status can_implement(Arguments const& args)
-    {
-        static int const kAlignmentA
-            = (platform::is_same<typename Mma::IteratorA::Layout, layout::ColumnMajorInterleaved<32>>::value) ? 32
-            : (platform::is_same<typename Mma::IteratorA::Layout, layout::ColumnMajorInterleaved<64>>::value)
+    CUTLASS_HOST_DEVICE
+    Arguments(cutlass::gemm::GemmCoord const& problem_size,
+              int const group_size,
+              typename Mma::IteratorA::TensorRef ref_A,
+              typename Mma::IteratorB::TensorRef ref_B,
+              typename Mma::IteratorScale::TensorRef ref_scale,
+              typename Mma::IteratorScale::TensorRef ref_zero,
+              typename Epilogue::OutputTileIterator::TensorRef ref_C,
+              typename Epilogue::OutputTileIterator::TensorRef ref_D,
+              int serial_split_k_factor,
+              typename EpilogueOutputOp::Params output_op =
+                  typename EpilogueOutputOp::Params(),
+              int const* gather_A_indices = nullptr,
+              int const* gather_B_indices = nullptr,
+              int const* scatter_D_indices = nullptr)
+        : problem_size(problem_size),
+          group_size(group_size),
+          ref_A(ref_A),
+          ref_B(ref_B),
+          ref_scale(ref_scale),
+          ref_zero(ref_zero),
+          ref_C(ref_C),
+          ref_D(ref_D),
+          batch_count(serial_split_k_factor),
+          output_op(output_op),
+          gather_A_indices(gather_A_indices),
+          gather_B_indices(gather_B_indices),
+          scatter_D_indices(scatter_D_indices) {}
+  };
+
+  /// Parameters structure
+  struct Params {
+    cutlass::gemm::GemmCoord problem_size;
+    int group_size;
+    cutlass::gemm::GemmCoord grid_tiled_shape;
+    int swizzle_log_tile;
+    typename Mma::IteratorA::Params params_A;
+    typename Mma::IteratorA::TensorRef ref_A;
+    typename Mma::IteratorB::Params params_B;
+    typename Mma::IteratorB::TensorRef ref_B;
+    typename Mma::IteratorScale::Params params_scale;
+    typename Mma::IteratorScale::TensorRef ref_scale;
+    typename Mma::IteratorScale::TensorRef ref_zero;
+    typename Epilogue::OutputTileIterator::Params params_C;
+    typename Epilogue::OutputTileIterator::TensorRef ref_C;
+    typename Epilogue::OutputTileIterator::Params params_D;
+    typename Epilogue::OutputTileIterator::TensorRef ref_D;
+    typename EpilogueOutputOp::Params output_op;
+    int* semaphore;
+    int gemm_k_size;
+    // For gather+scatter operations
+    int const* gather_A_indices;
+    int const* gather_B_indices;
+    int const* scatter_D_indices;
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Params() : swizzle_log_tile(0), semaphore(0), gemm_k_size(0) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(Arguments const& args,
+           cutlass::gemm::GemmCoord const& grid_tiled_shape,
+           int const gemm_k_size,
+           void* workspace = nullptr)
+        : problem_size(args.problem_size),
+          group_size(args.group_size),
+          grid_tiled_shape(grid_tiled_shape),
+          swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)),
+          params_A(args.ref_A.layout()),
+          ref_A(args.ref_A),
+          params_B(args.ref_B.layout()),
+          ref_B(args.ref_B),
+          params_scale(args.ref_scale.layout()),
+          ref_scale(args.ref_scale),
+          ref_zero(args.ref_zero),
+          params_C(args.ref_C.layout()),
+          ref_C(args.ref_C),
+          params_D(args.ref_D.layout()),
+          ref_D(args.ref_D),
+          output_op(args.output_op),
+          semaphore(static_cast<int*>(workspace)),
+          gemm_k_size(gemm_k_size),
+          gather_A_indices(args.gather_A_indices),
+          gather_B_indices(args.gather_B_indices),
+          scatter_D_indices(args.scatter_D_indices) {}
+  };
+
+  /// Shared memory storage structure
+  union SharedStorage {
+    typename Mma::SharedStorage main_loop;
+    typename Epilogue::SharedStorage epilogue;
+  };
+
+  //
+  // Methods
+  //
+
+  CUTLASS_HOST_DEVICE
+  GemmFpAIntB() {}
+
+  /// Determines whether kernel satisfies alignment
+  static Status can_implement(Arguments const& args) {
+    static int const kAlignmentA =
+        (platform::is_same<typename Mma::IteratorA::Layout,
+                           layout::ColumnMajorInterleaved<32>>::value)
+            ? 32
+        : (platform::is_same<typename Mma::IteratorA::Layout,
+                             layout::ColumnMajorInterleaved<64>>::value)
             ? 64
             : Mma::IteratorA::AccessType::kElements;
-        static int const kAlignmentB
-            = (platform::is_same<typename Mma::IteratorB::Layout, layout::RowMajorInterleaved<32>>::value) ? 32
-            : (platform::is_same<typename Mma::IteratorB::Layout, layout::RowMajorInterleaved<64>>::value)
+    static int const kAlignmentB =
+        (platform::is_same<typename Mma::IteratorB::Layout,
+                           layout::RowMajorInterleaved<32>>::value)
+            ? 32
+        : (platform::is_same<typename Mma::IteratorB::Layout,
+                             layout::RowMajorInterleaved<64>>::value)
             ? 64
             : Mma::IteratorB::AccessType::kElements;
 
-        static int const kAlignmentScale = Mma::IteratorScale::AccessType::kElements;
+    static int const kAlignmentScale =
+        Mma::IteratorScale::AccessType::kElements;
 
-        static int const kAlignmentC = (platform::is_same<typename Epilogue::OutputTileIterator::Layout,
-                                           layout::ColumnMajorInterleaved<32>>::value)
+    static int const kAlignmentC =
+        (platform::is_same<typename Epilogue::OutputTileIterator::Layout,
+                           layout::ColumnMajorInterleaved<32>>::value)
             ? 32
-            : (platform::is_same<typename Epilogue::OutputTileIterator::Layout,
-                  layout::ColumnMajorInterleaved<64>>::value)
+        : (platform::is_same<typename Epilogue::OutputTileIterator::Layout,
+                             layout::ColumnMajorInterleaved<64>>::value)
             ? 64
             : Epilogue::OutputTileIterator::kElementsPerAccess;
 
-        if (!TensorRef_aligned(args.ref_A, kAlignmentA))
-        {
-            return Status::kErrorMisalignedOperand;
-        }
-
-        if (!TensorRef_aligned(args.ref_B, kAlignmentB))
-        {
-            return Status::kErrorMisalignedOperand;
-        }
-
-        if (!TensorRef_aligned(args.ref_scale, kAlignmentScale))
-        {
-            return Status::kErrorMisalignedOperand;
-        }
-
-        if (!TensorRef_aligned(args.ref_zero, kAlignmentScale))
-        {
-            return Status::kErrorMisalignedOperand;
-        }
-
-        if (!TensorRef_aligned(args.ref_C, kAlignmentC))
-        {
-            return Status::kErrorMisalignedOperand;
-        }
-
-        if (!TensorRef_aligned(args.ref_D, kAlignmentC))
-        {
-            return Status::kErrorMisalignedOperand;
-        }
-
-        if (!args.ref_scale.good())
-        {
-            return Status::kErrorNotSupported;
-        }
-
-        if constexpr (hasZero(Mma::QuantOp))
-        {
-            if (!args.ref_zero.good())
-            {
-                return Status::kErrorNotSupported;
-            }
-        }
-        else
-        {
-            if (args.ref_zero.good())
-            {
-                return Status::kErrorNotSupported;
-            }
-        }
-
-        if constexpr (isFinegrained(Mma::QuantOp))
-        {
-            if (args.group_size != 64 && args.group_size != 128)
-            {
-                return Status::kErrorNotSupported;
-            }
-        }
-
-        return Status::kSuccess;
+    if (!TensorRef_aligned(args.ref_A, kAlignmentA)) {
+      return Status::kErrorMisalignedOperand;
     }
 
-    static size_t get_extra_workspace_size(Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape)
-    {
-
-        return 0;
+    if (!TensorRef_aligned(args.ref_B, kAlignmentB)) {
+      return Status::kErrorMisalignedOperand;
     }
 
-    // Initializes the fine grained scale+bias iterator. Needed since the fine grained iterator
-    // has a different constructor signature than a regular cutlass iterator
-    template <typename IteratorScale, WeightOnlyQuantOp op, std::enable_if_t<isFinegrained(op), bool> = true>
-    CUTLASS_DEVICE static IteratorScale initialize_scale(typename IteratorScale::Params const& params,
-        typename IteratorScale::Pointer pointer_scale, typename IteratorScale::Pointer pointer_zero,
-        typename IteratorScale::TensorCoord extent, int thread_id,
-        typename IteratorScale::TensorCoord const& threadblock_offset, int group_size)
-    {
-
-        return IteratorScale(params, pointer_scale, pointer_zero, extent, thread_id, threadblock_offset, group_size);
+    if (!TensorRef_aligned(args.ref_scale, kAlignmentScale)) {
+      return Status::kErrorMisalignedOperand;
     }
 
-    template <typename IteratorScale, WeightOnlyQuantOp op, std::enable_if_t<!isFinegrained(op), bool> = true>
-    CUTLASS_DEVICE static IteratorScale initialize_scale(typename IteratorScale::Params const& params,
-        typename IteratorScale::Pointer pointer_scale, typename IteratorScale::Pointer pointer_zero,
-        typename IteratorScale::TensorCoord extent, int thread_id,
-        typename IteratorScale::TensorCoord const& threadblock_offset, int group_size)
-    {
-
-        return IteratorScale(params, pointer_scale, extent, thread_id, threadblock_offset);
+    if (!TensorRef_aligned(args.ref_zero, kAlignmentScale)) {
+      return Status::kErrorMisalignedOperand;
     }
 
-    CUTLASS_DEVICE
-    void run_kernel_(Params const& params, SharedStorage& shared_storage)
-    {
-        using LayoutB = typename Mma::IteratorB::Layout;
-        static_assert(platform::is_same<LayoutB, layout::RowMajor>::value && kInterleave == 1
-                || platform::is_same<LayoutB, layout::ColumnMajor>::value && kInterleave >= 1,
-            "B must be row major/col major OR col major interleaved.");
-
-        // Compute threadblock location
-        ThreadblockSwizzle threadblock_swizzle;
-
-        cutlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
-
-        // Early exit if CTA is out of range
-        if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m()
-            || params.grid_tiled_shape.n() <= threadblock_tile_offset.n())
-        {
-
-            return;
-        }
-
-        // Compute initial location in logical coordinates
-        cutlass::MatrixCoord tb_offset_A{
-            threadblock_tile_offset.m() * Mma::Shape::kM,
-            threadblock_tile_offset.k() * params.gemm_k_size,
-        };
-
-        cutlass::MatrixCoord tb_offset_B{threadblock_tile_offset.k() * params.gemm_k_size * kInterleave,
-            threadblock_tile_offset.n() * Mma::Shape::kN / kInterleave};
-
-        typename MatrixCoord::Index fg_row_offset = threadblock_tile_offset.k() * params.gemm_k_size / 64;
-        typename MatrixCoord::Index scale_row_offset = isFinegrained(Mma::QuantOp) ? fg_row_offset : 0;
-        cutlass::MatrixCoord tb_offset_scale{scale_row_offset, threadblock_tile_offset.n() * Mma::Shape::kN};
-
-        // Problem size is a function of threadblock index in the K dimension
-        int problem_size_k = min(params.problem_size.k(), (threadblock_tile_offset.k() + 1) * params.gemm_k_size);
-
-        // Compute threadblock-scoped matrix multiply-add
-        int gemm_k_iterations = (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) / Mma::Shape::kK;
-
-        // Compute position within threadblock
-        int thread_idx = threadIdx.x;
-
-        // Construct iterators to A and B operands
-        typename Mma::IteratorA iterator_A(params.params_A, params.ref_A.data(),
-            {params.problem_size.m(), problem_size_k}, thread_idx, tb_offset_A, params.gather_A_indices);
-
-        typename Mma::IteratorB iterator_B(params.params_B, params.ref_B.data(),
-            {problem_size_k * kInterleave, params.problem_size.n() / kInterleave}, thread_idx, tb_offset_B,
-            params.gather_B_indices);
-
-        typename MatrixCoord::Index scale_row_extent = isFinegrained(Mma::QuantOp) ? problem_size_k / 64 : 1;
-        typename Mma::IteratorScale iterator_scale = initialize_scale<typename Mma::IteratorScale, Mma::QuantOp>(
-            params.params_scale, params.ref_scale.data(), params.ref_zero.data(),
-            {scale_row_extent, params.problem_size.n()}, thread_idx, tb_offset_scale, params.group_size);
-
-        // Broadcast the warp_id computed by lane 0 to ensure dependent code
-        // is compiled as warp-uniform.
-        int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
-        int lane_idx = threadIdx.x % 32;
-
-        //
-        // Main loop
-        //
-        // Construct thread-scoped matrix multiply
-        Mma mma(shared_storage.main_loop, params.group_size, thread_idx, warp_idx, lane_idx);
-
-        typename Mma::FragmentC accumulators;
-
-        accumulators.clear();
-
-        if (!kSplitKSerial || gemm_k_iterations > 0)
-        {
-            // Compute threadblock-scoped matrix multiply-add
-            mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, iterator_scale, accumulators);
-        }
-
-        //
-        // Epilogue
-        //
-
-        EpilogueOutputOp output_op(params.output_op);
-
-        //
-        // Masked tile iterators constructed from members
-        //
-
-        threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
-
-        // assume identity swizzle
-        MatrixCoord threadblock_offset(
-            threadblock_tile_offset.m() * Mma::Shape::kM, threadblock_tile_offset.n() * Mma::Shape::kN);
-
-        int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m();
-
-        // Construct the semaphore.
-        Semaphore semaphore(params.semaphore + block_idx, thread_idx);
-
-        // If performing a reduction via split-K, fetch the initial synchronization
-        if (kSplitKSerial && params.grid_tiled_shape.k() > 1)
-        {
-
-            // Fetch the synchronization lock initially but do not block.
-            semaphore.fetch();
-
-            // Indicate which position in a serial reduction the output operator is currently updating
-            output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k());
-        }
-
-        // Tile iterator loading from source tensor.
-        typename Epilogue::OutputTileIterator iterator_C(params.params_C, params.ref_C.data(), params.problem_size.mn(),
-            thread_idx, threadblock_offset, params.scatter_D_indices);
-
-        // Tile iterator writing to destination tensor.
-        typename Epilogue::OutputTileIterator iterator_D(params.params_D, params.ref_D.data(), params.problem_size.mn(),
-            thread_idx, threadblock_offset, params.scatter_D_indices);
-
-        Epilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx);
-
-        // Wait on the semaphore - this latency may have been covered by iterator construction
-        if (kSplitKSerial && params.grid_tiled_shape.k() > 1)
-        {
-
-            // For subsequent threadblocks, the source matrix is held in the 'D' tensor.
-            if (threadblock_tile_offset.k())
-            {
-                iterator_C = iterator_D;
-            }
-
-            semaphore.wait(threadblock_tile_offset.k());
-        }
-
-        // Execute the epilogue operator to update the destination tensor.
-        epilogue(output_op, iterator_D, accumulators, iterator_C);
-
-        //
-        // Release the semaphore
-        //
-
-        if (kSplitKSerial && params.grid_tiled_shape.k() > 1)
-        {
-
-            int lock = 0;
-            if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1)
-            {
-
-                // The final threadblock resets the semaphore for subsequent grids.
-                lock = 0;
-            }
-            else
-            {
-                // Otherwise, the semaphore is incremented
-                lock = threadblock_tile_offset.k() + 1;
-            }
-
-            semaphore.release(lock);
-        }
+    if (!TensorRef_aligned(args.ref_C, kAlignmentC)) {
+      return Status::kErrorMisalignedOperand;
     }
 
-    template <typename CompilationArch>
-    CUTLASS_DEVICE void run_kernel(Params const& params, SharedStorage& shared_storage)
-    {
-        if constexpr (platform::is_same<KernelArch, CompilationArch>::value)
-        {
-            run_kernel_(params, shared_storage);
-        }
-        else
-        {
-            CUTLASS_NOT_IMPLEMENTED();
-        }
+    if (!TensorRef_aligned(args.ref_D, kAlignmentC)) {
+      return Status::kErrorMisalignedOperand;
     }
 
-    /*
-        To improve compilation speed, we do not compile the device operator if the CUDA_ARCH does not correspond
-        to the ArchTag of the cutlass kernel operator.
-      */
-    /// Executes one GEMM
-    CUTLASS_DEVICE
-    void operator()(Params const& params, SharedStorage& shared_storage)
-    {
+    if (!args.ref_scale.good()) {
+      return Status::kErrorNotSupported;
+    }
+
+    if constexpr (hasZero(Mma::QuantOp)) {
+      if (!args.ref_zero.good()) {
+        return Status::kErrorNotSupported;
+      }
+    } else {
+      if (args.ref_zero.good()) {
+        return Status::kErrorNotSupported;
+      }
+    }
+
+    if constexpr (isFinegrained(Mma::QuantOp)) {
+      if (args.group_size != 64 && args.group_size != 128) {
+        return Status::kErrorNotSupported;
+      }
+    }
+
+    return Status::kSuccess;
+  }
+
+  static size_t get_extra_workspace_size(
+      Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape) {
+    return 0;
+  }
+
+  // Initializes the fine grained scale+bias iterator. Needed since the fine
+  // grained iterator has a different constructor signature than a regular
+  // cutlass iterator
+  template <typename IteratorScale,
+            WeightOnlyQuantOp op,
+            std::enable_if_t<isFinegrained(op), bool> = true>
+  CUTLASS_DEVICE static IteratorScale initialize_scale(
+      typename IteratorScale::Params const& params,
+      typename IteratorScale::Pointer pointer_scale,
+      typename IteratorScale::Pointer pointer_zero,
+      typename IteratorScale::TensorCoord extent,
+      int thread_id,
+      typename IteratorScale::TensorCoord const& threadblock_offset,
+      int group_size) {
+    return IteratorScale(params,
+                         pointer_scale,
+                         pointer_zero,
+                         extent,
+                         thread_id,
+                         threadblock_offset,
+                         group_size);
+  }
+
+  template <typename IteratorScale,
+            WeightOnlyQuantOp op,
+            std::enable_if_t<!isFinegrained(op), bool> = true>
+  CUTLASS_DEVICE static IteratorScale initialize_scale(
+      typename IteratorScale::Params const& params,
+      typename IteratorScale::Pointer pointer_scale,
+      typename IteratorScale::Pointer pointer_zero,
+      typename IteratorScale::TensorCoord extent,
+      int thread_id,
+      typename IteratorScale::TensorCoord const& threadblock_offset,
+      int group_size) {
+    return IteratorScale(
+        params, pointer_scale, extent, thread_id, threadblock_offset);
+  }
+
+  CUTLASS_DEVICE
+  void run_kernel_(Params const& params, SharedStorage& shared_storage) {
+    using LayoutB = typename Mma::IteratorB::Layout;
+    static_assert(platform::is_same<LayoutB, layout::RowMajor>::value &&
+                          kInterleave == 1 ||
+                      platform::is_same<LayoutB, layout::ColumnMajor>::value &&
+                          kInterleave >= 1,
+                  "B must be row major/col major OR col major interleaved.");
+
+    // Compute threadblock location
+    ThreadblockSwizzle threadblock_swizzle;
+
+    cutlass::gemm::GemmCoord threadblock_tile_offset =
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+
+    // Early exit if CTA is out of range
+    if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
+        params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
+      return;
+    }
+
+    // Compute initial location in logical coordinates
+    cutlass::MatrixCoord tb_offset_A{
+        threadblock_tile_offset.m() * Mma::Shape::kM,
+        threadblock_tile_offset.k() * params.gemm_k_size,
+    };
+
+    cutlass::MatrixCoord tb_offset_B{
+        threadblock_tile_offset.k() * params.gemm_k_size * kInterleave,
+        threadblock_tile_offset.n() * Mma::Shape::kN / kInterleave};
+
+    typename MatrixCoord::Index fg_row_offset =
+        threadblock_tile_offset.k() * params.gemm_k_size / 64;
+    typename MatrixCoord::Index scale_row_offset =
+        isFinegrained(Mma::QuantOp) ? fg_row_offset : 0;
+    cutlass::MatrixCoord tb_offset_scale{
+        scale_row_offset, threadblock_tile_offset.n() * Mma::Shape::kN};
+
+    // Problem size is a function of threadblock index in the K dimension
+    int problem_size_k =
+        min(params.problem_size.k(),
+            (threadblock_tile_offset.k() + 1) * params.gemm_k_size);
+
+    // Compute threadblock-scoped matrix multiply-add
+    int gemm_k_iterations =
+        (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) /
+        Mma::Shape::kK;
+
+    // Compute position within threadblock
+    int thread_idx = threadIdx.x;
+
+    // Construct iterators to A and B operands
+    typename Mma::IteratorA iterator_A(
+        params.params_A,
+        params.ref_A.data(),
+        {params.problem_size.m(), problem_size_k},
+        thread_idx,
+        tb_offset_A,
+        params.gather_A_indices);
+
+    typename Mma::IteratorB iterator_B(
+        params.params_B,
+        params.ref_B.data(),
+        {problem_size_k * kInterleave, params.problem_size.n() / kInterleave},
+        thread_idx,
+        tb_offset_B,
+        params.gather_B_indices);
+
+    typename MatrixCoord::Index scale_row_extent =
+        isFinegrained(Mma::QuantOp) ? problem_size_k / 64 : 1;
+    typename Mma::IteratorScale iterator_scale =
+        initialize_scale<typename Mma::IteratorScale, Mma::QuantOp>(
+            params.params_scale,
+            params.ref_scale.data(),
+            params.ref_zero.data(),
+            {scale_row_extent, params.problem_size.n()},
+            thread_idx,
+            tb_offset_scale,
+            params.group_size);
+
+    // Broadcast the warp_id computed by lane 0 to ensure dependent code
+    // is compiled as warp-uniform.
+    int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+    int lane_idx = threadIdx.x % 32;
+
+    //
+    // Main loop
+    //
+    // Construct thread-scoped matrix multiply
+    Mma mma(shared_storage.main_loop,
+            params.group_size,
+            thread_idx,
+            warp_idx,
+            lane_idx);
+
+    typename Mma::FragmentC accumulators;
+
+    accumulators.clear();
+
+    if (!kSplitKSerial || gemm_k_iterations > 0) {
+      // Compute threadblock-scoped matrix multiply-add
+      mma(gemm_k_iterations,
+          accumulators,
+          iterator_A,
+          iterator_B,
+          iterator_scale,
+          accumulators);
+    }
+
+    //
+    // Epilogue
+    //
+
+    EpilogueOutputOp output_op(params.output_op);
+
+    //
+    // Masked tile iterators constructed from members
+    //
+
+    threadblock_tile_offset =
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+
+    // assume identity swizzle
+    MatrixCoord threadblock_offset(
+        threadblock_tile_offset.m() * Mma::Shape::kM,
+        threadblock_tile_offset.n() * Mma::Shape::kN);
+
+    int block_idx = threadblock_tile_offset.m() +
+                    threadblock_tile_offset.n() * params.grid_tiled_shape.m();
+
+    // Construct the semaphore.
+    Semaphore semaphore(params.semaphore + block_idx, thread_idx);
+
+    // If performing a reduction via split-K, fetch the initial synchronization
+    if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+      // Fetch the synchronization lock initially but do not block.
+      semaphore.fetch();
+
+      // Indicate which position in a serial reduction the output operator is
+      // currently updating
+      output_op.set_k_partition(threadblock_tile_offset.k(),
+                                params.grid_tiled_shape.k());
+    }
+
+    // Tile iterator loading from source tensor.
+    typename Epilogue::OutputTileIterator iterator_C(params.params_C,
+                                                     params.ref_C.data(),
+                                                     params.problem_size.mn(),
+                                                     thread_idx,
+                                                     threadblock_offset,
+                                                     params.scatter_D_indices);
+
+    // Tile iterator writing to destination tensor.
+    typename Epilogue::OutputTileIterator iterator_D(params.params_D,
+                                                     params.ref_D.data(),
+                                                     params.problem_size.mn(),
+                                                     thread_idx,
+                                                     threadblock_offset,
+                                                     params.scatter_D_indices);
+
+    Epilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx);
+
+    // Wait on the semaphore - this latency may have been covered by iterator
+    // construction
+    if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+      // For subsequent threadblocks, the source matrix is held in the 'D'
+      // tensor.
+      if (threadblock_tile_offset.k()) {
+        iterator_C = iterator_D;
+      }
+
+      semaphore.wait(threadblock_tile_offset.k());
+    }
+
+    // Execute the epilogue operator to update the destination tensor.
+    epilogue(output_op, iterator_D, accumulators, iterator_C);
+
+    //
+    // Release the semaphore
+    //
+
+    if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+      int lock = 0;
+      if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) {
+        // The final threadblock resets the semaphore for subsequent grids.
+        lock = 0;
+      } else {
+        // Otherwise, the semaphore is incremented
+        lock = threadblock_tile_offset.k() + 1;
+      }
+
+      semaphore.release(lock);
+    }
+  }
+
+  template <typename CompilationArch>
+  CUTLASS_DEVICE void run_kernel(Params const& params,
+                                 SharedStorage& shared_storage) {
+    if constexpr (platform::is_same<KernelArch, CompilationArch>::value) {
+      run_kernel_(params, shared_storage);
+    } else {
+      CUTLASS_NOT_IMPLEMENTED();
+    }
+  }
+
+  /*
+      To improve compilation speed, we do not compile the device operator if the
+     CUDA_ARCH does not correspond to the ArchTag of the cutlass kernel
+     operator.
+    */
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const& params, SharedStorage& shared_storage) {
 #if defined(__CUDA_ARCH__)
 #if (__CUDA_ARCH__ >= 700) && (__CUDA_ARCH__ < 750)
-        run_kernel<arch::Sm70>(params, shared_storage);
+    run_kernel<arch::Sm70>(params, shared_storage);
 #elif (__CUDA_ARCH__ >= 750) && (__CUDA_ARCH__ < 800)
-        run_kernel<arch::Sm75>(params, shared_storage);
+    run_kernel<arch::Sm75>(params, shared_storage);
 #elif (__CUDA_ARCH__ >= 800) && (__CUDA_ARCH__ < 890)
-        run_kernel<arch::Sm80>(params, shared_storage);
+    run_kernel<arch::Sm80>(params, shared_storage);
 #elif (__CUDA_ARCH__ == 890)
-        run_kernel<arch::Sm89>(params, shared_storage);
+    run_kernel<arch::Sm89>(params, shared_storage);
 #elif (__CUDA_ARCH__ >= 900)
-        CUTLASS_NOT_IMPLEMENTED(); // Don't compile these for Hopper or later. Use CUTLASS 3.x kernels.
+    CUTLASS_NOT_IMPLEMENTED();  // Don't compile these for Hopper or later. Use
+                                // CUTLASS 3.x kernels.
 #else
-        static_assert(
-            false, "Invalid architecture being compiled. Only Volta+ supported in weight-only quantization kernels.");
+    static_assert(false,
+                  "Invalid architecture being compiled. Only Volta+ supported "
+                  "in weight-only quantization kernels.");
 #endif
 #else
-        CUTLASS_NOT_IMPLEMENTED();
+    CUTLASS_NOT_IMPLEMENTED();
 #endif
-    }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace kernel
-} // namespace gemm
-} // namespace cutlass
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/gemm_universal_gated.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/gemm_universal_gated.hpp
index 15faad26ee..cc143751bc 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/gemm_universal_gated.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/gemm_universal_gated.hpp
@@ -54,15 +54,16 @@ namespace cutlass::gemm::kernel {
  * 2.x API type argument order. Template arguments without two names
  * belong to the 3.x API only.
  **/
-template <class ProblemShapeOrThreadblockMma_, // (m, n, k) or (m, n, k, l)
+template <class ProblemShapeOrThreadblockMma_,  // (m, n, k) or (m, n, k, l)
           class CollectiveMainloopOrEpilogue_,
           class CollectiveEpilogueOrThreadblockSwizzle_,
-          class TileScheduler_ = void, class Enable = void>
+          class TileScheduler_ = void,
+          class Enable = void>
 class GemmUniversalGated;
 
 ////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm::kernel
+}  // namespace cutlass::gemm::kernel
 
 ////////////////////////////////////////////////////////////////////////////////
 
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor.h
index 6c4c578c9c..be80a3feec 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,25 +18,27 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
     \brief GEMM kernel to support the epilogue visitor model
     for customized softmax partial reduction epilogue fusion.
 
-    This source file will likely be moved to `include/cutlass/gemm/kernel/` in the future once
-    its usage has been stabilized. For now, it is included in this example to demonstrate
-    some basic output fusion options.
+    This source file will likely be moved to `include/cutlass/gemm/kernel/` in
+   the future once its usage has been stabilized. For now, it is included in
+   this example to demonstrate some basic output fusion options.
 
-    original file: 3rdparty/cutlass/examples/35_gemm_softmax/gemm_with_epilogue_visitor.h
+    original file:
+   3rdparty/cutlass/examples/35_gemm_softmax/gemm_with_epilogue_visitor.h
 */
 
 #pragma once
@@ -55,533 +57,527 @@ namespace tk = common;
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace kernel
-{
+namespace cutlass {
+namespace gemm {
+namespace kernel {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-template <typename Mma_,         ///! Threadblock-scoped matrix multiply-accumulate
-    typename Epilogue_,          ///! Epilogue
-    typename ThreadblockSwizzle_ ///! Threadblock swizzling function
-    >
-struct GemmWithEpilogueVisitor
-{
-public:
-    using Mma = Mma_;
-    using Epilogue = Epilogue_;
-    using EpilogueVisitor = typename Epilogue::Visitor;
-    using ThreadblockSwizzle = ThreadblockSwizzle_;
+template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
+          typename Epilogue_,           ///! Epilogue
+          typename ThreadblockSwizzle_  ///! Threadblock swizzling function
+          >
+struct GemmWithEpilogueVisitor {
+ public:
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using EpilogueVisitor = typename Epilogue::Visitor;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
 
-    using ElementA = typename Mma::IteratorA::Element;
-    using LayoutA = typename Mma::IteratorA::Layout;
-    using TensorRefA = TensorRef<ElementA, LayoutA>;
+  using ElementA = typename Mma::IteratorA::Element;
+  using LayoutA = typename Mma::IteratorA::Layout;
+  using TensorRefA = TensorRef<ElementA, LayoutA>;
 
-    using ElementB = typename Mma::IteratorB::Element;
-    using LayoutB = typename Mma::IteratorB::Layout;
-    using TensorRefB = TensorRef<ElementB, LayoutB>;
+  using ElementB = typename Mma::IteratorB::Element;
+  using LayoutB = typename Mma::IteratorB::Layout;
+  using TensorRefB = TensorRef<ElementB, LayoutB>;
 
-    using ElementCompute = typename EpilogueVisitor::ElementCompute;
-    using LayoutAlphaCol = cutlass::layout::RowMajor;
-    using LayoutAlphaRow = cutlass::layout::ColumnMajor;
-    using TensorRefAlphaCol = TensorRef<ElementCompute, LayoutAlphaCol>;
-    using TensorRefAlphaRow = TensorRef<ElementCompute, LayoutAlphaRow>;
+  using ElementCompute = typename EpilogueVisitor::ElementCompute;
+  using LayoutAlphaCol = cutlass::layout::RowMajor;
+  using LayoutAlphaRow = cutlass::layout::ColumnMajor;
+  using TensorRefAlphaCol = TensorRef<ElementCompute, LayoutAlphaCol>;
+  using TensorRefAlphaRow = TensorRef<ElementCompute, LayoutAlphaRow>;
 
-    using ElementC = typename EpilogueVisitor::ElementOutput;
-    using LayoutC = typename Epilogue::Layout;
-    using TensorRefC = TensorRef<ElementC, LayoutC>;
+  using ElementC = typename EpilogueVisitor::ElementOutput;
+  using LayoutC = typename Epilogue::Layout;
+  using TensorRefC = TensorRef<ElementC, LayoutC>;
 
-    static ComplexTransform const kTransformA = Mma::kTransformA;
-    static ComplexTransform const kTransformB = Mma::kTransformB;
-    using Operator = typename Mma::Operator;
+  static ComplexTransform const kTransformA = Mma::kTransformA;
+  static ComplexTransform const kTransformB = Mma::kTransformB;
+  using Operator = typename Mma::Operator;
 
-    using OperatorClass = typename Mma::Operator::OperatorClass;
-    using ThreadblockShape = typename Mma::Shape;
-    using WarpShape = typename Mma::Operator::Shape;
-    using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
-    using ArchTag = typename Mma::ArchTag;
-    using EpilogueOutputOp =
-        typename Epilogue::Visitor::ElementwiseFunctor; // Define type so GemmUniversalBase doesn't complain
+  using OperatorClass = typename Mma::Operator::OperatorClass;
+  using ThreadblockShape = typename Mma::Shape;
+  using WarpShape = typename Mma::Operator::Shape;
+  using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
+  using ArchTag = typename Mma::ArchTag;
+  using EpilogueOutputOp = typename Epilogue::Visitor::
+      ElementwiseFunctor;  // Define type so GemmUniversalBase doesn't complain
 
-    static int const kStages = Mma::kStages;
-    static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
-    static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
-    static int const kAlignmentC = EpilogueVisitor::kElementsPerAccess;
+  static int const kStages = Mma::kStages;
+  static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+  static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+  static int const kAlignmentC = EpilogueVisitor::kElementsPerAccess;
 
-    /// Warp count (concept: GemmShape)
-    using WarpCount = typename Mma::WarpCount;
-    static int const kThreadCount = 32 * WarpCount::kCount;
+  /// Warp count (concept: GemmShape)
+  using WarpCount = typename Mma::WarpCount;
+  static int const kThreadCount = 32 * WarpCount::kCount;
 
-    /// Split-K preserves splits that are 128b aligned
-    static int const kSplitKAlignment
-        = const_max(128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value);
+  /// Split-K preserves splits that are 128b aligned
+  static int const kSplitKAlignment = const_max(
+      128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value);
 
+  //
+  // Structures
+  //
+
+  /// Argument structure
+  struct Arguments {
     //
-    // Structures
+    // Data members
     //
 
-    /// Argument structure
-    struct Arguments
-    {
+    GemmUniversalMode mode;
+    GemmCoord problem_size;
+    int batch_count;
 
-        //
-        // Data members
-        //
+    TensorRefA ref_A;
+    TensorRefB ref_B;
+    tk::QuantMode quant_option;
+    TensorRefAlphaCol ref_alpha_col;
+    TensorRefAlphaRow ref_alpha_row;
+    TensorRefC ref_C;
+    TensorRefC ref_D;
 
-        GemmUniversalMode mode;
-        GemmCoord problem_size;
-        int batch_count;
+    int64_t batch_stride_A;
+    int64_t batch_stride_B;
+    int64_t batch_stride_D;
 
-        TensorRefA ref_A;
-        TensorRefB ref_B;
-        tk::QuantMode quant_option;
-        TensorRefAlphaCol ref_alpha_col;
-        TensorRefAlphaRow ref_alpha_row;
-        TensorRefC ref_C;
-        TensorRefC ref_D;
+    typename EpilogueVisitor::Arguments epilogue_visitor;
 
-        int64_t batch_stride_A;
-        int64_t batch_stride_B;
-        int64_t batch_stride_D;
-
-        typename EpilogueVisitor::Arguments epilogue_visitor;
-
-        //
-        // Methods
-        //
-
-        Arguments()
-            : mode(GemmUniversalMode::kGemm)
-            , batch_count(1)
-        {
-        }
-
-        /// constructs an arguments structure
-        Arguments(GemmUniversalMode mode_, GemmCoord problem_size_, int batch_count_, TensorRefA ref_A_,
-            TensorRefB ref_B_, tk::QuantMode quant_option_, TensorRefAlphaCol ref_alpha_col_,
-            TensorRefAlphaRow ref_alpha_row_, TensorRefC ref_C_, TensorRefC ref_D_, int64_t batch_stride_A_,
-            int64_t batch_stride_B_, typename EpilogueVisitor::Arguments epilogue_visitor_)
-            : mode(mode_)
-            , problem_size(problem_size_)
-            , batch_count(batch_count_)
-            , ref_A(ref_A_)
-            , ref_B(ref_B_)
-            , quant_option(quant_option_)
-            , ref_alpha_col(ref_alpha_col_)
-            , ref_alpha_row(ref_alpha_row_)
-            , ref_C(ref_C_)
-            , ref_D(ref_D_)
-            , batch_stride_A(batch_stride_A_)
-            , batch_stride_B(batch_stride_B_)
-            , batch_stride_D(0)
-            , epilogue_visitor(epilogue_visitor_)
-        {
-        }
-    };
-
-    //
-    // Structure for precomputing values in host memory and passing to kernels
-    //
-
-    /// Parameters structure
-    struct Params
-    {
-
-        cutlass::gemm::GemmCoord problem_size;
-        cutlass::gemm::GemmCoord grid_tiled_shape;
-        int swizzle_log_tile;
-
-        typename Mma::IteratorA::Params params_A;
-        typename Mma::IteratorB::Params params_B;
-        typename EpilogueVisitor::ScaleTileIterator::Params params_alpha_col;
-        typename EpilogueVisitor::ScaleTileIterator::Params params_alpha_row;
-        typename EpilogueVisitor::OutputTileIterator::Params params_C;
-        typename EpilogueVisitor::OutputTileIterator::Params params_D;
-
-        GemmUniversalMode mode;
-        int batch_count;
-        int gemm_k_size;
-
-        void* ptr_A;
-        void* ptr_B;
-        tk::QuantMode quant_option;
-        typename EpilogueVisitor::ScaleTileIterator::Element* ptr_alpha_col;
-        typename EpilogueVisitor::ScaleTileIterator::Element* ptr_alpha_row;
-        ElementC* ptr_C;
-        ElementC* ptr_D;
-
-        int64_t batch_stride_A;
-        int64_t batch_stride_B;
-
-        typename EpilogueVisitor::Params epilogue_visitor;
-
-        //
-        // Methods
-        //
-
-        CUTLASS_HOST_DEVICE
-        Params()
-            : swizzle_log_tile(0)
-            , params_A(0)
-            , params_B(0)
-            , params_alpha_col(0)
-            , params_C(0)
-            , params_D(0)
-            , batch_count(0)
-            , gemm_k_size(0)
-            , mode(cutlass::gemm::GemmUniversalMode::kGemm)
-            , ptr_A(nullptr)
-            , ptr_B(nullptr)
-            , ptr_alpha_col(nullptr)
-            , ptr_alpha_row(nullptr)
-            , ptr_C(nullptr)
-            , ptr_D(nullptr)
-            , batch_stride_A(0)
-            , batch_stride_B(0)
-        {
-        }
-
-        Params(
-            Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape_, int gemm_k_size_, int* workspace_)
-            : problem_size(args.problem_size)
-            , swizzle_log_tile(0)
-            , params_A(args.ref_A.layout())
-            , params_B(args.ref_B.layout())
-            , params_alpha_col(args.ref_alpha_col.layout())
-            , params_alpha_row(args.ref_alpha_col.layout())
-            , params_C(args.ref_C.layout())
-            , params_D(args.ref_D.layout())
-            , mode(args.mode)
-            , batch_count(args.batch_count)
-            , gemm_k_size(args.problem_size.k())
-            , ptr_A(args.ref_A.data())
-            , ptr_B(args.ref_B.data())
-            , quant_option(args.quant_option)
-            , ptr_alpha_col(args.ref_alpha_col.data())
-            , ptr_alpha_row(args.ref_alpha_row.data())
-            , ptr_C(args.ref_C.data())
-            , ptr_D(args.ref_D.data())
-            , batch_stride_A(args.batch_stride_A)
-            , batch_stride_B(args.batch_stride_B)
-            , epilogue_visitor(args.epilogue_visitor)
-        {
-
-            ThreadblockSwizzle threadblock_swizzle;
-
-            grid_tiled_shape = threadblock_swizzle.get_tiled_shape(args.problem_size,
-                {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, args.batch_count);
-
-            if (args.mode == GemmUniversalMode::kGemm || args.mode == GemmUniversalMode::kGemmSplitKParallel)
-            {
-
-                int const kAlignK
-                    = const_max(const_max(128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value), 1);
-
-                gemm_k_size = round_up(ceil_div(args.problem_size.k(), args.batch_count), kAlignK);
-
-                if (gemm_k_size)
-                {
-                    grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size);
-                }
-            }
-
-            swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape);
-        }
-    };
-
-    /// Shared memory storage structure
-    union SharedStorage
-    {
-
-        typename Mma::SharedStorage main_loop;
-
-        struct
-        {
-            typename Epilogue::SharedStorage epilogue;
-            typename EpilogueVisitor::SharedStorage visitor;
-        } epilogue;
-    };
-
-public:
     //
     // Methods
     //
 
-    CUTLASS_DEVICE
-    GemmWithEpilogueVisitor() {}
+    Arguments() : mode(GemmUniversalMode::kGemm), batch_count(1) {}
 
-    /// Determines whether kernel satisfies alignment
-    static Status can_implement(cutlass::gemm::GemmCoord const& problem_size)
-    {
+    /// constructs an arguments structure
+    Arguments(GemmUniversalMode mode_,
+              GemmCoord problem_size_,
+              int batch_count_,
+              TensorRefA ref_A_,
+              TensorRefB ref_B_,
+              tk::QuantMode quant_option_,
+              TensorRefAlphaCol ref_alpha_col_,
+              TensorRefAlphaRow ref_alpha_row_,
+              TensorRefC ref_C_,
+              TensorRefC ref_D_,
+              int64_t batch_stride_A_,
+              int64_t batch_stride_B_,
+              typename EpilogueVisitor::Arguments epilogue_visitor_)
+        : mode(mode_),
+          problem_size(problem_size_),
+          batch_count(batch_count_),
+          ref_A(ref_A_),
+          ref_B(ref_B_),
+          quant_option(quant_option_),
+          ref_alpha_col(ref_alpha_col_),
+          ref_alpha_row(ref_alpha_row_),
+          ref_C(ref_C_),
+          ref_D(ref_D_),
+          batch_stride_A(batch_stride_A_),
+          batch_stride_B(batch_stride_B_),
+          batch_stride_D(0),
+          epilogue_visitor(epilogue_visitor_) {}
+  };
 
-        CUTLASS_TRACE_HOST("GemmWithEpilogueVisitor::can_implement()");
+  //
+  // Structure for precomputing values in host memory and passing to kernels
+  //
 
-        static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
-        static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
-        static int const kAlignmentC = EpilogueVisitor::OutputTileIterator::kElementsPerAccess;
+  /// Parameters structure
+  struct Params {
+    cutlass::gemm::GemmCoord problem_size;
+    cutlass::gemm::GemmCoord grid_tiled_shape;
+    int swizzle_log_tile;
 
-        bool isAMisaligned = false;
-        bool isBMisaligned = false;
-        bool isCMisaligned = false;
+    typename Mma::IteratorA::Params params_A;
+    typename Mma::IteratorB::Params params_B;
+    typename EpilogueVisitor::ScaleTileIterator::Params params_alpha_col;
+    typename EpilogueVisitor::ScaleTileIterator::Params params_alpha_row;
+    typename EpilogueVisitor::OutputTileIterator::Params params_C;
+    typename EpilogueVisitor::OutputTileIterator::Params params_D;
 
-        if (platform::is_same<LayoutA, layout::RowMajor>::value)
-        {
-            isAMisaligned = problem_size.k() % kAlignmentA;
-        }
-        else if (platform::is_same<LayoutA, layout::ColumnMajor>::value)
-        {
-            isAMisaligned = problem_size.m() % kAlignmentA;
-        }
-        else if (platform::is_same<LayoutA, layout::ColumnMajorInterleaved<32>>::value
-            || platform::is_same<LayoutA, layout::ColumnMajorInterleaved<64>>::value)
-        {
-            isAMisaligned = problem_size.k() % kAlignmentA;
-        }
+    GemmUniversalMode mode;
+    int batch_count;
+    int gemm_k_size;
 
-        if (platform::is_same<LayoutB, layout::RowMajor>::value)
-        {
-            isBMisaligned = problem_size.n() % kAlignmentB;
-        }
-        else if (platform::is_same<LayoutB, layout::ColumnMajor>::value)
-        {
-            isBMisaligned = problem_size.k() % kAlignmentB;
-        }
-        else if (platform::is_same<LayoutB, layout::RowMajorInterleaved<32>>::value
-            || platform::is_same<LayoutB, layout::RowMajorInterleaved<64>>::value)
-        {
-            isBMisaligned = problem_size.k() % kAlignmentB;
-        }
+    void* ptr_A;
+    void* ptr_B;
+    tk::QuantMode quant_option;
+    typename EpilogueVisitor::ScaleTileIterator::Element* ptr_alpha_col;
+    typename EpilogueVisitor::ScaleTileIterator::Element* ptr_alpha_row;
+    ElementC* ptr_C;
+    ElementC* ptr_D;
 
-        if (platform::is_same<LayoutC, layout::RowMajor>::value)
-        {
-            isCMisaligned = problem_size.n() % kAlignmentC;
-        }
-        else if (platform::is_same<LayoutC, layout::ColumnMajor>::value)
-        {
-            isCMisaligned = problem_size.m() % kAlignmentC;
-        }
-        else if (platform::is_same<LayoutC, layout::ColumnMajorInterleaved<32>>::value
-            || platform::is_same<LayoutC, layout::ColumnMajorInterleaved<64>>::value)
-        {
-            isCMisaligned = problem_size.n() % kAlignmentC;
-        }
+    int64_t batch_stride_A;
+    int64_t batch_stride_B;
 
-        if (isAMisaligned)
-        {
-            CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for A operand");
-            return Status::kErrorMisalignedOperand;
+    typename EpilogueVisitor::Params epilogue_visitor;
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : swizzle_log_tile(0),
+          params_A(0),
+          params_B(0),
+          params_alpha_col(0),
+          params_C(0),
+          params_D(0),
+          batch_count(0),
+          gemm_k_size(0),
+          mode(cutlass::gemm::GemmUniversalMode::kGemm),
+          ptr_A(nullptr),
+          ptr_B(nullptr),
+          ptr_alpha_col(nullptr),
+          ptr_alpha_row(nullptr),
+          ptr_C(nullptr),
+          ptr_D(nullptr),
+          batch_stride_A(0),
+          batch_stride_B(0) {}
+
+    Params(Arguments const& args,
+           cutlass::gemm::GemmCoord const& grid_tiled_shape_,
+           int gemm_k_size_,
+           int* workspace_)
+        : problem_size(args.problem_size),
+          swizzle_log_tile(0),
+          params_A(args.ref_A.layout()),
+          params_B(args.ref_B.layout()),
+          params_alpha_col(args.ref_alpha_col.layout()),
+          params_alpha_row(args.ref_alpha_col.layout()),
+          params_C(args.ref_C.layout()),
+          params_D(args.ref_D.layout()),
+          mode(args.mode),
+          batch_count(args.batch_count),
+          gemm_k_size(args.problem_size.k()),
+          ptr_A(args.ref_A.data()),
+          ptr_B(args.ref_B.data()),
+          quant_option(args.quant_option),
+          ptr_alpha_col(args.ref_alpha_col.data()),
+          ptr_alpha_row(args.ref_alpha_row.data()),
+          ptr_C(args.ref_C.data()),
+          ptr_D(args.ref_D.data()),
+          batch_stride_A(args.batch_stride_A),
+          batch_stride_B(args.batch_stride_B),
+          epilogue_visitor(args.epilogue_visitor) {
+      ThreadblockSwizzle threadblock_swizzle;
+
+      grid_tiled_shape = threadblock_swizzle.get_tiled_shape(
+          args.problem_size,
+          {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+          args.batch_count);
+
+      if (args.mode == GemmUniversalMode::kGemm ||
+          args.mode == GemmUniversalMode::kGemmSplitKParallel) {
+        int const kAlignK =
+            const_max(const_max(128 / sizeof_bits<ElementA>::value,
+                                128 / sizeof_bits<ElementB>::value),
+                      1);
+
+        gemm_k_size = round_up(
+            ceil_div(args.problem_size.k(), args.batch_count), kAlignK);
+
+        if (gemm_k_size) {
+          grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size);
         }
+      }
 
-        if (isBMisaligned)
-        {
-            CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for B operand");
-            return Status::kErrorMisalignedOperand;
-        }
+      swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape);
+    }
+  };
 
-        if (isCMisaligned)
-        {
-            CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for C operand");
-            return Status::kErrorMisalignedOperand;
-        }
+  /// Shared memory storage structure
+  union SharedStorage {
+    typename Mma::SharedStorage main_loop;
 
-        CUTLASS_TRACE_HOST("  returning kSuccess");
+    struct {
+      typename Epilogue::SharedStorage epilogue;
+      typename EpilogueVisitor::SharedStorage visitor;
+    } epilogue;
+  };
 
-        return Status::kSuccess;
+ public:
+  //
+  // Methods
+  //
+
+  CUTLASS_DEVICE
+  GemmWithEpilogueVisitor() {}
+
+  /// Determines whether kernel satisfies alignment
+  static Status can_implement(cutlass::gemm::GemmCoord const& problem_size) {
+    CUTLASS_TRACE_HOST("GemmWithEpilogueVisitor::can_implement()");
+
+    static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+    static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+    static int const kAlignmentC =
+        EpilogueVisitor::OutputTileIterator::kElementsPerAccess;
+
+    bool isAMisaligned = false;
+    bool isBMisaligned = false;
+    bool isCMisaligned = false;
+
+    if (platform::is_same<LayoutA, layout::RowMajor>::value) {
+      isAMisaligned = problem_size.k() % kAlignmentA;
+    } else if (platform::is_same<LayoutA, layout::ColumnMajor>::value) {
+      isAMisaligned = problem_size.m() % kAlignmentA;
+    } else if (platform::is_same<LayoutA,
+                                 layout::ColumnMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutA,
+                                 layout::ColumnMajorInterleaved<64>>::value) {
+      isAMisaligned = problem_size.k() % kAlignmentA;
     }
 
-    static Status can_implement(Arguments const& args)
-    {
-        return can_implement(args.problem_size);
+    if (platform::is_same<LayoutB, layout::RowMajor>::value) {
+      isBMisaligned = problem_size.n() % kAlignmentB;
+    } else if (platform::is_same<LayoutB, layout::ColumnMajor>::value) {
+      isBMisaligned = problem_size.k() % kAlignmentB;
+    } else if (platform::is_same<LayoutB,
+                                 layout::RowMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutB,
+                                 layout::RowMajorInterleaved<64>>::value) {
+      isBMisaligned = problem_size.k() % kAlignmentB;
     }
 
-    static size_t get_extra_workspace_size(Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape)
-    {
-
-        return 0;
+    if (platform::is_same<LayoutC, layout::RowMajor>::value) {
+      isCMisaligned = problem_size.n() % kAlignmentC;
+    } else if (platform::is_same<LayoutC, layout::ColumnMajor>::value) {
+      isCMisaligned = problem_size.m() % kAlignmentC;
+    } else if (platform::is_same<LayoutC,
+                                 layout::ColumnMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutC,
+                                 layout::ColumnMajorInterleaved<64>>::value) {
+      isCMisaligned = problem_size.n() % kAlignmentC;
     }
 
+    if (isAMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for A operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (isBMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for B operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (isCMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for C operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    CUTLASS_TRACE_HOST("  returning kSuccess");
+
+    return Status::kSuccess;
+  }
+
+  static Status can_implement(Arguments const& args) {
+    return can_implement(args.problem_size);
+  }
+
+  static size_t get_extra_workspace_size(
+      Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape) {
+    return 0;
+  }
+
 #define SPLIT_K_ENABLED 1
 
-    /// Executes one GEMM
-    CUTLASS_DEVICE
-    void run_kernel_(Params const& params, SharedStorage& shared_storage)
-    {
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void run_kernel_(Params const& params, SharedStorage& shared_storage) {
+    // Compute threadblock location
+    ThreadblockSwizzle threadblock_swizzle;
 
-        // Compute threadblock location
-        ThreadblockSwizzle threadblock_swizzle;
+    cutlass::gemm::GemmCoord threadblock_tile_offset =
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-        cutlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+    // Early exit if CTA is out of range
+    if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
+        params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
+      return;
+    }
 
-        // Early exit if CTA is out of range
-        if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m()
-            || params.grid_tiled_shape.n() <= threadblock_tile_offset.n())
-        {
+    int offset_k = 0;
+    int problem_size_k = params.problem_size.k();
 
-            return;
-        }
-
-        int offset_k = 0;
-        int problem_size_k = params.problem_size.k();
-
-        ElementA* ptr_A = static_cast<ElementA*>(params.ptr_A);
-        ElementB* ptr_B = static_cast<ElementB*>(params.ptr_B);
+    ElementA* ptr_A = static_cast<ElementA*>(params.ptr_A);
+    ElementB* ptr_B = static_cast<ElementB*>(params.ptr_B);
 
 #if SPLIT_K_ENABLED
-        //
-        // Fetch pointers based on mode.
-        //
-        if (params.mode == GemmUniversalMode::kGemm || params.mode == GemmUniversalMode::kGemmSplitKParallel)
-        {
+    //
+    // Fetch pointers based on mode.
+    //
+    if (params.mode == GemmUniversalMode::kGemm ||
+        params.mode == GemmUniversalMode::kGemmSplitKParallel) {
+      if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
+        problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
+      }
 
-            if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k())
-            {
-
-                problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
-            }
-
-            offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
-        }
-        else if (params.mode == GemmUniversalMode::kBatched)
-        {
-            ptr_A += threadblock_tile_offset.k() * params.batch_stride_A;
-            ptr_B += threadblock_tile_offset.k() * params.batch_stride_B;
-        }
-        else if (params.mode == GemmUniversalMode::kArray)
-        {
-            ptr_A = static_cast<ElementA* const*>(params.ptr_A)[threadblock_tile_offset.k()];
-            ptr_B = static_cast<ElementB* const*>(params.ptr_B)[threadblock_tile_offset.k()];
-        }
+      offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
+    } else if (params.mode == GemmUniversalMode::kBatched) {
+      ptr_A += threadblock_tile_offset.k() * params.batch_stride_A;
+      ptr_B += threadblock_tile_offset.k() * params.batch_stride_B;
+    } else if (params.mode == GemmUniversalMode::kArray) {
+      ptr_A = static_cast<ElementA* const*>(
+          params.ptr_A)[threadblock_tile_offset.k()];
+      ptr_B = static_cast<ElementB* const*>(
+          params.ptr_B)[threadblock_tile_offset.k()];
+    }
 #endif
 
-        // Compute initial location in logical coordinates
-        cutlass::MatrixCoord tb_offset_A{
-            threadblock_tile_offset.m() * Mma::Shape::kM,
-            offset_k,
-        };
+    // Compute initial location in logical coordinates
+    cutlass::MatrixCoord tb_offset_A{
+        threadblock_tile_offset.m() * Mma::Shape::kM,
+        offset_k,
+    };
 
-        cutlass::MatrixCoord tb_offset_B{offset_k, threadblock_tile_offset.n() * Mma::Shape::kN};
+    cutlass::MatrixCoord tb_offset_B{
+        offset_k, threadblock_tile_offset.n() * Mma::Shape::kN};
 
-        // Compute position within threadblock
-        int thread_idx = threadIdx.x;
+    // Compute position within threadblock
+    int thread_idx = threadIdx.x;
 
-        // Construct iterators to A and B operands
-        typename Mma::IteratorA iterator_A(
-            params.params_A, ptr_A, {params.problem_size.m(), problem_size_k}, thread_idx, tb_offset_A);
+    // Construct iterators to A and B operands
+    typename Mma::IteratorA iterator_A(
+        params.params_A,
+        ptr_A,
+        {params.problem_size.m(), problem_size_k},
+        thread_idx,
+        tb_offset_A);
 
-        typename Mma::IteratorB iterator_B(
-            params.params_B, ptr_B, {problem_size_k, params.problem_size.n()}, thread_idx, tb_offset_B);
+    typename Mma::IteratorB iterator_B(
+        params.params_B,
+        ptr_B,
+        {problem_size_k, params.problem_size.n()},
+        thread_idx,
+        tb_offset_B);
 
-        // Broadcast the warp_id computed by lane 0 to ensure dependent code
-        // is compiled as warp-uniform.
-        int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+    // Broadcast the warp_id computed by lane 0 to ensure dependent code
+    // is compiled as warp-uniform.
+    int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
 
-        int lane_idx = threadIdx.x % 32;
+    int lane_idx = threadIdx.x % 32;
 
-        //
-        // Main loop
-        //
+    //
+    // Main loop
+    //
 
-        // Construct thread-scoped matrix multiply
-        Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
+    // Construct thread-scoped matrix multiply
+    Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
 
-        typename Mma::FragmentC accumulators;
+    typename Mma::FragmentC accumulators;
 
-        accumulators.clear();
+    accumulators.clear();
 
-        // Compute threadblock-scoped matrix multiply-add
-        int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
+    // Compute threadblock-scoped matrix multiply-add
+    int gemm_k_iterations =
+        (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
 
-        // Compute threadblock-scoped matrix multiply-add
-        mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators);
+    // Compute threadblock-scoped matrix multiply-add
+    mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators);
 
-        //
-        // Masked tile iterators constructed from members
-        //
+    //
+    // Masked tile iterators constructed from members
+    //
 
-        threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+    threadblock_tile_offset =
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-        // assume identity swizzle
-        MatrixCoord threadblock_offset(
-            threadblock_tile_offset.m() * Mma::Shape::kM, threadblock_tile_offset.n() * Mma::Shape::kN);
+    // assume identity swizzle
+    MatrixCoord threadblock_offset(
+        threadblock_tile_offset.m() * Mma::Shape::kM,
+        threadblock_tile_offset.n() * Mma::Shape::kN);
 
-        int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m();
+    int block_idx = threadblock_tile_offset.m() +
+                    threadblock_tile_offset.n() * params.grid_tiled_shape.m();
 
-        //
-        // Construct the epilogue visitor
-        //
+    //
+    // Construct the epilogue visitor
+    //
 
-        EpilogueVisitor epilogue_visitor(params.epilogue_visitor, shared_storage.epilogue.visitor,
-            params.problem_size.mn(), thread_idx, warp_idx, lane_idx, params.params_alpha_col, params.params_C,
-            params.params_D, params.quant_option, params.ptr_alpha_row, params.ptr_alpha_col, params.ptr_C,
-            params.ptr_D, threadblock_offset, blockIdx.y * params.problem_size.m());
+    EpilogueVisitor epilogue_visitor(params.epilogue_visitor,
+                                     shared_storage.epilogue.visitor,
+                                     params.problem_size.mn(),
+                                     thread_idx,
+                                     warp_idx,
+                                     lane_idx,
+                                     params.params_alpha_col,
+                                     params.params_C,
+                                     params.params_D,
+                                     params.quant_option,
+                                     params.ptr_alpha_row,
+                                     params.ptr_alpha_col,
+                                     params.ptr_C,
+                                     params.ptr_D,
+                                     threadblock_offset,
+                                     blockIdx.y * params.problem_size.m());
 
-        if (params.mode == GemmUniversalMode::kGemm)
-        {
-            // Indicate which position in a serial reduction the output operator is currently updating
-            epilogue_visitor.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k());
-        }
-        else if (params.mode == GemmUniversalMode::kBatched || params.mode == GemmUniversalMode::kArray)
-        {
-            epilogue_visitor.set_batch_index(threadblock_tile_offset.k());
-        }
-
-        // Construct the epilogue
-        Epilogue epilogue(shared_storage.epilogue.epilogue, thread_idx, warp_idx, lane_idx);
-
-        // Execute the epilogue operator to update the destination tensor.
-        epilogue(epilogue_visitor, accumulators);
+    if (params.mode == GemmUniversalMode::kGemm) {
+      // Indicate which position in a serial reduction the output operator is
+      // currently updating
+      epilogue_visitor.set_k_partition(threadblock_tile_offset.k(),
+                                       params.grid_tiled_shape.k());
+    } else if (params.mode == GemmUniversalMode::kBatched ||
+               params.mode == GemmUniversalMode::kArray) {
+      epilogue_visitor.set_batch_index(threadblock_tile_offset.k());
     }
 
-    template <typename CompilationArch>
-    CUTLASS_DEVICE void run_kernel(Params const& params, SharedStorage& shared_storage)
-    {
-        if constexpr (platform::is_same<ArchTag, CompilationArch>::value)
-        {
-            run_kernel_(params, shared_storage);
-        }
-        else
-        {
-            CUTLASS_NOT_IMPLEMENTED();
-        }
-    }
+    // Construct the epilogue
+    Epilogue epilogue(
+        shared_storage.epilogue.epilogue, thread_idx, warp_idx, lane_idx);
 
-    /*
-        To improve compilation speed, we do not compile the device operator if the CUDA_ARCH does not correspond
-        to the ArchTag of the cutlass kernel operator.
-      */
-    /// Executes one GEMM
-    CUTLASS_DEVICE
-    void operator()(Params const& params, SharedStorage& shared_storage)
-    {
+    // Execute the epilogue operator to update the destination tensor.
+    epilogue(epilogue_visitor, accumulators);
+  }
+
+  template <typename CompilationArch>
+  CUTLASS_DEVICE void run_kernel(Params const& params,
+                                 SharedStorage& shared_storage) {
+    if constexpr (platform::is_same<ArchTag, CompilationArch>::value) {
+      run_kernel_(params, shared_storage);
+    } else {
+      CUTLASS_NOT_IMPLEMENTED();
+    }
+  }
+
+  /*
+      To improve compilation speed, we do not compile the device operator if the
+     CUDA_ARCH does not correspond to the ArchTag of the cutlass kernel
+     operator.
+    */
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const& params, SharedStorage& shared_storage) {
 #if defined(__CUDA_ARCH__)
 #if (__CUDA_ARCH__ >= 700) && (__CUDA_ARCH__ < 720)
-        run_kernel<arch::Sm70>(params, shared_storage);
+    run_kernel<arch::Sm70>(params, shared_storage);
 #elif (__CUDA_ARCH__ >= 720) && (__CUDA_ARCH__ < 750)
-        run_kernel<arch::Sm72>(params, shared_storage);
+    run_kernel<arch::Sm72>(params, shared_storage);
 #elif (__CUDA_ARCH__ >= 750) && (__CUDA_ARCH__ < 800)
-        run_kernel<arch::Sm75>(params, shared_storage);
+    run_kernel<arch::Sm75>(params, shared_storage);
 #elif (__CUDA_ARCH__ >= 800) && (__CUDA_ARCH__ < 900)
-        run_kernel<arch::Sm80>(params, shared_storage);
+    run_kernel<arch::Sm80>(params, shared_storage);
 #elif (__CUDA_ARCH__ >= 900)
-        // TODO - replace with CUTLASS_NOT_IMPLEMENTED() and upgrade to 3.x kernels.
-        run_kernel<arch::Sm80>(params, shared_storage);
+    // TODO - replace with CUTLASS_NOT_IMPLEMENTED() and upgrade to 3.x kernels.
+    run_kernel<arch::Sm80>(params, shared_storage);
 #else
-        static_assert(
-            false, "Invalid architecture being compiled. Only Volta+ supported in weight-only quantization kernels.");
+    static_assert(false,
+                  "Invalid architecture being compiled. Only Volta+ supported "
+                  "in weight-only quantization kernels.");
 #endif
 #else
-        CUTLASS_NOT_IMPLEMENTED();
+    CUTLASS_NOT_IMPLEMENTED();
 #endif
-    }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace kernel
-} // namespace gemm
-} // namespace cutlass
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h
index 8f61c6d9c4..56140c81d6 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,9 +15,10 @@
  * limitations under the License.
  */
 /*
-  This file exists so that we use the same weight layout for MoE grouped gemm and regular gemm when the weight is
-  quantized. The preprocessing code reads this template to know how to organize the quantized weight matrices
-  to be consumed by CUTLASS.
+  This file exists so that we use the same weight layout for MoE grouped gemm
+  and regular gemm when the weight is quantized. The preprocessing code reads
+  this template to know how to organize the quantized weight matrices to be
+  consumed by CUTLASS.
 
   Note that for int4, ThreadBlockK MUST be 64.
 
@@ -35,136 +36,172 @@
 #include "cutlass_extensions/arch/mma.h"
 #include "cutlass_extensions/tile_interleaved_layout.h"
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace kernel
-{
+namespace cutlass {
+namespace gemm {
+namespace kernel {
 
 template <typename TypeA, typename TypeB, typename Arch, typename Enable = void>
-struct LayoutDetailsB
-{
-};
+struct LayoutDetailsB {};
 
-// Volta specialiations. Volta will dequantize before STS, so we need a different operator
+// Volta specialiations. Volta will dequantize before STS, so we need a
+// different operator
 template <typename TypeA, typename TypeB>
-struct LayoutDetailsB<TypeA, TypeB, arch::Sm70>
-{
-    static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value;
-    using Layout = layout::ColumnMajor;
-    static constexpr int ElementsPerAccess = 8;
-    using Operator = cutlass::arch::OpMultiplyAdd;
+struct LayoutDetailsB<TypeA, TypeB, arch::Sm70> {
+  static constexpr int ThreadblockK =
+      128 * 8 / cutlass::sizeof_bits<TypeA>::value;
+  using Layout = layout::ColumnMajor;
+  static constexpr int ElementsPerAccess = 8;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
-// Specializations for Turing+ when B is FP16. These are currently only used for MoE networks.
-// TODO - Switch this to column major for weights since gemms should be more performant.
+// Specializations for Turing+ when B is FP16. These are currently only used for
+// MoE networks.
+// TODO - Switch this to column major for weights since gemms should be more
+// performant.
 template <typename TypeA, typename Arch>
-struct LayoutDetailsB<TypeA, half_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type>
-{
-    static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value;
-    using Layout = layout::ColumnMajor;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<half_t>::value;
-    using Operator = cutlass::arch::OpMultiplyAdd;
+struct LayoutDetailsB<
+    TypeA,
+    half_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK =
+      128 * 8 / cutlass::sizeof_bits<TypeA>::value;
+  using Layout = layout::ColumnMajor;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<half_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
 template <typename TypeA, typename Arch>
-struct LayoutDetailsB<TypeA, bfloat16_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type>
-{
-    static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value;
-    using Layout = layout::ColumnMajor;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<bfloat16_t>::value;
-    using Operator = cutlass::arch::OpMultiplyAdd;
+struct LayoutDetailsB<
+    TypeA,
+    bfloat16_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK =
+      128 * 8 / cutlass::sizeof_bits<TypeA>::value;
+  using Layout = layout::ColumnMajor;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<bfloat16_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
 template <typename TypeA>
-struct LayoutDetailsB<TypeA, cutlass::float_e4m3_t, arch::Sm89>
-{
-    static constexpr int ThreadblockK = 64;
+struct LayoutDetailsB<TypeA, cutlass::float_e4m3_t, arch::Sm89> {
+  static constexpr int ThreadblockK = 64;
 
-private:
-    static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits<uint8_t>::value;
-    static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
+ private:
+  static constexpr int ElementsPerCacheLine =
+      128 * 8 / sizeof_bits<uint8_t>::value;
+  static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
 
-public:
-    using Layout = layout::ColumnMajor;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<cutlass::float_e4m3_t>::value;
-    using Operator = cutlass::arch::OpMultiplyAdd;
-    // for fast accumulation
-    // using Operator = cutlass::arch::OpMultiplyAddFastAccum;
+ public:
+  using Layout = layout::ColumnMajor;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<cutlass::float_e4m3_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAdd;
+  // for fast accumulation
+  // using Operator = cutlass::arch::OpMultiplyAddFastAccum;
 };
 
-// Specializations for Turing+ when B is quantized. These can use the operator OpMultiplyAddDequantizeInterleavedBToA,
-// which signals that we want to dequantize after loading from smem.
+// Specializations for Turing+ when B is quantized. These can use the operator
+// OpMultiplyAddDequantizeInterleavedBToA, which signals that we want to
+// dequantize after loading from smem.
 template <typename TypeA, typename Arch>
     struct LayoutDetailsB < TypeA,
     uint8_t, Arch,
-    typename platform::enable_if<Arch::kMinComputeCapability >= 75 && Arch::kMinComputeCapability<90>::type>
-{
-    static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value;
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75 &&
+                                 Arch::kMinComputeCapability<90>::type> {
+  static constexpr int ThreadblockK =
+      128 * 8 / cutlass::sizeof_bits<TypeA>::value;
 
-private:
-    static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits<uint8_t>::value;
-    static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
+ private:
+  static constexpr int ElementsPerCacheLine =
+      128 * 8 / sizeof_bits<uint8_t>::value;
+  static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
 
-public:
-    using Layout = layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<uint8_t>::value;
-    using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
+ public:
+  using Layout =
+      layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<uint8_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
 };
 
 template <typename TypeA, typename Arch>
     struct LayoutDetailsB < TypeA,
     uint4b_t, Arch,
-    typename platform::enable_if<Arch::kMinComputeCapability >= 75 && Arch::kMinComputeCapability<90>::type>
-{
-    static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value;
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75 &&
+                                 Arch::kMinComputeCapability<90>::type> {
+  static constexpr int ThreadblockK =
+      128 * 8 / cutlass::sizeof_bits<TypeA>::value;
 
-private:
-    static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits<uint4b_t>::value;
-    static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
+ private:
+  static constexpr int ElementsPerCacheLine =
+      128 * 8 / sizeof_bits<uint4b_t>::value;
+  static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
 
-public:
-    using Layout = layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<uint4b_t>::value;
-    using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
+ public:
+  using Layout =
+      layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<uint4b_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
 };
 
 template <typename TypeA, typename Arch>
-struct LayoutDetailsB<TypeA, uint2b_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type>
-{
-    static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value; // 64
+struct LayoutDetailsB<
+    TypeA,
+    uint2b_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK =
+      128 * 8 / cutlass::sizeof_bits<TypeA>::value;  // 64
 
-private:
-    static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits<uint2b_t>::value;
-    static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK; // 8
+ private:
+  static constexpr int ElementsPerCacheLine =
+      128 * 8 / sizeof_bits<uint2b_t>::value;
+  static constexpr int ColumnsInterleaved =
+      ElementsPerCacheLine / ThreadblockK;  // 8
 
-public:
-    // using Layout = layout::ColumnMajor;
-    // static constexpr int ElementsPerAccess = 16; // at least 4-bytes
-    using Layout = layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<uint2b_t>::value; // 64
-    using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
+ public:
+  // using Layout = layout::ColumnMajor;
+  // static constexpr int ElementsPerAccess = 16; // at least 4-bytes
+  using Layout =
+      layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<uint2b_t>::value;  // 64
+  using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
 };
 
 template <typename TypeA, typename Arch>
-struct LayoutDetailsB<TypeA, uint8_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 90>::type>
-{
-    static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value;
-    using Layout = layout::ColumnMajor;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<half_t>::value;
-    using Operator = cutlass::arch::OpMultiplyAdd;
+struct LayoutDetailsB<
+    TypeA,
+    uint8_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 90>::type> {
+  static constexpr int ThreadblockK =
+      128 * 8 / cutlass::sizeof_bits<TypeA>::value;
+  using Layout = layout::ColumnMajor;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<half_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
 template <typename TypeA, typename Arch>
-struct LayoutDetailsB<TypeA, uint4b_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 90>::type>
-{
-    static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value;
-    using Layout = layout::ColumnMajor;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<half_t>::value;
-    using Operator = cutlass::arch::OpMultiplyAdd;
+struct LayoutDetailsB<
+    TypeA,
+    uint4b_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 90>::type> {
+  static constexpr int ThreadblockK =
+      128 * 8 / cutlass::sizeof_bits<TypeA>::value;
+  using Layout = layout::ColumnMajor;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<half_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
-} // namespace kernel
-} // namespace gemm
-} // namespace cutlass
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/moe_problem_visitor.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/moe_problem_visitor.h
index b9126e3500..a2376e6b8b 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/moe_problem_visitor.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/moe_problem_visitor.h
@@ -160,11 +160,11 @@ struct BaseMoeProblemVisitor {
 
   CUTLASS_HOST_DEVICE
   cutlass::gemm::GemmCoord problem_size(int idx) const {
-
     int64_t gemm_m = 0;
 
     if (params.total_rows < 0) {
-      const int64_t prev_problem_row = idx == 0 ? 0 : params.last_row_for_problem[idx - 1];
+      const int64_t prev_problem_row =
+          idx == 0 ? 0 : params.last_row_for_problem[idx - 1];
       const int64_t current_problem_row = params.last_row_for_problem[idx];
       gemm_m = current_problem_row - prev_problem_row;
     } else {
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_cooperative.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_cooperative.hpp
index 843529cde5..efb90a5e7c 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_cooperative.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_cooperative.hpp
@@ -53,15 +53,20 @@ namespace cutlass::gemm::kernel {
 
 ///////////////////////////////////////////////////////////////////////////////
 
-template <class ProblemShape_, class CollectiveMainloop_,
-          class CollectiveEpilogue_, class TileScheduler_>
+template <class ProblemShape_,
+          class CollectiveMainloop_,
+          class CollectiveEpilogue_,
+          class TileScheduler_>
 class GemmUniversalGated<
-    ProblemShape_, CollectiveMainloop_, CollectiveEpilogue_, TileScheduler_,
+    ProblemShape_,
+    CollectiveMainloop_,
+    CollectiveEpilogue_,
+    TileScheduler_,
     cute::enable_if_t<cute::is_base_of_v<KernelTmaWarpSpecializedCooperative,
                                          typename CollectiveMainloop_::
                                              DispatchPolicy::Schedule> &&
                       CollectiveMainloop_::isGated>> {
-public:
+ public:
   //
   // Type Aliases
   //
@@ -98,7 +103,9 @@ public:
 
   using TileSchedulerTag = TileScheduler_;
   using TileScheduler =
-      typename detail::TileSchedulerSelector<TileScheduler_, ArchTag, TileShape,
+      typename detail::TileSchedulerSelector<TileScheduler_,
+                                             ArchTag,
+                                             TileShape,
                                              ClusterShape>::Scheduler;
   using TileSchedulerArguments = typename TileScheduler::Arguments;
   using TileSchedulerParams = typename TileScheduler::Params;
@@ -222,9 +229,14 @@ public:
     // used, therefore separate reduction will not be enabled.
     constexpr uint32_t NumEpilogueSubTiles =
         CollectiveEpilogue::get_store_pipe_increment(TileShape{});
-    TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
-        problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info,
-        args.scheduler, scheduler_workspace, NumEpilogueSubTiles);
+    TileSchedulerParams scheduler =
+        TileScheduler::to_underlying_arguments(problem_shape_MNKL,
+                                               TileShape{},
+                                               ClusterShape{},
+                                               hw_info,
+                                               args.scheduler,
+                                               scheduler_workspace,
+                                               NumEpilogueSubTiles);
 
     return {args.mode,
             problem_shape,
@@ -242,8 +254,9 @@ public:
                          (args.mode == GemmUniversalMode::kBatched &&
                           cute::rank(ProblemShape{}) == 4);
     if (!implementable) {
-      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Arguments or Problem Shape don't "
-                         "meet the requirements.\n");
+      CUTLASS_TRACE_HOST(
+          "  CAN IMPLEMENT: Arguments or Problem Shape don't "
+          "meet the requirements.\n");
       return implementable;
     }
     implementable &=
@@ -262,7 +275,10 @@ public:
     workspace_size +=
         TileScheduler::template get_workspace_size<ProblemShape,
                                                    ElementAccumulator>(
-            args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups,
+            args.scheduler,
+            args.problem_shape,
+            args.hw_info,
+            NumMmaWarpGroups,
             NumEpilogueSubTiles);
     workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
 
@@ -273,10 +289,11 @@ public:
     return workspace_size;
   }
 
-  static cutlass::Status
-  initialize_workspace(Arguments const &args, void *workspace = nullptr,
-                       cudaStream_t stream = nullptr,
-                       CudaHostAdapter *cuda_adapter = nullptr) {
+  static cutlass::Status initialize_workspace(
+      Arguments const &args,
+      void *workspace = nullptr,
+      cudaStream_t stream = nullptr,
+      CudaHostAdapter *cuda_adapter = nullptr) {
     Status status = Status::kSuccess;
     uint8_t *workspace_ptr = reinterpret_cast<uint8_t *>(workspace);
     size_t workspace_offset = 0;
@@ -285,13 +302,20 @@ public:
 
     status = TileScheduler::template initialize_workspace<ProblemShape,
                                                           ElementAccumulator>(
-        args.scheduler, workspace_ptr + workspace_offset, stream,
-        args.problem_shape, args.hw_info, NumMmaWarpGroups,
+        args.scheduler,
+        workspace_ptr + workspace_offset,
+        stream,
+        args.problem_shape,
+        args.hw_info,
+        NumMmaWarpGroups,
         NumEpilogueSubTiles);
     workspace_offset +=
         TileScheduler::template get_workspace_size<ProblemShape,
                                                    ElementAccumulator>(
-            args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups,
+            args.scheduler,
+            args.problem_shape,
+            args.hw_info,
+            NumMmaWarpGroups,
             NumEpilogueSubTiles);
     workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
     if (status != Status::kSuccess) {
@@ -299,8 +323,11 @@ public:
     }
 
     status = CollectiveEpilogue::initialize_workspace(
-        args.problem_shape, args.epilogue, workspace_ptr + workspace_offset,
-        stream, cuda_adapter);
+        args.problem_shape,
+        args.epilogue,
+        workspace_ptr + workspace_offset,
+        stream,
+        cuda_adapter);
     workspace_offset += CollectiveEpilogue::get_workspace_size(
         args.problem_shape, args.epilogue);
     workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
@@ -323,9 +350,12 @@ public:
         params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN
             ? TileScheduler::RasterOrderOptions::AlongN
             : TileScheduler::RasterOrderOptions::AlongM;
-    return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape,
-                                         TileShape{}, ClusterShape{},
-                                         params.hw_info, args);
+    return TileScheduler::get_grid_shape(params.scheduler,
+                                         params.problem_shape,
+                                         TileShape{},
+                                         ClusterShape{},
+                                         params.hw_info,
+                                         args);
   }
 
   static dim3 get_block_shape() { return dim3(MaxThreadsPerBlock, 1, 1); }
@@ -337,8 +367,9 @@ public:
 
 // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
 #if !defined(__CUDA_ARCH_FEAT_SM90_ALL)
-    printf("ERROR : Arch conditional MMA instruction used without targeting "
-           "sm90a compute capability. Aborting.\n");
+    printf(
+        "ERROR : Arch conditional MMA instruction used without targeting "
+        "sm90a compute capability. Aborting.\n");
 #else
 
     // Preconditions
@@ -469,7 +500,7 @@ public:
         return []() { cute::cluster_wait(); };
       } else {
         __syncthreads();
-        return []() {}; // do nothing
+        return []() {};  // do nothing
       }
     }();
 
@@ -480,7 +511,7 @@ public:
     // Get the appropriate blocks for this thread block -- potential for thread
     // block locality
     TiledMma tiled_mma;
-    auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
+    auto blk_shape = TileShape{};  // (BLK_M,BLK_N,BLK_K)
 
     TileScheduler scheduler{params.scheduler};
     auto work_tile_info = scheduler.get_current_work();
@@ -540,10 +571,16 @@ public:
           auto k_tile_iter = cute::make_coord_iterator(
               idx2crd(work_k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl));
 
-          collective_mainloop.load(
-              params.mainloop, mainloop_pipeline, mainloop_pipe_producer_state,
-              load_inputs, blk_coord, k_tile_iter, work_k_tile_count, lane_idx,
-              block_rank_in_cluster, shared_storage.tensors.mainloop);
+          collective_mainloop.load(params.mainloop,
+                                   mainloop_pipeline,
+                                   mainloop_pipe_producer_state,
+                                   load_inputs,
+                                   blk_coord,
+                                   k_tile_iter,
+                                   work_k_tile_count,
+                                   lane_idx,
+                                   block_rank_in_cluster,
+                                   shared_storage.tensors.mainloop);
           // Update starting pipeline state for the next tile
           mainloop_pipe_producer_state.advance(work_k_tile_count);
 
@@ -555,12 +592,12 @@ public:
 
           // Get next work tile
           work_tile_info = fetch_next_work(work_tile_info, scheduler);
-        } // Scheduler work fetch loop
+        }  // Scheduler work fetch loop
 
         // Make sure all Consumer Warp Groups have been waited upon
         collective_mainloop.load_tail(mainloop_pipeline,
                                       mainloop_pipe_producer_state);
-      } // Mainloop Producer Warp End
+      }  // Mainloop Producer Warp End
 
       // Epilogue Producer Warp
       else if (producer_warp_role == ProducerWarpRole::Epilogue &&
@@ -579,21 +616,26 @@ public:
             auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
 
             epi_load_pipe_producer_state = collective_epilogue.load(
-                epi_load_pipeline, epi_load_pipe_producer_state,
-                problem_shape_MNKL, blk_shape, blk_coord, tiled_mma, lane_idx,
+                epi_load_pipeline,
+                epi_load_pipe_producer_state,
+                problem_shape_MNKL,
+                blk_shape,
+                blk_coord,
+                tiled_mma,
+                lane_idx,
                 shared_storage.tensors.epilogue,
                 work_tile_info.reduction_subtile_idx());
           }
 
           // Get next work tile
           work_tile_info = fetch_next_work(work_tile_info, scheduler);
-        } // Scheduler work fetch loop
+        }  // Scheduler work fetch loop
 
         // Make sure all Consumer Warp Groups have been waited upon
         collective_epilogue.load_tail(epi_load_pipeline,
                                       epi_load_pipe_producer_state);
-      } // Epilogue Producer Warp End
-    } // Producer Warp Group End
+      }  // Epilogue Producer Warp End
+    }    // Producer Warp Group End
 
     else if (warp_group_role == WarpGroupRole::Consumer0 ||
              warp_group_role == WarpGroupRole::Consumer1) {
@@ -618,14 +660,18 @@ public:
         //
         // MSVC CTAD breaks if we say "Tensor" here, so we use "auto" instead.
         auto accumulators0 = partition_fragment_C(
-            tiled_mma, take<0, 2>(blk_shape)); // (MMA,MMA_M,MMA_N)
+            tiled_mma, take<0, 2>(blk_shape));  // (MMA,MMA_M,MMA_N)
         auto accumulators1 = partition_fragment_C(
-            tiled_mma, take<0, 2>(blk_shape)); // (MMA,MMA_M,MMA_N)
+            tiled_mma, take<0, 2>(blk_shape));  // (MMA,MMA_M,MMA_N)
         if (TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
-          collective_mainloop.mma(
-              mainloop_pipeline, mainloop_pipe_consumer_state, accumulators0,
-              accumulators1, work_k_tile_count, mma_thread_idx,
-              shared_storage.tensors.mainloop, params.mainloop);
+          collective_mainloop.mma(mainloop_pipeline,
+                                  mainloop_pipe_consumer_state,
+                                  accumulators0,
+                                  accumulators1,
+                                  work_k_tile_count,
+                                  mma_thread_idx,
+                                  shared_storage.tensors.mainloop,
+                                  params.mainloop);
 
           // Make sure the math instructions are done and free buffers before
           // entering the epilogue
@@ -641,10 +687,16 @@ public:
             canonical_warp_group_idx() - NumLoadWarpGroups;
 
         // Perform reduction across splits, if needed
-        TileScheduler::fixup(params.scheduler, work_tile_info, accumulators0,
-                             NumMmaWarpGroups, consumer_warp_group_idx);
-        TileScheduler::fixup(params.scheduler, work_tile_info, accumulators1,
-                             NumMmaWarpGroups, consumer_warp_group_idx);
+        TileScheduler::fixup(params.scheduler,
+                             work_tile_info,
+                             accumulators0,
+                             NumMmaWarpGroups,
+                             consumer_warp_group_idx);
+        TileScheduler::fixup(params.scheduler,
+                             work_tile_info,
+                             accumulators1,
+                             NumMmaWarpGroups,
+                             consumer_warp_group_idx);
 
         Activation elt_op;
         CUTLASS_PRAGMA_UNROLL
@@ -657,12 +709,18 @@ public:
           // Epilogue and write to gD
           auto [epi_load_pipe_consumer_state_next,
                 epi_store_pipe_producer_state_next] =
-              collective_epilogue.store(
-                  epi_load_pipeline, epi_load_pipe_consumer_state,
-                  epi_store_pipeline, epi_store_pipe_producer_state,
-                  problem_shape_MNKL, blk_shape, blk_coord, accumulators0,
-                  tiled_mma, mma_thread_idx, shared_storage.tensors.epilogue,
-                  work_tile_info.reduction_subtile_idx());
+              collective_epilogue.store(epi_load_pipeline,
+                                        epi_load_pipe_consumer_state,
+                                        epi_store_pipeline,
+                                        epi_store_pipe_producer_state,
+                                        problem_shape_MNKL,
+                                        blk_shape,
+                                        blk_coord,
+                                        accumulators0,
+                                        tiled_mma,
+                                        mma_thread_idx,
+                                        shared_storage.tensors.epilogue,
+                                        work_tile_info.reduction_subtile_idx());
           epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next;
           epi_store_pipe_producer_state = epi_store_pipe_producer_state_next;
           do_store_tail = true;
@@ -670,23 +728,24 @@ public:
 
         // Get next work tile
         work_tile_info = fetch_next_work(work_tile_info, scheduler);
-      } // Scheduler work fetch loop
+      }  // Scheduler work fetch loop
 
       if (do_store_tail) {
-        collective_epilogue.store_tail(
-            epi_load_pipeline, epi_load_pipe_consumer_state, epi_store_pipeline,
-            epi_store_pipe_producer_state);
+        collective_epilogue.store_tail(epi_load_pipeline,
+                                       epi_load_pipe_consumer_state,
+                                       epi_store_pipeline,
+                                       epi_store_pipe_producer_state);
       }
-    } // Consumer Warp Groups End
+    }  // Consumer Warp Groups End
 #endif
   }
 
-private:
+ private:
   // Kernel helper function to get next work unit
   CUTLASS_DEVICE
-  typename TileScheduler::WorkTileInfo
-  fetch_next_work(typename TileScheduler::WorkTileInfo &work_tile_info,
-                  TileScheduler &scheduler) const {
+  typename TileScheduler::WorkTileInfo fetch_next_work(
+      typename TileScheduler::WorkTileInfo &work_tile_info,
+      TileScheduler &scheduler) const {
     // Check whether we should continue on with the current work unit. If this
     // is the case, the work unit will have been updated in
     // continue_current_work to reflect the new tile to be computed.
@@ -702,4 +761,4 @@ private:
 
 ///////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm::kernel
+}  // namespace cutlass::gemm::kernel
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_pingpong.hpp b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_pingpong.hpp
index e6cc7de5c6..9609adc32a 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_pingpong.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_pingpong.hpp
@@ -56,15 +56,20 @@ namespace cutlass::gemm::kernel {
 
 ///////////////////////////////////////////////////////////////////////////////
 
-template <class ProblemShape_, class CollectiveMainloop_,
-          class CollectiveEpilogue_, class TileScheduler_>
+template <class ProblemShape_,
+          class CollectiveMainloop_,
+          class CollectiveEpilogue_,
+          class TileScheduler_>
 class GemmUniversalGated<
-    ProblemShape_, CollectiveMainloop_, CollectiveEpilogue_, TileScheduler_,
+    ProblemShape_,
+    CollectiveMainloop_,
+    CollectiveEpilogue_,
+    TileScheduler_,
     cute::enable_if_t<cute::is_base_of_v<KernelTmaWarpSpecializedPingpong,
                                          typename CollectiveMainloop_::
                                              DispatchPolicy::Schedule> &&
                       CollectiveMainloop_::isGated>> {
-public:
+ public:
   //
   // Type Aliases
   //
@@ -103,7 +108,9 @@ public:
       "Ping-pong kernel does not currently support stream-K scheduler.");
   using TileSchedulerTag = TileScheduler_;
   using TileScheduler =
-      typename detail::TileSchedulerSelector<TileScheduler_, ArchTag, TileShape,
+      typename detail::TileSchedulerSelector<TileScheduler_,
+                                             ArchTag,
+                                             TileShape,
                                              ClusterShape>::Scheduler;
   using TileSchedulerArguments = typename TileScheduler::Arguments;
   using TileSchedulerParams = typename TileScheduler::Params;
@@ -236,9 +243,12 @@ public:
             CollectiveEpilogue::to_underlying_arguments(
                 args.problem_shape, args.epilogue, epilogue_workspace),
             hw_info,
-            TileScheduler::to_underlying_arguments(
-                problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info,
-                args.scheduler, scheduler_workspace)};
+            TileScheduler::to_underlying_arguments(problem_shape_MNKL,
+                                                   TileShape{},
+                                                   ClusterShape{},
+                                                   hw_info,
+                                                   args.scheduler,
+                                                   scheduler_workspace)};
   }
 
   static bool can_implement(Arguments const &args) {
@@ -246,8 +256,9 @@ public:
                          (args.mode == GemmUniversalMode::kBatched &&
                           cute::rank(ProblemShape{}) == 4);
     if (!implementable) {
-      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Arguments or Problem Shape don't "
-                         "meet the requirements.\n");
+      CUTLASS_TRACE_HOST(
+          "  CAN IMPLEMENT: Arguments or Problem Shape don't "
+          "meet the requirements.\n");
       return implementable;
     }
     implementable &=
@@ -273,18 +284,23 @@ public:
     return workspace_size;
   }
 
-  static cutlass::Status
-  initialize_workspace(Arguments const &args, void *workspace = nullptr,
-                       cudaStream_t stream = nullptr,
-                       CudaHostAdapter *cuda_adapter = nullptr) {
+  static cutlass::Status initialize_workspace(
+      Arguments const &args,
+      void *workspace = nullptr,
+      cudaStream_t stream = nullptr,
+      CudaHostAdapter *cuda_adapter = nullptr) {
     Status status = Status::kSuccess;
     uint8_t *workspace_ptr = reinterpret_cast<uint8_t *>(workspace);
     size_t workspace_offset = 0;
 
     status = TileScheduler::template initialize_workspace<ProblemShape,
                                                           ElementAccumulator>(
-        args.scheduler, workspace_ptr + workspace_offset, stream,
-        args.problem_shape, args.hw_info, NumMmaWarpGroups);
+        args.scheduler,
+        workspace_ptr + workspace_offset,
+        stream,
+        args.problem_shape,
+        args.hw_info,
+        NumMmaWarpGroups);
     workspace_offset +=
         TileScheduler::template get_workspace_size<ProblemShape,
                                                    ElementAccumulator>(
@@ -295,8 +311,11 @@ public:
     }
 
     status = CollectiveEpilogue::initialize_workspace(
-        args.problem_shape, args.epilogue, workspace_ptr + workspace_offset,
-        stream, cuda_adapter);
+        args.problem_shape,
+        args.epilogue,
+        workspace_ptr + workspace_offset,
+        stream,
+        cuda_adapter);
     workspace_offset += CollectiveEpilogue::get_workspace_size(
         args.problem_shape, args.epilogue);
     workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
@@ -319,9 +338,12 @@ public:
         params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN
             ? TileScheduler::RasterOrderOptions::AlongN
             : TileScheduler::RasterOrderOptions::AlongM;
-    return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape,
-                                         TileShape{}, ClusterShape{},
-                                         params.hw_info, args);
+    return TileScheduler::get_grid_shape(params.scheduler,
+                                         params.problem_shape,
+                                         TileShape{},
+                                         ClusterShape{},
+                                         params.hw_info,
+                                         args);
   }
 
   static dim3 get_block_shape() { return dim3(MaxThreadsPerBlock, 1, 1); }
@@ -333,8 +355,9 @@ public:
 
 // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
 #if !defined(__CUDA_ARCH_FEAT_SM90_ALL)
-    printf("ERROR : Arch conditional MMA instruction used without targeting "
-           "sm90a compute capability. Aborting.\n");
+    printf(
+        "ERROR : Arch conditional MMA instruction used without targeting "
+        "sm90a compute capability. Aborting.\n");
 #else
 
     // Preconditions
@@ -437,7 +460,7 @@ public:
     params_math_wg_order_barrier.group_id =
         canonical_warp_group_idx() - static_cast<int>(WarpGroupRole::Consumer0);
     params_math_wg_order_barrier.group_size =
-        NumThreadsPerWarpGroup; // Number of threads / participants in a group
+        NumThreadsPerWarpGroup;  // Number of threads / participants in a group
     MathWarpGroupOrderBarrier math_wg_order_barrier(
         shared_storage.pipelines.math_wg_order, params_math_wg_order_barrier);
 
@@ -464,7 +487,7 @@ public:
         return []() { cute::cluster_wait(); };
       } else {
         __syncthreads();
-        return []() {}; // do nothing
+        return []() {};  // do nothing
       }
     }();
 
@@ -476,7 +499,7 @@ public:
     // Get the appropriate blocks for this thread block -- potential for thread
     // block locality
     TiledMma tiled_mma;
-    auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
+    auto blk_shape = TileShape{};  // (BLK_M,BLK_N,BLK_K)
 
     // In a warp specialized kernel, collectives expose data movement and
     // compute operations separately
@@ -535,10 +558,16 @@ public:
 
           auto k_tile_iter = cute::make_coord_iterator(shape<3>(gA_mkl));
 
-          collective_mainloop.load(
-              params.mainloop, mainloop_pipeline, mainloop_pipe_producer_state,
-              load_inputs, blk_coord, k_tile_iter, k_tile_count, lane_idx,
-              block_rank_in_cluster, shared_storage.tensors.mainloop);
+          collective_mainloop.load(params.mainloop,
+                                   mainloop_pipeline,
+                                   mainloop_pipe_producer_state,
+                                   load_inputs,
+                                   blk_coord,
+                                   k_tile_iter,
+                                   k_tile_count,
+                                   lane_idx,
+                                   block_rank_in_cluster,
+                                   shared_storage.tensors.mainloop);
           // Update starting pipeline state for the next tile
           mainloop_pipe_producer_state.advance(k_tile_count);
 
@@ -551,12 +580,12 @@ public:
           // Get next work tile
           scheduler.advance_to_next_work();
           work_tile_info = scheduler.get_current_work();
-        } // Scheduler work fetch loop
+        }  // Scheduler work fetch loop
 
         // Make sure all Consumer Warp Groups have been waited upon
         collective_mainloop.load_tail(mainloop_pipeline,
                                       mainloop_pipe_producer_state);
-      } // Mainloop Producer Warp End
+      }  // Mainloop Producer Warp End
 
       // Epilogue Producer Warp
       else if (producer_warp_role == ProducerWarpRole::Epilogue &&
@@ -570,21 +599,26 @@ public:
           auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
           auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
 
-          epi_load_pipe_producer_state = collective_epilogue.load(
-              epi_load_pipeline, epi_load_pipe_producer_state,
-              problem_shape_MNKL, blk_shape, blk_coord, tiled_mma, lane_idx,
-              shared_storage.tensors.epilogue);
+          epi_load_pipe_producer_state =
+              collective_epilogue.load(epi_load_pipeline,
+                                       epi_load_pipe_producer_state,
+                                       problem_shape_MNKL,
+                                       blk_shape,
+                                       blk_coord,
+                                       tiled_mma,
+                                       lane_idx,
+                                       shared_storage.tensors.epilogue);
 
           // Get next work tile
           scheduler.advance_to_next_work();
           work_tile_info = scheduler.get_current_work();
-        } // Scheduler work fetch loop
+        }  // Scheduler work fetch loop
 
         // Make sure all Consumer Warp Groups have been waited upon
         collective_epilogue.load_tail(epi_load_pipeline,
                                       epi_load_pipe_producer_state);
-      } // Epilogue Producer Warp End
-    } // Producer Warp Group End
+      }  // Epilogue Producer Warp End
+    }    // Producer Warp Group End
 
     else if (warp_group_role == WarpGroupRole::Consumer0 ||
              warp_group_role == WarpGroupRole::Consumer1) {
@@ -602,17 +636,21 @@ public:
 
         // Allocate the accumulators for the (M,N) blk_shape
         Tensor accumulators0 = partition_fragment_C(
-            tiled_mma, take<0, 2>(blk_shape)); // (MMA,MMA_M,MMA_N)
+            tiled_mma, take<0, 2>(blk_shape));  // (MMA,MMA_M,MMA_N)
         Tensor accumulators1 = partition_fragment_C(
-            tiled_mma, take<0, 2>(blk_shape)); // (MMA,MMA_M,MMA_N)
+            tiled_mma, take<0, 2>(blk_shape));  // (MMA,MMA_M,MMA_N)
 
         // Order two Math WG's MMA one after the other, helps hide Epilogue
         math_wg_order_barrier.wait();
 
-        collective_mainloop.mma(
-            mainloop_pipeline, mainloop_pipe_consumer_state, accumulators0,
-            accumulators1, k_tile_count, warp_group_thread_idx,
-            shared_storage.tensors.mainloop, params.mainloop);
+        collective_mainloop.mma(mainloop_pipeline,
+                                mainloop_pipe_consumer_state,
+                                accumulators0,
+                                accumulators1,
+                                k_tile_count,
+                                warp_group_thread_idx,
+                                shared_storage.tensors.mainloop,
+                                params.mainloop);
 
         // Cue for next Math WG's MMA to start
         math_wg_order_barrier.arrive();
@@ -637,12 +675,17 @@ public:
         // Epilogue and write to gD
         auto [epi_load_pipe_consumer_state_next,
               epi_store_pipe_producer_state_next] =
-            collective_epilogue.store(
-                epi_load_pipeline, epi_load_pipe_consumer_state,
-                epi_store_pipeline, epi_store_pipe_producer_state,
-                problem_shape_MNKL, blk_shape, blk_coord, accumulators0,
-                tiled_mma, warp_group_thread_idx,
-                shared_storage.tensors.epilogue);
+            collective_epilogue.store(epi_load_pipeline,
+                                      epi_load_pipe_consumer_state,
+                                      epi_store_pipeline,
+                                      epi_store_pipe_producer_state,
+                                      problem_shape_MNKL,
+                                      blk_shape,
+                                      blk_coord,
+                                      accumulators0,
+                                      tiled_mma,
+                                      warp_group_thread_idx,
+                                      shared_storage.tensors.epilogue);
 
         // TMA store pipeline wait is only visible to TMA-issuing warp, so for
         // multiple-consumer kernels we need to wait for all TMA stores to
@@ -651,9 +694,10 @@ public:
         // current consumer.
         auto [epi_load_pipe_consumer_state_next_,
               epi_store_pipe_producer_state_next_] =
-            collective_epilogue.store_tail(
-                epi_load_pipeline, epi_load_pipe_consumer_state_next,
-                epi_store_pipeline, epi_store_pipe_producer_state_next);
+            collective_epilogue.store_tail(epi_load_pipeline,
+                                           epi_load_pipe_consumer_state_next,
+                                           epi_store_pipeline,
+                                           epi_store_pipe_producer_state_next);
 
         // Update starting load/store pipeline states for the next tile
         // state has already been incremented by 1 tile in collective calls,
@@ -669,12 +713,12 @@ public:
         // Get next work tile
         scheduler.advance_to_next_work(NumMmaWarpGroups);
         work_tile_info = scheduler.get_current_work();
-      } // Scheduler work fetch loop
-    } // Consumer Warp Groups End
+      }  // Scheduler work fetch loop
+    }    // Consumer Warp Groups End
 #endif
   }
 };
 
 ///////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm::kernel
+}  // namespace cutlass::gemm::kernel
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/splitk_gemm_grouped.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/splitk_gemm_grouped.h
index 5e3531f093..5d68ea26d9 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/splitk_gemm_grouped.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/kernel/splitk_gemm_grouped.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
@@ -49,446 +50,471 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace kernel
-{
+namespace cutlass {
+namespace gemm {
+namespace kernel {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-template <typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate
-    typename Epilogue_,                   ///! Epilogue
-    typename ThreadblockSwizzle_,         ///! Threadblock swizzling function
-    GroupScheduleMode GroupScheduleMode_, ///! Type of scheduling to perform
-    bool Transposed = false>
-struct SplitkGemmGrouped
-{
-public:
-    using Mma = Mma_;
-    using Epilogue = Epilogue_;
-    using EpilogueOutputOp = typename Epilogue::OutputOp;
-    using ThreadblockSwizzle = ThreadblockSwizzle_;
-    static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_;
-    static bool const kTransposed = Transposed;
+template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
+          typename Epilogue_,            ///! Epilogue
+          typename ThreadblockSwizzle_,  ///! Threadblock swizzling function
+          GroupScheduleMode GroupScheduleMode_,  ///! Type of scheduling to
+                                                 /// perform
+          bool Transposed = false>
+struct SplitkGemmGrouped {
+ public:
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using EpilogueOutputOp = typename Epilogue::OutputOp;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_;
+  static bool const kTransposed = Transposed;
 
-    // Optional transpose
-    using MapArguments = kernel::detail::MapArguments<typename Mma::IteratorA::Element, typename Mma::IteratorA::Layout,
-        Mma::kTransformA, Mma::IteratorA::AccessType::kElements, typename Mma::IteratorB::Element,
-        typename Mma::IteratorB::Layout, Mma::kTransformB, Mma::IteratorB::AccessType::kElements, typename Mma::LayoutC,
-        kTransposed>;
+  // Optional transpose
+  using MapArguments =
+      kernel::detail::MapArguments<typename Mma::IteratorA::Element,
+                                   typename Mma::IteratorA::Layout,
+                                   Mma::kTransformA,
+                                   Mma::IteratorA::AccessType::kElements,
+                                   typename Mma::IteratorB::Element,
+                                   typename Mma::IteratorB::Layout,
+                                   Mma::kTransformB,
+                                   Mma::IteratorB::AccessType::kElements,
+                                   typename Mma::LayoutC,
+                                   kTransposed>;
 
-    // Public-facing type definitions related to operand element type, layout, and complex conjugate
-    // operation. Must interact with the 'kTransposed' notion.
-    using ElementA = typename MapArguments::ElementA;
-    using LayoutA = typename MapArguments::LayoutA;
-    using ElementB = typename MapArguments::ElementB;
-    using LayoutB = typename MapArguments::LayoutB;
-    using ElementC = typename Epilogue::OutputTileIterator::Element;
-    using LayoutC = typename MapArguments::LayoutC;
+  // Public-facing type definitions related to operand element type, layout, and
+  // complex conjugate operation. Must interact with the 'kTransposed' notion.
+  using ElementA = typename MapArguments::ElementA;
+  using LayoutA = typename MapArguments::LayoutA;
+  using ElementB = typename MapArguments::ElementB;
+  using LayoutB = typename MapArguments::LayoutB;
+  using ElementC = typename Epilogue::OutputTileIterator::Element;
+  using LayoutC = typename MapArguments::LayoutC;
 
-    using ElementFinalOutput = typename MapArguments::ElementA;
+  using ElementFinalOutput = typename MapArguments::ElementA;
 
-    static ComplexTransform const kTransformA = MapArguments::kTransformA;
-    static ComplexTransform const kTransformB = MapArguments::kTransformB;
+  static ComplexTransform const kTransformA = MapArguments::kTransformA;
+  static ComplexTransform const kTransformB = MapArguments::kTransformB;
 
-    // Type definitions about the mainloop.
-    using Operator = typename Mma::Operator;
-    using OperatorClass = typename Mma::Operator::OperatorClass;
-    using ThreadblockShape = typename Mma::Shape;
-    using WarpShape = typename Mma::Operator::Shape;
-    using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
-    using ArchTag = typename Mma::ArchTag;
+  // Type definitions about the mainloop.
+  using Operator = typename Mma::Operator;
+  using OperatorClass = typename Mma::Operator::OperatorClass;
+  using ThreadblockShape = typename Mma::Shape;
+  using WarpShape = typename Mma::Operator::Shape;
+  using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
+  using ArchTag = typename Mma::ArchTag;
 
-    static int const kStages = Mma::kStages;
-    static int const kAlignmentA = MapArguments::kAlignmentA;
-    static int const kAlignmentB = MapArguments::kAlignmentB;
-    static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess;
+  static int const kStages = Mma::kStages;
+  static int const kAlignmentA = MapArguments::kAlignmentA;
+  static int const kAlignmentB = MapArguments::kAlignmentB;
+  static int const kAlignmentC =
+      Epilogue::OutputTileIterator::kElementsPerAccess;
 
-    /// Warp count (concept: GemmShape)
-    using WarpCount = typename Mma::WarpCount;
-    static int const kThreadCount = 32 * WarpCount::kCount;
+  /// Warp count (concept: GemmShape)
+  using WarpCount = typename Mma::WarpCount;
+  static int const kThreadCount = 32 * WarpCount::kCount;
 
-    using ProblemVisitor
-        = GemmGroupedProblemVisitor<ThreadblockShape, kGroupScheduleMode, kThreadCount, kThreadCount, kTransposed>;
+  using ProblemVisitor = GemmGroupedProblemVisitor<ThreadblockShape,
+                                                   kGroupScheduleMode,
+                                                   kThreadCount,
+                                                   kThreadCount,
+                                                   kTransposed>;
 
+  //
+  // Structures
+  //
+
+  /// Argument structure
+  struct Arguments {
     //
-    // Structures
+    // Data members
     //
 
-    /// Argument structure
-    struct Arguments
-    {
+    GemmCoord* problem_sizes;
+    int problem_count;
+    int threadblock_count;
 
-        //
-        // Data members
-        //
+    typename EpilogueOutputOp::Params output_op;
 
-        GemmCoord* problem_sizes;
-        int problem_count;
-        int threadblock_count;
+    ElementA** ptr_A;
+    ElementB** ptr_B;
+    ElementFinalOutput** ptr_C;
+    ElementFinalOutput** ptr_D;
 
-        typename EpilogueOutputOp::Params output_op;
+    typename LayoutA::Stride::LongIndex* lda;
+    typename LayoutB::Stride::LongIndex* ldb;
+    typename LayoutC::Stride::LongIndex* ldc;
+    typename LayoutC::Stride::LongIndex* ldd;
 
-        ElementA** ptr_A;
-        ElementB** ptr_B;
-        ElementFinalOutput** ptr_C;
-        ElementFinalOutput** ptr_D;
+    // Only used by device-level operator
+    GemmCoord* host_problem_sizes;
 
-        typename LayoutA::Stride::LongIndex* lda;
-        typename LayoutB::Stride::LongIndex* ldb;
-        typename LayoutC::Stride::LongIndex* ldc;
-        typename LayoutC::Stride::LongIndex* ldd;
+    // splitK
+    int split_k_slices;
+    int64_t* splitk_buffer_offsets;
 
-        // Only used by device-level operator
-        GemmCoord* host_problem_sizes;
-
-        // splitK
-        int split_k_slices;
-        int64_t* splitk_buffer_offsets;
-
-        //
-        // Methods
-        //
-
-        /// Default ctor
-        CUTLASS_HOST_DEVICE
-        Arguments()
-            : problem_count(0)
-            , threadblock_count(0)
-            , ptr_A(nullptr)
-            , ptr_B(nullptr)
-            , ptr_C(nullptr)
-            , ptr_D(nullptr)
-            , lda(nullptr)
-            , ldb(nullptr)
-            , ldc(nullptr)
-            , ldd(nullptr)
-            , host_problem_sizes(nullptr)
-            , split_k_slices(1)
-            , splitk_buffer_offsets(nullptr)
-        {
-        }
-
-        /// Ctor
-        CUTLASS_HOST_DEVICE
-        Arguments(GemmCoord* problem_sizes, int problem_count, int threadblock_count,
-            typename EpilogueOutputOp::Params output_op, ElementA** ptr_A, ElementB** ptr_B, ElementFinalOutput** ptr_C,
-            ElementFinalOutput** ptr_D, typename LayoutA::Stride::LongIndex* lda,
-            typename LayoutB::Stride::LongIndex* ldb, typename LayoutC::Stride::LongIndex* ldc,
-            typename LayoutC::Stride::LongIndex* ldd, GemmCoord* host_problem_sizes, int split_k_slices,
-            int64_t* splitk_buffer_offsets)
-            : problem_sizes(problem_sizes)
-            , problem_count(problem_count)
-            , threadblock_count(threadblock_count)
-            , output_op(output_op)
-            , ptr_A(ptr_A)
-            , ptr_B(ptr_B)
-            , ptr_C(ptr_C)
-            , ptr_D(ptr_D)
-            , lda(lda)
-            , ldb(ldb)
-            , ldc(ldc)
-            , ldd(ldd)
-            , host_problem_sizes(host_problem_sizes)
-            , split_k_slices(split_k_slices)
-            , splitk_buffer_offsets(splitk_buffer_offsets)
-        {
-        }
-    };
-
-    //
-    // Structure for precomputing values in host memory and passing to kernels
-    //
-
-    /// Parameters structure
-    struct Params
-    {
-
-        typename ProblemVisitor::Params problem_visitor;
-        int threadblock_count;
-
-        typename EpilogueOutputOp::Params output_op;
-
-        ElementA** ptr_A;
-        ElementB** ptr_B;
-        ElementFinalOutput** ptr_C;
-        ElementFinalOutput** ptr_D;
-        ElementC* ptr_C_split;
-        ElementC* ptr_D_split;
-
-        typename LayoutA::Stride::LongIndex* lda;
-        typename LayoutB::Stride::LongIndex* ldb;
-        typename LayoutC::Stride::LongIndex* ldc;
-        typename LayoutC::Stride::LongIndex* ldd;
-
-        //
-        // Methods
-        //
-
-        // splitk
-        GemmCoord grid_tiled_shape;
-        int swizzle_log_tile;
-        int gemm_k_size;
-        GemmCoord* host_problem_sizes;
-        int split_k_slices;
-        int64_t* splitk_buffer_offsets;
-
-        CUTLASS_HOST_DEVICE
-        Params()
-            : ptr_A(nullptr)
-            , ptr_B(nullptr)
-            , ptr_C(nullptr)
-            , ptr_D(nullptr)
-            , ptr_C_split(nullptr)
-            , ptr_D_split(nullptr)
-            , lda(nullptr)
-            , ldb(nullptr)
-            , ldc(nullptr)
-            , ldd(nullptr)
-            , swizzle_log_tile(0)
-            , gemm_k_size(0)
-            , host_problem_sizes(nullptr)
-            , split_k_slices(1)
-            , splitk_buffer_offsets(nullptr)
-        {
-        }
-
-        CUTLASS_HOST_DEVICE
-        Params(Arguments const& args, void* workspace = nullptr, int tile_count = 0)
-            : problem_visitor(args.problem_sizes, args.problem_count, workspace, tile_count)
-            , host_problem_sizes(args.host_problem_sizes)
-            , threadblock_count(args.threadblock_count)
-            , output_op(args.output_op)
-            , ptr_A(args.ptr_A)
-            , ptr_B(args.ptr_B)
-            , ptr_C(args.ptr_C)
-            , ptr_D(args.ptr_D)
-            , ptr_C_split((ElementC*) workspace)
-            , ptr_D_split((ElementC*) workspace)
-            , lda(args.lda)
-            , ldb(args.ldb)
-            , ldc(args.ldc)
-            , ldd(args.ldd)
-            , split_k_slices(args.split_k_slices)
-            , splitk_buffer_offsets(args.splitk_buffer_offsets)
-        {
-            // Determine grid shape
-            ThreadblockSwizzle threadblock_swizzle;
-            grid_tiled_shape = threadblock_swizzle.get_tiled_shape(args.host_problem_sizes[0],
-                {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, args.split_k_slices);
-            swizzle_log_tile = ThreadblockSwizzle().get_log_tile(grid_tiled_shape);
-
-            // only support same k
-            int full_gemm_k_iterations = args.host_problem_sizes[0].k() / Mma::Shape::kK;
-            int gemm_k_iterations = full_gemm_k_iterations / grid_tiled_shape.k();
-
-            gemm_k_size = gemm_k_iterations * Mma::Shape::kK;
-        }
-
-        CUTLASS_HOST_DEVICE
-        void update(Arguments const& args, void* workspace = nullptr, int tile_count = 0)
-        {
-
-            problem_visitor =
-                typename ProblemVisitor::Params(args.problem_sizes, args.problem_count, workspace, tile_count);
-            threadblock_count = args.threadblock_count;
-            output_op = args.output_op;
-            ptr_A = args.ptr_A;
-            ptr_B = args.ptr_B;
-            ptr_C = args.ptr_C;
-            ptr_D = args.ptr_D;
-            ptr_C_split = workspace;
-            ptr_D_split = workspace;
-
-            lda = args.lda;
-            ldb = args.ldb;
-            ldc = args.ldc;
-            ldd = args.ldd;
-        }
-    };
-
-    /// Shared memory storage structure
-    struct SharedStorage
-    {
-        union
-        {
-            typename Mma::SharedStorage main_loop;
-            typename Epilogue::SharedStorage epilogue;
-        } kernel;
-
-        // ProblemVisitor shared storage can't be overlapped with others
-        typename ProblemVisitor::SharedStorage problem_visitor;
-    };
-
-public:
     //
     // Methods
     //
 
-    CUTLASS_DEVICE
-    SplitkGemmGrouped() {}
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments()
+        : problem_count(0),
+          threadblock_count(0),
+          ptr_A(nullptr),
+          ptr_B(nullptr),
+          ptr_C(nullptr),
+          ptr_D(nullptr),
+          lda(nullptr),
+          ldb(nullptr),
+          ldc(nullptr),
+          ldd(nullptr),
+          host_problem_sizes(nullptr),
+          split_k_slices(1),
+          splitk_buffer_offsets(nullptr) {}
 
-    /// Determines whether kernel satisfies alignment
-    static Status can_implement(cutlass::gemm::GemmCoord const& problem_size)
-    {
-        return Status::kSuccess;
+    /// Ctor
+    CUTLASS_HOST_DEVICE
+    Arguments(GemmCoord* problem_sizes,
+              int problem_count,
+              int threadblock_count,
+              typename EpilogueOutputOp::Params output_op,
+              ElementA** ptr_A,
+              ElementB** ptr_B,
+              ElementFinalOutput** ptr_C,
+              ElementFinalOutput** ptr_D,
+              typename LayoutA::Stride::LongIndex* lda,
+              typename LayoutB::Stride::LongIndex* ldb,
+              typename LayoutC::Stride::LongIndex* ldc,
+              typename LayoutC::Stride::LongIndex* ldd,
+              GemmCoord* host_problem_sizes,
+              int split_k_slices,
+              int64_t* splitk_buffer_offsets)
+        : problem_sizes(problem_sizes),
+          problem_count(problem_count),
+          threadblock_count(threadblock_count),
+          output_op(output_op),
+          ptr_A(ptr_A),
+          ptr_B(ptr_B),
+          ptr_C(ptr_C),
+          ptr_D(ptr_D),
+          lda(lda),
+          ldb(ldb),
+          ldc(ldc),
+          ldd(ldd),
+          host_problem_sizes(host_problem_sizes),
+          split_k_slices(split_k_slices),
+          splitk_buffer_offsets(splitk_buffer_offsets) {}
+  };
+
+  //
+  // Structure for precomputing values in host memory and passing to kernels
+  //
+
+  /// Parameters structure
+  struct Params {
+    typename ProblemVisitor::Params problem_visitor;
+    int threadblock_count;
+
+    typename EpilogueOutputOp::Params output_op;
+
+    ElementA** ptr_A;
+    ElementB** ptr_B;
+    ElementFinalOutput** ptr_C;
+    ElementFinalOutput** ptr_D;
+    ElementC* ptr_C_split;
+    ElementC* ptr_D_split;
+
+    typename LayoutA::Stride::LongIndex* lda;
+    typename LayoutB::Stride::LongIndex* ldb;
+    typename LayoutC::Stride::LongIndex* ldc;
+    typename LayoutC::Stride::LongIndex* ldd;
+
+    //
+    // Methods
+    //
+
+    // splitk
+    GemmCoord grid_tiled_shape;
+    int swizzle_log_tile;
+    int gemm_k_size;
+    GemmCoord* host_problem_sizes;
+    int split_k_slices;
+    int64_t* splitk_buffer_offsets;
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : ptr_A(nullptr),
+          ptr_B(nullptr),
+          ptr_C(nullptr),
+          ptr_D(nullptr),
+          ptr_C_split(nullptr),
+          ptr_D_split(nullptr),
+          lda(nullptr),
+          ldb(nullptr),
+          ldc(nullptr),
+          ldd(nullptr),
+          swizzle_log_tile(0),
+          gemm_k_size(0),
+          host_problem_sizes(nullptr),
+          split_k_slices(1),
+          splitk_buffer_offsets(nullptr) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(Arguments const& args, void* workspace = nullptr, int tile_count = 0)
+        : problem_visitor(
+              args.problem_sizes, args.problem_count, workspace, tile_count),
+          host_problem_sizes(args.host_problem_sizes),
+          threadblock_count(args.threadblock_count),
+          output_op(args.output_op),
+          ptr_A(args.ptr_A),
+          ptr_B(args.ptr_B),
+          ptr_C(args.ptr_C),
+          ptr_D(args.ptr_D),
+          ptr_C_split((ElementC*)workspace),
+          ptr_D_split((ElementC*)workspace),
+          lda(args.lda),
+          ldb(args.ldb),
+          ldc(args.ldc),
+          ldd(args.ldd),
+          split_k_slices(args.split_k_slices),
+          splitk_buffer_offsets(args.splitk_buffer_offsets) {
+      // Determine grid shape
+      ThreadblockSwizzle threadblock_swizzle;
+      grid_tiled_shape = threadblock_swizzle.get_tiled_shape(
+          args.host_problem_sizes[0],
+          {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+          args.split_k_slices);
+      swizzle_log_tile = ThreadblockSwizzle().get_log_tile(grid_tiled_shape);
+
+      // only support same k
+      int full_gemm_k_iterations =
+          args.host_problem_sizes[0].k() / Mma::Shape::kK;
+      int gemm_k_iterations = full_gemm_k_iterations / grid_tiled_shape.k();
+
+      gemm_k_size = gemm_k_iterations * Mma::Shape::kK;
     }
 
-    static Status can_implement(Arguments const& args)
-    {
-        return Status::kSuccess;
+    CUTLASS_HOST_DEVICE
+    void update(Arguments const& args,
+                void* workspace = nullptr,
+                int tile_count = 0) {
+      problem_visitor = typename ProblemVisitor::Params(
+          args.problem_sizes, args.problem_count, workspace, tile_count);
+      threadblock_count = args.threadblock_count;
+      output_op = args.output_op;
+      ptr_A = args.ptr_A;
+      ptr_B = args.ptr_B;
+      ptr_C = args.ptr_C;
+      ptr_D = args.ptr_D;
+      ptr_C_split = workspace;
+      ptr_D_split = workspace;
+
+      lda = args.lda;
+      ldb = args.ldb;
+      ldc = args.ldc;
+      ldd = args.ldd;
     }
+  };
 
-    /// Executes one GEMM
-    CUTLASS_DEVICE
-    void operator()(Params const& params, SharedStorage& shared_storage)
-    {
+  /// Shared memory storage structure
+  struct SharedStorage {
+    union {
+      typename Mma::SharedStorage main_loop;
+      typename Epilogue::SharedStorage epilogue;
+    } kernel;
 
-        //
-        // These types shadow the type-level definitions and support the ability to implement
-        // a 'transposed' GEMM that computes the transposed problems.
-        //
-        using ElementA = typename Mma::IteratorA::Element;
-        using LayoutA = typename Mma::IteratorA::Layout;
-        using ElementB = typename Mma::IteratorB::Element;
-        using LayoutB = typename Mma::IteratorB::Layout;
-        using ElementC = typename Epilogue::OutputTileIterator::Element;
-        using LayoutC = typename Epilogue::OutputTileIterator::Layout;
+    // ProblemVisitor shared storage can't be overlapped with others
+    typename ProblemVisitor::SharedStorage problem_visitor;
+  };
 
-        //
-        // Problem visitor.
-        //
-        ProblemVisitor problem_visitor(params.problem_visitor, shared_storage.problem_visitor, blockIdx.x);
+ public:
+  //
+  // Methods
+  //
 
-        // Outer 'persistent' loop to iterate over tiles
-        while (problem_visitor.next_tile())
-        {
+  CUTLASS_DEVICE
+  SplitkGemmGrouped() {}
 
-            GemmCoord problem_size = problem_visitor.problem_size();
-            int32_t problem_idx = problem_visitor.problem_index();
-            int32_t threadblock_idx = int32_t(problem_visitor.threadblock_idx());
+  /// Determines whether kernel satisfies alignment
+  static Status can_implement(cutlass::gemm::GemmCoord const& problem_size) {
+    return Status::kSuccess;
+  }
 
-            GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
+  static Status can_implement(Arguments const& args) {
+    return Status::kSuccess;
+  }
 
-            // Load element pointers. Exchange pointers and strides if working on the transpose
-            ElementA* ptr_A
-                = reinterpret_cast<ElementA*>((kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx]));
-            typename LayoutA::LongIndex ldm_A = (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]);
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const& params, SharedStorage& shared_storage) {
+    //
+    // These types shadow the type-level definitions and support the ability to
+    // implement a 'transposed' GEMM that computes the transposed problems.
+    //
+    using ElementA = typename Mma::IteratorA::Element;
+    using LayoutA = typename Mma::IteratorA::Layout;
+    using ElementB = typename Mma::IteratorB::Element;
+    using LayoutB = typename Mma::IteratorB::Layout;
+    using ElementC = typename Epilogue::OutputTileIterator::Element;
+    using LayoutC = typename Epilogue::OutputTileIterator::Layout;
 
-            ElementB* ptr_B
-                = reinterpret_cast<ElementB*>((kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx]));
-            typename LayoutB::LongIndex ldm_B = (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]);
+    //
+    // Problem visitor.
+    //
+    ProblemVisitor problem_visitor(
+        params.problem_visitor, shared_storage.problem_visitor, blockIdx.x);
 
-            // Compute threadblock location
-            ThreadblockSwizzle threadblock_swizzle;
-            GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+    // Outer 'persistent' loop to iterate over tiles
+    while (problem_visitor.next_tile()) {
+      GemmCoord problem_size = problem_visitor.problem_size();
+      int32_t problem_idx = problem_visitor.problem_index();
+      int32_t threadblock_idx = int32_t(problem_visitor.threadblock_idx());
 
-            cutlass::gemm::GemmCoord threadblock_offset(int(threadblock_idx / grid_shape.n()) * Mma::Shape::kM,
-                int(threadblock_idx % grid_shape.n()) * Mma::Shape::kN, 0);
+      GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
 
-            // Compute initial location in logical coordinates
-            cutlass::MatrixCoord tb_offset_A{
-                threadblock_offset.m(),
-                threadblock_tile_offset.k() * params.gemm_k_size,
-            };
+      // Load element pointers. Exchange pointers and strides if working on the
+      // transpose
+      ElementA* ptr_A = reinterpret_cast<ElementA*>((
+          kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx]));
+      typename LayoutA::LongIndex ldm_A =
+          (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]);
 
-            cutlass::MatrixCoord tb_offset_B{threadblock_tile_offset.k() * params.gemm_k_size, threadblock_offset.n()};
+      ElementB* ptr_B = reinterpret_cast<ElementB*>((
+          kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx]));
+      typename LayoutB::LongIndex ldm_B =
+          (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]);
 
-            // Problem size is a function of threadblock index in the K dimension
-            int problem_size_k;
-            if (threadblock_tile_offset.k() + 1 == params.grid_tiled_shape.k())
-            {
-                problem_size_k = problem_size.k();
-            }
-            else
-            {
-                problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
-            }
+      // Compute threadblock location
+      ThreadblockSwizzle threadblock_swizzle;
+      GemmCoord threadblock_tile_offset =
+          threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-            // Compute threadblock-scoped matrix multiply-add
-            int gemm_k_iterations = (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) / Mma::Shape::kK;
+      cutlass::gemm::GemmCoord threadblock_offset(
+          int(threadblock_idx / grid_shape.n()) * Mma::Shape::kM,
+          int(threadblock_idx % grid_shape.n()) * Mma::Shape::kN,
+          0);
 
-            // Compute position within threadblock
-            int thread_idx = threadIdx.x;
+      // Compute initial location in logical coordinates
+      cutlass::MatrixCoord tb_offset_A{
+          threadblock_offset.m(),
+          threadblock_tile_offset.k() * params.gemm_k_size,
+      };
 
-            // Construct iterators to A and B operands
-            typename Mma::IteratorA iterator_A(
-                LayoutA(ldm_A), ptr_A, {problem_size.m(), problem_size_k}, thread_idx, tb_offset_A);
+      cutlass::MatrixCoord tb_offset_B{
+          threadblock_tile_offset.k() * params.gemm_k_size,
+          threadblock_offset.n()};
 
-            typename Mma::IteratorB iterator_B(
-                LayoutB(ldm_B), ptr_B, {problem_size_k, problem_size.n()}, thread_idx, tb_offset_B);
+      // Problem size is a function of threadblock index in the K dimension
+      int problem_size_k;
+      if (threadblock_tile_offset.k() + 1 == params.grid_tiled_shape.k()) {
+        problem_size_k = problem_size.k();
+      } else {
+        problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
+      }
 
-            typename Mma::FragmentC accumulators;
+      // Compute threadblock-scoped matrix multiply-add
+      int gemm_k_iterations =
+          (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) /
+          Mma::Shape::kK;
 
-            accumulators.clear();
+      // Compute position within threadblock
+      int thread_idx = threadIdx.x;
 
-            // Broadcast the warp_id computed by lane 0 to ensure dependent code
-            // is compiled as warp-uniform.
-            int warp_idx = canonical_warp_idx_sync();
+      // Construct iterators to A and B operands
+      typename Mma::IteratorA iterator_A(LayoutA(ldm_A),
+                                         ptr_A,
+                                         {problem_size.m(), problem_size_k},
+                                         thread_idx,
+                                         tb_offset_A);
 
-            int lane_idx = threadIdx.x % 32;
+      typename Mma::IteratorB iterator_B(LayoutB(ldm_B),
+                                         ptr_B,
+                                         {problem_size_k, problem_size.n()},
+                                         thread_idx,
+                                         tb_offset_B);
 
-            //
-            // Matrix multiply phase
-            //
+      typename Mma::FragmentC accumulators;
 
-            // Construct thread-scoped matrix multiply
-            Mma mma(shared_storage.kernel.main_loop, thread_idx, warp_idx, lane_idx);
+      accumulators.clear();
 
-            // Wait for all threads to finish their epilogue phases from the previous tile.
-            __syncthreads();
+      // Broadcast the warp_id computed by lane 0 to ensure dependent code
+      // is compiled as warp-uniform.
+      int warp_idx = canonical_warp_idx_sync();
 
-            // Compute threadblock-scoped matrix multiply-add
-            mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators);
+      int lane_idx = threadIdx.x % 32;
 
-            //
-            // Epilogue
-            //
+      //
+      // Matrix multiply phase
+      //
 
-            EpilogueOutputOp output_op(params.output_op);
+      // Construct thread-scoped matrix multiply
+      Mma mma(shared_storage.kernel.main_loop, thread_idx, warp_idx, lane_idx);
 
-            ElementC* ptr_C = params.ptr_C_split;
-            ElementC* ptr_D = params.ptr_D_split;
+      // Wait for all threads to finish their epilogue phases from the previous
+      // tile.
+      __syncthreads();
 
-            LayoutC layout_C(params.ldc[problem_idx]);
-            LayoutC layout_D(params.ldd[problem_idx]);
+      // Compute threadblock-scoped matrix multiply-add
+      mma(gemm_k_iterations,
+          accumulators,
+          iterator_A,
+          iterator_B,
+          accumulators);
 
-            typename Epilogue::OutputTileIterator::Params params_C(layout_C);
-            typename Epilogue::OutputTileIterator::Params params_D(layout_D);
+      //
+      // Epilogue
+      //
 
-            // assume identity swizzle
-            MatrixCoord threadblock_offset_C(threadblock_offset.m(), threadblock_offset.n());
+      EpilogueOutputOp output_op(params.output_op);
 
-            // Tile iterator loading from source tensor.
-            typename Epilogue::OutputTileIterator iterator_C(
-                params_C, ptr_C, problem_size.mn(), thread_idx, threadblock_offset_C);
+      ElementC* ptr_C = params.ptr_C_split;
+      ElementC* ptr_D = params.ptr_D_split;
 
-            iterator_C.add_pointer_offset(problem_size.m() * problem_size.n() * threadblock_tile_offset.k()
-                + gridDim.z * params.splitk_buffer_offsets[problem_idx]);
+      LayoutC layout_C(params.ldc[problem_idx]);
+      LayoutC layout_D(params.ldd[problem_idx]);
 
-            // Tile iterator writing to destination tensor.
-            typename Epilogue::OutputTileIterator iterator_D(
-                params_D, ptr_D, problem_size.mn(), thread_idx, threadblock_offset_C);
-            iterator_D.add_pointer_offset(problem_size.m() * problem_size.n() * threadblock_tile_offset.k()
-                + gridDim.z * params.splitk_buffer_offsets[problem_idx]);
+      typename Epilogue::OutputTileIterator::Params params_C(layout_C);
+      typename Epilogue::OutputTileIterator::Params params_D(layout_D);
 
-            Epilogue epilogue(shared_storage.kernel.epilogue, thread_idx, warp_idx, lane_idx);
+      // assume identity swizzle
+      MatrixCoord threadblock_offset_C(threadblock_offset.m(),
+                                       threadblock_offset.n());
 
-            // Execute the epilogue operator to update the destination tensor.
-            epilogue(output_op, iterator_D, accumulators, iterator_C);
+      // Tile iterator loading from source tensor.
+      typename Epilogue::OutputTileIterator iterator_C(
+          params_C, ptr_C, problem_size.mn(), thread_idx, threadblock_offset_C);
 
-            // Next tile
-            problem_visitor.advance(gridDim.x);
-        }
+      iterator_C.add_pointer_offset(
+          problem_size.m() * problem_size.n() * threadblock_tile_offset.k() +
+          gridDim.z * params.splitk_buffer_offsets[problem_idx]);
+
+      // Tile iterator writing to destination tensor.
+      typename Epilogue::OutputTileIterator iterator_D(
+          params_D, ptr_D, problem_size.mn(), thread_idx, threadblock_offset_C);
+      iterator_D.add_pointer_offset(
+          problem_size.m() * problem_size.n() * threadblock_tile_offset.k() +
+          gridDim.z * params.splitk_buffer_offsets[problem_idx]);
+
+      Epilogue epilogue(
+          shared_storage.kernel.epilogue, thread_idx, warp_idx, lane_idx);
+
+      // Execute the epilogue operator to update the destination tensor.
+      epilogue(output_op, iterator_D, accumulators, iterator_C);
+
+      // Next tile
+      problem_visitor.advance(gridDim.x);
     }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace kernel
-} // namespace gemm
-} // namespace cutlass
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma.h
index ed5e3e4daf..a268861a0d 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,18 +19,15 @@
 #include "cutlass_extensions/arch/mma.h"
 #include "cutlass_extensions/interleaved_numeric_conversion.h"
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 ////////////////////////////////////////////////////////////////////////////////
 
-// We need to distinguish here, since we want volta support. It is too much effort
-// to write shared memory iterators that are probably needed for volta to function
-// properly. As a result, we allow converters both after the LDG (for volta) and after
-// the LDS for Turing+.
+// We need to distinguish here, since we want volta support. It is too much
+// effort to write shared memory iterators that are probably needed for volta to
+// function properly. As a result, we allow converters both after the LDG (for
+// volta) and after the LDS for Turing+.
 template <
     /// Iterator for B matrix in global memory
     typename IteratorB,
@@ -38,9 +35,7 @@ template <
     typename MmaOperator,
     /// Math operation perform by warp level operator
     typename MathOperator>
-struct SetConverters
-{
-};
+struct SetConverters {};
 
 // Dequantize after LDG, so set transforms accordingly
 template <
@@ -48,14 +43,16 @@ template <
     typename IteratorB,
     /// Mma Policy
     typename MmaOperator>
-struct SetConverters<IteratorB, MmaOperator, arch::OpMultiplyAdd>
-{
-    using TransformAfterLDG
-        = FastInterleavedAndBiasedNumericArrayConverter<typename MmaOperator::ArchMmaOperator::ElementB,
-            typename IteratorB::Element, IteratorB::Fragment::kElements>;
+struct SetConverters<IteratorB, MmaOperator, arch::OpMultiplyAdd> {
+  using TransformAfterLDG = FastInterleavedAndBiasedNumericArrayConverter<
+      typename MmaOperator::ArchMmaOperator::ElementB,
+      typename IteratorB::Element,
+      IteratorB::Fragment::kElements>;
 
-    using TransformAfterLDS = NumericArrayConverter<typename MmaOperator::ArchMmaOperator::ElementB,
-        typename MmaOperator::ArchMmaOperator::ElementB, MmaOperator::FragmentB::kElements>;
+  using TransformAfterLDS =
+      NumericArrayConverter<typename MmaOperator::ArchMmaOperator::ElementB,
+                            typename MmaOperator::ArchMmaOperator::ElementB,
+                            MmaOperator::FragmentB::kElements>;
 };
 
 // Dequantize after LDS, so set transforms accordingly
@@ -65,14 +62,18 @@ template <
     typename IteratorB,
     /// Mma Policy
     typename MmaOperator>
-struct SetConverters<IteratorB, MmaOperator, arch::OpMultiplyAddDequantizeInterleavedBToA>
-{
-    using TransformAfterLDG = NumericArrayConverter<typename IteratorB::Element, typename IteratorB::Element,
-        IteratorB::Fragment::kElements>;
+struct SetConverters<IteratorB,
+                     MmaOperator,
+                     arch::OpMultiplyAddDequantizeInterleavedBToA> {
+  using TransformAfterLDG =
+      NumericArrayConverter<typename IteratorB::Element,
+                            typename IteratorB::Element,
+                            IteratorB::Fragment::kElements>;
 
-    using TransformAfterLDS
-        = FastInterleavedAndBiasedNumericArrayConverter<typename MmaOperator::ArchMmaOperator::ElementB,
-            typename TransformAfterLDG::result_type::Element, MmaOperator::FragmentB::kElements>;
+  using TransformAfterLDS = FastInterleavedAndBiasedNumericArrayConverter<
+      typename MmaOperator::ArchMmaOperator::ElementB,
+      typename TransformAfterLDG::result_type::Element,
+      MmaOperator::FragmentB::kElements>;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -120,6 +121,6 @@ template <
     typename Enable = void>
 struct DqMma;
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h
index 17c6346553..566cd379b2 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,49 +27,77 @@
 #include "cutlass_extensions/gemm/threadblock/default_dq_mma.h"
 #include "cutlass_extensions/transform/threadblock/fine_grained_scale_zero_iterator.h"
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template <typename MmaShape, typename Element, typename Layout, WeightOnlyQuantOp QuantOp, int Alignment,
-    typename Enable = void>
+template <typename MmaShape,
+          typename Element,
+          typename Layout,
+          WeightOnlyQuantOp QuantOp,
+          int Alignment,
+          typename Enable = void>
 struct DefaultScaleIteratorsMultistage;
 
 // Fine grained iterators
-template <typename MmaShape, typename Element, typename Layout, WeightOnlyQuantOp QuantOp, int Alignment>
-struct DefaultScaleIteratorsMultistage<MmaShape, Element, Layout, QuantOp, Alignment,
-    std::enable_if_t<isFinegrained(QuantOp)>>
-{
-    using IteratorScale
-        = cutlass::transform::threadblock::FineGrainedScaleZeroIterator<cutlass::MatrixShape<1, MmaShape::kN>, Element,
-            Layout, 0, Alignment>;
+template <typename MmaShape,
+          typename Element,
+          typename Layout,
+          WeightOnlyQuantOp QuantOp,
+          int Alignment>
+struct DefaultScaleIteratorsMultistage<
+    MmaShape,
+    Element,
+    Layout,
+    QuantOp,
+    Alignment,
+    std::enable_if_t<isFinegrained(QuantOp)>> {
+  using IteratorScale =
+      cutlass::transform::threadblock::FineGrainedScaleZeroIterator<
+          cutlass::MatrixShape<1, MmaShape::kN>,
+          Element,
+          Layout,
+          0,
+          Alignment>;
 
-    using SmemIteratorScale = IteratorScale;
+  using SmemIteratorScale = IteratorScale;
 };
 
 // Per column iterators
-template <typename MmaShape, typename Element, typename Layout, WeightOnlyQuantOp QuantOp, int Alignment>
-struct DefaultScaleIteratorsMultistage<MmaShape, Element, Layout, QuantOp, Alignment,
-    std::enable_if_t<!isFinegrained(QuantOp)>>
-{
-    // ThreadMap for scale iterator
-    static_assert((MmaShape::kN % Alignment) == 0, "");
+template <typename MmaShape,
+          typename Element,
+          typename Layout,
+          WeightOnlyQuantOp QuantOp,
+          int Alignment>
+struct DefaultScaleIteratorsMultistage<
+    MmaShape,
+    Element,
+    Layout,
+    QuantOp,
+    Alignment,
+    std::enable_if_t<!isFinegrained(QuantOp)>> {
+  // ThreadMap for scale iterator
+  static_assert((MmaShape::kN % Alignment) == 0, "");
 
-private:
-    using IteratorScaleThreadMap = transform::PitchLinearStripminedThreadMap<layout::PitchLinearShape<MmaShape::kN, 1>,
-        MmaShape::kN / Alignment, Alignment>;
+ private:
+  using IteratorScaleThreadMap = transform::PitchLinearStripminedThreadMap<
+      layout::PitchLinearShape<MmaShape::kN, 1>,
+      MmaShape::kN / Alignment,
+      Alignment>;
 
-public:
-    // Define iterators over tiles from the scale operand
-    using IteratorScale = cutlass::transform::threadblock::PredicatedTileIterator<cutlass::MatrixShape<1, MmaShape::kN>,
-        Element, Layout, 0, IteratorScaleThreadMap, Alignment>;
+ public:
+  // Define iterators over tiles from the scale operand
+  using IteratorScale = cutlass::transform::threadblock::PredicatedTileIterator<
+      cutlass::MatrixShape<1, MmaShape::kN>,
+      Element,
+      Layout,
+      0,
+      IteratorScaleThreadMap,
+      Alignment>;
 
-    using SmemIteratorScale = IteratorScale;
+  using SmemIteratorScale = IteratorScale;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -111,69 +139,133 @@ template <
     typename Operator_,
     /// Use zfill or predicate for out-of-bound cp.async
     SharedMemoryClearOption SharedMemoryClear>
-struct DqMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementScale, LayoutScale, kAlignmentScale,
-    ElementAccumulator, layout::RowMajor, OperatorClass, ArchTag, ThreadblockShape, WarpShape, InstructionShape,
-    kStages, Operator_, SharedMemoryClear,
-    typename platform::enable_if<(
-        ArchTag::kMinComputeCapability >= 80 && !layout::IsColumnMajorTileInterleave<LayoutB>::value)>::type>
-{
+struct DqMma<ElementA,
+             LayoutA,
+             kAlignmentA,
+             ElementB,
+             LayoutB,
+             kAlignmentB,
+             ElementScale,
+             LayoutScale,
+             kAlignmentScale,
+             ElementAccumulator,
+             layout::RowMajor,
+             OperatorClass,
+             ArchTag,
+             ThreadblockShape,
+             WarpShape,
+             InstructionShape,
+             kStages,
+             Operator_,
+             SharedMemoryClear,
+             typename platform::enable_if<(
+                 ArchTag::kMinComputeCapability >= 80 &&
+                 !layout::IsColumnMajorTileInterleave<LayoutB>::value)>::type> {
+  static_assert(platform::is_same<ElementA, half_t>::value ||
+                    platform::is_same<ElementA, bfloat16_t>::value ||
+                    platform::is_same<ElementA, float_e4m3_t>::value,
+                "Element A must be fp16, fp8 or bf16");
 
-    static_assert(platform::is_same<ElementA, half_t>::value || platform::is_same<ElementA, bfloat16_t>::value
-            || platform::is_same<ElementA, float_e4m3_t>::value,
-        "Element A must be fp16, fp8 or bf16");
+  using OperatorInfo = arch::DetagOperator<Operator_>;
+  using Operator = typename OperatorInfo::Operator;
+  static_assert(
+      platform::is_same<Operator,
+                        arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
+      "Mma multistage must dequantize after ldsm");
 
-    using OperatorInfo = arch::DetagOperator<Operator_>;
-    using Operator = typename OperatorInfo::Operator;
-    static_assert(platform::is_same<Operator, arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
-        "Mma multistage must dequantize after ldsm");
+  static_assert(platform::is_same<ElementB, uint8_t>::value ||
+                    platform::is_same<ElementB, uint4b_t>::value,
+                "Element B must be uint8 or uint4");
 
-    static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-        "Element B must be uint8 or uint4");
+  static cutlass::arch::CacheOperation::Kind const CacheOpA =
+      ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
-        ? cutlass::arch::CacheOperation::Global
-        : cutlass::arch::CacheOperation::Always;
+  static cutlass::arch::CacheOperation::Kind const CacheOpB =
+      ((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpB = ((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
-        ? cutlass::arch::CacheOperation::Global
-        : cutlass::arch::CacheOperation::Always;
+  // Define the MmaCore components
+  // Mma core does not depend on stages, so pass in at least 3 here to mma
+  // multistage pieces are created
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          ElementA,
+                                                          LayoutA,
+                                                          ElementB,
+                                                          LayoutB,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          OperatorClass,
+                                                          std::max(kStages, 3),
+                                                          Operator,
+                                                          false,
+                                                          CacheOpA,
+                                                          CacheOpB>;
 
-    // Define the MmaCore components
-    // Mma core does not depend on stages, so pass in at least 3 here to mma multistage pieces are created
-    using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape, WarpShape, InstructionShape,
-        ElementA, LayoutA, ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, std::max(kStages, 3),
-        Operator, false, CacheOpA, CacheOpB>;
+  // Define iterators over tiles from the A operand
+  using ThreadMapA = typename MmaCore::IteratorThreadMapA;
+  using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
+  using IteratorA =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
+          ElementA,
+          LayoutA,
+          1,
+          ThreadMapA,
+          AccessTypeA>;
 
-    // Define iterators over tiles from the A operand
-    using ThreadMapA = typename MmaCore::IteratorThreadMapA;
-    using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
-    using IteratorA = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>, ElementA, LayoutA, 1, ThreadMapA,
-        AccessTypeA>;
+  // Define iterators over tiles from the B operand
+  using ThreadMapB = typename MmaCore::IteratorThreadMapB;
+  using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
+  using IteratorB =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>,
+          ElementB,
+          LayoutB,
+          0,
+          ThreadMapB,
+          AccessTypeB>;
 
-    // Define iterators over tiles from the B operand
-    using ThreadMapB = typename MmaCore::IteratorThreadMapB;
-    using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
-    using IteratorB = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>, ElementB, LayoutB, 0, ThreadMapB,
-        AccessTypeB>;
+  using ScaleIterators =
+      DefaultScaleIteratorsMultistage<typename MmaCore::Shape,
+                                      ElementScale,
+                                      LayoutScale,
+                                      OperatorInfo::QuantOp,
+                                      kAlignmentScale>;
 
-    using ScaleIterators = DefaultScaleIteratorsMultistage<typename MmaCore::Shape, ElementScale, LayoutScale,
-        OperatorInfo::QuantOp, kAlignmentScale>;
+  // Define iterators over tiles from the scale operand
+  using IteratorScale = typename ScaleIterators::IteratorScale;
 
-    // Define iterators over tiles from the scale operand
-    using IteratorScale = typename ScaleIterators::IteratorScale;
+  using SmemIteratorScale = typename ScaleIterators::SmemIteratorScale;
 
-    using SmemIteratorScale = typename ScaleIterators::SmemIteratorScale;
+  using Converter = FastInterleavedAndBiasedNumericArrayConverter<
+      ElementScale,
+      ElementB,
+      MmaCore::MmaPolicy::Operator::FragmentB::kElements>;
 
-    using Converter = FastInterleavedAndBiasedNumericArrayConverter<ElementScale, ElementB,
-        MmaCore::MmaPolicy::Operator::FragmentB::kElements>;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::DqMmaMultistage<typename MmaCore::Shape, IteratorA,
-        typename MmaCore::SmemIteratorA, MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB,
-        MmaCore::kCacheOpB, IteratorScale, SmemIteratorScale, ElementAccumulator, layout::RowMajor,
-        typename MmaCore::MmaPolicy, kStages, Converter, OperatorInfo::QuantOp, SharedMemoryClear>;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = cutlass::gemm::threadblock::DqMmaMultistage<
+      typename MmaCore::Shape,
+      IteratorA,
+      typename MmaCore::SmemIteratorA,
+      MmaCore::kCacheOpA,
+      IteratorB,
+      typename MmaCore::SmemIteratorB,
+      MmaCore::kCacheOpB,
+      IteratorScale,
+      SmemIteratorScale,
+      ElementAccumulator,
+      layout::RowMajor,
+      typename MmaCore::MmaPolicy,
+      kStages,
+      Converter,
+      OperatorInfo::QuantOp,
+      SharedMemoryClear>;
 };
 
 // Specialization to handle column major interleave B
@@ -214,89 +306,159 @@ template <
     typename Operator_,
     /// Use zfill or predicate for out-of-bound cp.async
     SharedMemoryClearOption SharedMemoryClear>
-struct DqMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementScale, LayoutScale, kAlignmentScale,
-    ElementAccumulator, layout::RowMajor, OperatorClass, ArchTag, ThreadblockShape, WarpShape, InstructionShape,
-    kStages, Operator_, SharedMemoryClear,
-    typename platform::enable_if<(
-        ArchTag::kMinComputeCapability >= 80 && layout::IsColumnMajorTileInterleave<LayoutB>::value)>::type>
-{
+struct DqMma<ElementA,
+             LayoutA,
+             kAlignmentA,
+             ElementB,
+             LayoutB,
+             kAlignmentB,
+             ElementScale,
+             LayoutScale,
+             kAlignmentScale,
+             ElementAccumulator,
+             layout::RowMajor,
+             OperatorClass,
+             ArchTag,
+             ThreadblockShape,
+             WarpShape,
+             InstructionShape,
+             kStages,
+             Operator_,
+             SharedMemoryClear,
+             typename platform::enable_if<(
+                 ArchTag::kMinComputeCapability >= 80 &&
+                 layout::IsColumnMajorTileInterleave<LayoutB>::value)>::type> {
+  static_assert(platform::is_same<ElementA, half_t>::value ||
+                    platform::is_same<ElementA, bfloat16_t>::value ||
+                    platform::is_same<ElementA, float_e4m3_t>::value,
+                "Element A must be fp16, fp8 or bf16");
 
-    static_assert(platform::is_same<ElementA, half_t>::value || platform::is_same<ElementA, bfloat16_t>::value
-            || platform::is_same<ElementA, float_e4m3_t>::value,
-        "Element A must be fp16, fp8 or bf16");
+  using OperatorInfo = arch::DetagOperator<Operator_>;
+  using Operator = typename OperatorInfo::Operator;
+  static_assert(
+      platform::is_same<Operator,
+                        arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
+      "Mma multistage must dequantize after ldsm");
 
-    using OperatorInfo = arch::DetagOperator<Operator_>;
-    using Operator = typename OperatorInfo::Operator;
-    static_assert(platform::is_same<Operator, arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
-        "Mma multistage must dequantize after ldsm");
+  static_assert(platform::is_same<ElementB, uint8_t>::value ||
+                    platform::is_same<ElementB, uint4b_t>::value,
+                "Element B must be uint8 or uint4");
 
-    static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-        "Element B must be uint8 or uint4");
+  static cutlass::arch::CacheOperation::Kind const CacheOpA =
+      ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
-        ? cutlass::arch::CacheOperation::Global
-        : cutlass::arch::CacheOperation::Always;
+  static cutlass::arch::CacheOperation::Kind const CacheOpB =
+      ((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpB = ((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
-        ? cutlass::arch::CacheOperation::Global
-        : cutlass::arch::CacheOperation::Always;
+  // Define the MmaCore components
+  // Mma core does not depend on stages, so pass in at least 3 here to mma
+  // multistage pieces are created
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          ElementA,
+                                                          LayoutA,
+                                                          ElementB,
+                                                          layout::ColumnMajor,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          OperatorClass,
+                                                          std::max(kStages, 3),
+                                                          Operator,
+                                                          false,
+                                                          CacheOpA,
+                                                          CacheOpB>;
 
-    // Define the MmaCore components
-    // Mma core does not depend on stages, so pass in at least 3 here to mma multistage pieces are created
-    using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape, WarpShape, InstructionShape,
-        ElementA, LayoutA, ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, OperatorClass,
-        std::max(kStages, 3), Operator, false, CacheOpA, CacheOpB>;
+  // Define iterators over tiles from the A operand
+  using ThreadMapA = typename MmaCore::IteratorThreadMapA;
+  using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
+  using IteratorA =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
+          ElementA,
+          LayoutA,
+          1,
+          ThreadMapA,
+          AccessTypeA>;
 
-    // Define iterators over tiles from the A operand
-    using ThreadMapA = typename MmaCore::IteratorThreadMapA;
-    using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
-    using IteratorA = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>, ElementA, LayoutA, 1, ThreadMapA,
-        AccessTypeA>;
+ private:
+  static constexpr int ColumnsInterleaved = LayoutB::kColumnsInterleaved;
+  static constexpr int RowsPerTile = LayoutB::kRowsPerTile;
+  static_assert(!(MmaCore::Shape::kN % ColumnsInterleaved), "");
+  static_assert(RowsPerTile == MmaCore::Shape::kK, "");
 
-private:
-    static constexpr int ColumnsInterleaved = LayoutB::kColumnsInterleaved;
-    static constexpr int RowsPerTile = LayoutB::kRowsPerTile;
-    static_assert(!(MmaCore::Shape::kN % ColumnsInterleaved), "");
-    static_assert(RowsPerTile == MmaCore::Shape::kK, "");
+  using OriginalThreadMap = typename MmaCore::IteratorThreadMapB;
+  using OriginalWarpArrangement =
+      typename OriginalThreadMap::Detail::WarpThreadArrangement;
+  static_assert(!(OriginalWarpArrangement::kStrided % ColumnsInterleaved), "");
 
-    using OriginalThreadMap = typename MmaCore::IteratorThreadMapB;
-    using OriginalWarpArrangement = typename OriginalThreadMap::Detail::WarpThreadArrangement;
-    static_assert(!(OriginalWarpArrangement::kStrided % ColumnsInterleaved), "");
+  using GmemIteratorShape =
+      MatrixShape<MmaCore::Shape::kK * ColumnsInterleaved,
+                  MmaCore::Shape::kN / ColumnsInterleaved>;
+  using GmemThreadMapB = transform::PitchLinearWarpRakedThreadMap<
+      layout::PitchLinearShape<GmemIteratorShape::kRow,
+                               GmemIteratorShape::kColumn>,
+      OriginalThreadMap::kThreads,
+      layout::PitchLinearShape<
+          OriginalWarpArrangement::kContiguous * ColumnsInterleaved,
+          OriginalWarpArrangement::kStrided / ColumnsInterleaved>,
+      MmaCore::kAccessSizeInBits / sizeof_bits<ElementB>::value>;
 
-    using GmemIteratorShape
-        = MatrixShape<MmaCore::Shape::kK * ColumnsInterleaved, MmaCore::Shape::kN / ColumnsInterleaved>;
-    using GmemThreadMapB = transform::PitchLinearWarpRakedThreadMap<
-        layout::PitchLinearShape<GmemIteratorShape::kRow, GmemIteratorShape::kColumn>, OriginalThreadMap::kThreads,
-        layout::PitchLinearShape<OriginalWarpArrangement::kContiguous * ColumnsInterleaved,
-            OriginalWarpArrangement::kStrided / ColumnsInterleaved>,
-        MmaCore::kAccessSizeInBits / sizeof_bits<ElementB>::value>;
+ public:
+  // Define iterators over tiles from the B operand
+  using ThreadMapB = typename MmaCore::IteratorThreadMapB;
+  using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
+  using IteratorB =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          GmemIteratorShape,
+          ElementB,
+          layout::ColumnMajor,
+          0,
+          GmemThreadMapB,
+          AccessTypeB>;
 
-public:
-    // Define iterators over tiles from the B operand
-    using ThreadMapB = typename MmaCore::IteratorThreadMapB;
-    using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
-    using IteratorB = cutlass::transform::threadblock::PredicatedTileAccessIterator<GmemIteratorShape, ElementB,
-        layout::ColumnMajor, 0, GmemThreadMapB, AccessTypeB>;
+  using ScaleIterators =
+      DefaultScaleIteratorsMultistage<typename MmaCore::Shape,
+                                      ElementScale,
+                                      LayoutScale,
+                                      OperatorInfo::QuantOp,
+                                      kAlignmentScale>;
 
-    using ScaleIterators = DefaultScaleIteratorsMultistage<typename MmaCore::Shape, ElementScale, LayoutScale,
-        OperatorInfo::QuantOp, kAlignmentScale>;
+  // Define iterators over tiles from the scale operand
+  using IteratorScale = typename ScaleIterators::IteratorScale;
 
-    // Define iterators over tiles from the scale operand
-    using IteratorScale = typename ScaleIterators::IteratorScale;
+  using SmemIteratorScale = typename ScaleIterators::SmemIteratorScale;
 
-    using SmemIteratorScale = typename ScaleIterators::SmemIteratorScale;
+  using Converter = FastInterleavedAndBiasedNumericArrayConverter<
+      ElementScale,
+      ElementB,
+      MmaCore::MmaPolicy::Operator::FragmentB::kElements>;
 
-    using Converter = FastInterleavedAndBiasedNumericArrayConverter<ElementScale, ElementB,
-        MmaCore::MmaPolicy::Operator::FragmentB::kElements>;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::DqMmaMultistage<typename MmaCore::Shape, IteratorA,
-        typename MmaCore::SmemIteratorA, MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB,
-        MmaCore::kCacheOpB, IteratorScale, SmemIteratorScale, ElementAccumulator, layout::RowMajor,
-        typename MmaCore::MmaPolicy, kStages, Converter, OperatorInfo::QuantOp, SharedMemoryClear>;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = cutlass::gemm::threadblock::DqMmaMultistage<
+      typename MmaCore::Shape,
+      IteratorA,
+      typename MmaCore::SmemIteratorA,
+      MmaCore::kCacheOpA,
+      IteratorB,
+      typename MmaCore::SmemIteratorB,
+      MmaCore::kCacheOpB,
+      IteratorScale,
+      SmemIteratorScale,
+      ElementAccumulator,
+      layout::RowMajor,
+      typename MmaCore::MmaPolicy,
+      kStages,
+      Converter,
+      OperatorInfo::QuantOp,
+      SharedMemoryClear>;
 };
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h
index 345cd2eec9..ba7f2863e0 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,58 +27,95 @@
 #include "cutlass_extensions/gemm/threadblock/default_dq_mma.h"
 #include "cutlass_extensions/transform/threadblock/fine_grained_scale_zero_iterator.h"
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template <typename MmaShape, typename Element, typename Layout, WeightOnlyQuantOp QuantOp, int Alignment,
-    typename Enable = void>
+template <typename MmaShape,
+          typename Element,
+          typename Layout,
+          WeightOnlyQuantOp QuantOp,
+          int Alignment,
+          typename Enable = void>
 struct DefaultScaleIteratorsPipelined;
 
 // Fine grained iterators
-template <typename MmaShape, typename Element, typename Layout, WeightOnlyQuantOp QuantOp, int Alignment>
-struct DefaultScaleIteratorsPipelined<MmaShape, Element, Layout, QuantOp, Alignment,
-    std::enable_if_t<isFinegrained(QuantOp)>>
-{
-private:
-    using SmemScaleType = half_t;
+template <typename MmaShape,
+          typename Element,
+          typename Layout,
+          WeightOnlyQuantOp QuantOp,
+          int Alignment>
+struct DefaultScaleIteratorsPipelined<
+    MmaShape,
+    Element,
+    Layout,
+    QuantOp,
+    Alignment,
+    std::enable_if_t<isFinegrained(QuantOp)>> {
+ private:
+  using SmemScaleType = half_t;
 
-public:
-    using IteratorScale
-        = cutlass::transform::threadblock::FineGrainedScaleZeroIterator<cutlass::MatrixShape<1, MmaShape::kN>, Element,
-            Layout, 0, Alignment>;
+ public:
+  using IteratorScale =
+      cutlass::transform::threadblock::FineGrainedScaleZeroIterator<
+          cutlass::MatrixShape<1, MmaShape::kN>,
+          Element,
+          Layout,
+          0,
+          Alignment>;
 
-    using SmemIteratorScale
-        = cutlass::transform::threadblock::FineGrainedScaleZeroIterator<cutlass::MatrixShape<1, MmaShape::kN>,
-            SmemScaleType, Layout, 0, Alignment>;
+  using SmemIteratorScale =
+      cutlass::transform::threadblock::FineGrainedScaleZeroIterator<
+          cutlass::MatrixShape<1, MmaShape::kN>,
+          SmemScaleType,
+          Layout,
+          0,
+          Alignment>;
 };
 
 // Per column iterators
-template <typename MmaShape, typename Element, typename Layout, WeightOnlyQuantOp QuantOp, int Alignment>
-struct DefaultScaleIteratorsPipelined<MmaShape, Element, Layout, QuantOp, Alignment,
-    std::enable_if_t<!isFinegrained(QuantOp)>>
-{
-    static_assert((MmaShape::kN % Alignment) == 0, "");
+template <typename MmaShape,
+          typename Element,
+          typename Layout,
+          WeightOnlyQuantOp QuantOp,
+          int Alignment>
+struct DefaultScaleIteratorsPipelined<
+    MmaShape,
+    Element,
+    Layout,
+    QuantOp,
+    Alignment,
+    std::enable_if_t<!isFinegrained(QuantOp)>> {
+  static_assert((MmaShape::kN % Alignment) == 0, "");
 
-private:
-    // ThreadMap for scale iterator
-    using IteratorScaleThreadMap = transform::PitchLinearStripminedThreadMap<layout::PitchLinearShape<MmaShape::kN, 1>,
-        MmaShape::kN / Alignment, Alignment>;
-    using SmemScaleType = half_t;
+ private:
+  // ThreadMap for scale iterator
+  using IteratorScaleThreadMap = transform::PitchLinearStripminedThreadMap<
+      layout::PitchLinearShape<MmaShape::kN, 1>,
+      MmaShape::kN / Alignment,
+      Alignment>;
+  using SmemScaleType = half_t;
 
-public:
-    // Define iterators over tiles from the scale operand
-    using IteratorScale = cutlass::transform::threadblock::PredicatedTileIterator<cutlass::MatrixShape<1, MmaShape::kN>,
-        Element, Layout, 0, IteratorScaleThreadMap, Alignment>;
+ public:
+  // Define iterators over tiles from the scale operand
+  using IteratorScale = cutlass::transform::threadblock::PredicatedTileIterator<
+      cutlass::MatrixShape<1, MmaShape::kN>,
+      Element,
+      Layout,
+      0,
+      IteratorScaleThreadMap,
+      Alignment>;
 
-    using SmemIteratorScale
-        = cutlass::transform::threadblock::PredicatedTileIterator<cutlass::MatrixShape<1, MmaShape::kN>, SmemScaleType,
-            Layout, 0, IteratorScaleThreadMap, Alignment>;
+  using SmemIteratorScale =
+      cutlass::transform::threadblock::PredicatedTileIterator<
+          cutlass::MatrixShape<1, MmaShape::kN>,
+          SmemScaleType,
+          Layout,
+          0,
+          IteratorScaleThreadMap,
+          Alignment>;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -116,57 +153,110 @@ template <
     typename InstructionShape,
     /// Operation performed by GEMM
     typename Operator_>
-struct DqMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementScale, LayoutScale, kAlignmentScale,
-    ElementAccumulator, layout::RowMajor, OperatorClass, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2,
-    Operator_, SharedMemoryClearOption::kNone,
-    typename platform::enable_if<(
-        ArchTag::kMinComputeCapability < 80 && !layout::IsColumnMajorTileInterleave<LayoutB>::value)>::type>
-{
+struct DqMma<ElementA,
+             LayoutA,
+             kAlignmentA,
+             ElementB,
+             LayoutB,
+             kAlignmentB,
+             ElementScale,
+             LayoutScale,
+             kAlignmentScale,
+             ElementAccumulator,
+             layout::RowMajor,
+             OperatorClass,
+             ArchTag,
+             ThreadblockShape,
+             WarpShape,
+             InstructionShape,
+             2,
+             Operator_,
+             SharedMemoryClearOption::kNone,
+             typename platform::enable_if<(
+                 ArchTag::kMinComputeCapability < 80 &&
+                 !layout::IsColumnMajorTileInterleave<LayoutB>::value)>::type> {
+  static_assert(platform::is_same<ElementA, half_t>::value ||
+                    platform::is_same<ElementA, bfloat16_t>::value,
+                "Element A must be fp16 or bf16");
 
-    static_assert(platform::is_same<ElementA, half_t>::value || platform::is_same<ElementA, bfloat16_t>::value,
-        "Element A must be fp16 or bf16");
+  static_assert(platform::is_same<ElementB, uint8_t>::value ||
+                    platform::is_same<ElementB, uint4b_t>::value,
+                "Element B must be uint8 or uint4");
 
-    static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-        "Element B must be uint8 or uint4");
+  using OperatorInfo = arch::DetagOperator<Operator_>;
+  using Operator = typename OperatorInfo::Operator;
+  static_assert(OperatorInfo::QuantOp ==
+                    WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY,
+                "");
 
-    using OperatorInfo = arch::DetagOperator<Operator_>;
-    using Operator = typename OperatorInfo::Operator;
-    static_assert(OperatorInfo::QuantOp == WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY, "");
+  static constexpr bool DqAfterLDG =
+      platform::is_same<arch::OpMultiplyAdd, Operator>::value;
+  using MmaCoreElementA = half_t;
+  using MmaCoreElementB = typename platform::
+      conditional<DqAfterLDG, MmaCoreElementA, ElementB>::type;
 
-    static constexpr bool DqAfterLDG = platform::is_same<arch::OpMultiplyAdd, Operator>::value;
-    using MmaCoreElementA = half_t;
-    using MmaCoreElementB = typename platform::conditional<DqAfterLDG, MmaCoreElementA, ElementB>::type;
+  // Define the MmaCore components
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          MmaCoreElementA,
+                                                          LayoutA,
+                                                          MmaCoreElementB,
+                                                          LayoutB,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          OperatorClass,
+                                                          2,
+                                                          Operator>;
 
-    // Define the MmaCore components
-    using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape, WarpShape, InstructionShape,
-        MmaCoreElementA, LayoutA, MmaCoreElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, 2,
-        Operator>;
+  // Define iterators over tiles from the A operand
+  using IteratorA = cutlass::transform::threadblock::PredicatedTileIterator<
+      cutlass::MatrixShape<MmaCore::Shape::kM, MmaCore::Shape::kK>,
+      ElementA,
+      LayoutA,
+      1,
+      typename MmaCore::IteratorThreadMapA,
+      kAlignmentA>;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = cutlass::transform::threadblock::PredicatedTileIterator<
-        cutlass::MatrixShape<MmaCore::Shape::kM, MmaCore::Shape::kK>, ElementA, LayoutA, 1,
-        typename MmaCore::IteratorThreadMapA, kAlignmentA>;
+  // Define iterators over tiles from the B operand
+  using IteratorB = cutlass::transform::threadblock::PredicatedTileIterator<
+      cutlass::MatrixShape<MmaCore::Shape::kK, MmaCore::Shape::kN>,
+      ElementB,
+      LayoutB,
+      0,
+      typename MmaCore::IteratorThreadMapB,
+      kAlignmentB>;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = cutlass::transform::threadblock::PredicatedTileIterator<
-        cutlass::MatrixShape<MmaCore::Shape::kK, MmaCore::Shape::kN>, ElementB, LayoutB, 0,
-        typename MmaCore::IteratorThreadMapB, kAlignmentB>;
+  using ScaleIterators = DefaultScaleIteratorsPipelined<typename MmaCore::Shape,
+                                                        ElementScale,
+                                                        LayoutScale,
+                                                        OperatorInfo::QuantOp,
+                                                        kAlignmentScale>;
 
-    using ScaleIterators = DefaultScaleIteratorsPipelined<typename MmaCore::Shape, ElementScale, LayoutScale,
-        OperatorInfo::QuantOp, kAlignmentScale>;
+  // Define iterators over tiles from the scale operand
+  using IteratorScale = typename ScaleIterators::IteratorScale;
 
-    // Define iterators over tiles from the scale operand
-    using IteratorScale = typename ScaleIterators::IteratorScale;
+  using SmemIteratorScale = typename ScaleIterators::SmemIteratorScale;
 
-    using SmemIteratorScale = typename ScaleIterators::SmemIteratorScale;
+  using Converters =
+      SetConverters<IteratorB, typename MmaCore::MmaPolicy::Operator, Operator>;
 
-    using Converters = SetConverters<IteratorB, typename MmaCore::MmaPolicy::Operator, Operator>;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::DqMmaPipelined<typename MmaCore::Shape, IteratorA,
-        typename MmaCore::SmemIteratorA, IteratorB, typename MmaCore::SmemIteratorB, IteratorScale, SmemIteratorScale,
-        ElementAccumulator, layout::RowMajor, typename MmaCore::MmaPolicy, typename Converters::TransformAfterLDG,
-        typename Converters::TransformAfterLDS, OperatorInfo::QuantOp>;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = cutlass::gemm::threadblock::DqMmaPipelined<
+      typename MmaCore::Shape,
+      IteratorA,
+      typename MmaCore::SmemIteratorA,
+      IteratorB,
+      typename MmaCore::SmemIteratorB,
+      IteratorScale,
+      SmemIteratorScale,
+      ElementAccumulator,
+      layout::RowMajor,
+      typename MmaCore::MmaPolicy,
+      typename Converters::TransformAfterLDG,
+      typename Converters::TransformAfterLDS,
+      OperatorInfo::QuantOp>;
 };
 
 // Specialization to handle column major interleave B
@@ -203,82 +293,140 @@ template <
     typename InstructionShape,
     /// Operation performed by GEMM
     typename Operator_>
-struct DqMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementScale, LayoutScale, kAlignmentScale,
-    ElementAccumulator, layout::RowMajor, OperatorClass, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2,
-    Operator_, SharedMemoryClearOption::kNone,
-    typename platform::enable_if<(
-        ArchTag::kMinComputeCapability < 80 && layout::IsColumnMajorTileInterleave<LayoutB>::value)>::type>
-{
+struct DqMma<ElementA,
+             LayoutA,
+             kAlignmentA,
+             ElementB,
+             LayoutB,
+             kAlignmentB,
+             ElementScale,
+             LayoutScale,
+             kAlignmentScale,
+             ElementAccumulator,
+             layout::RowMajor,
+             OperatorClass,
+             ArchTag,
+             ThreadblockShape,
+             WarpShape,
+             InstructionShape,
+             2,
+             Operator_,
+             SharedMemoryClearOption::kNone,
+             typename platform::enable_if<(
+                 ArchTag::kMinComputeCapability < 80 &&
+                 layout::IsColumnMajorTileInterleave<LayoutB>::value)>::type> {
+  static_assert(platform::is_same<ElementA, half_t>::value ||
+                    platform::is_same<ElementA, bfloat16_t>::value,
+                "Element A must be fp16 or bf16");
 
-    static_assert(platform::is_same<ElementA, half_t>::value || platform::is_same<ElementA, bfloat16_t>::value,
-        "Element A must be fp16 or bf16");
+  static_assert(platform::is_same<ElementB, uint8_t>::value ||
+                    platform::is_same<ElementB, uint4b_t>::value,
+                "Element B must be uint8 or uint4");
 
-    static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-        "Element B must be uint8 or uint4");
+  using OperatorInfo = arch::DetagOperator<Operator_>;
+  using Operator = typename OperatorInfo::Operator;
 
-    using OperatorInfo = arch::DetagOperator<Operator_>;
-    using Operator = typename OperatorInfo::Operator;
+  static constexpr bool DqAfterLDG =
+      platform::is_same<arch::OpMultiplyAdd, Operator>::value;
+  using MmaCoreElementA = half_t;
+  using MmaCoreElementB = typename platform::
+      conditional<DqAfterLDG, MmaCoreElementA, ElementB>::type;
 
-    static constexpr bool DqAfterLDG = platform::is_same<arch::OpMultiplyAdd, Operator>::value;
-    using MmaCoreElementA = half_t;
-    using MmaCoreElementB = typename platform::conditional<DqAfterLDG, MmaCoreElementA, ElementB>::type;
+  // Define the MmaCore components
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          MmaCoreElementA,
+                                                          LayoutA,
+                                                          MmaCoreElementB,
+                                                          layout::ColumnMajor,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          OperatorClass,
+                                                          2,
+                                                          Operator>;
 
-    // Define the MmaCore components
-    using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape, WarpShape, InstructionShape,
-        MmaCoreElementA, LayoutA, MmaCoreElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor,
-        OperatorClass, 2, Operator>;
+  // Define iterators over tiles from the A operand
+  using IteratorA = cutlass::transform::threadblock::PredicatedTileIterator<
+      cutlass::MatrixShape<MmaCore::Shape::kM, MmaCore::Shape::kK>,
+      ElementA,
+      LayoutA,
+      1,
+      typename MmaCore::IteratorThreadMapA,
+      kAlignmentA>;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = cutlass::transform::threadblock::PredicatedTileIterator<
-        cutlass::MatrixShape<MmaCore::Shape::kM, MmaCore::Shape::kK>, ElementA, LayoutA, 1,
-        typename MmaCore::IteratorThreadMapA, kAlignmentA>;
+ private:
+  static constexpr int ColumnsInterleaved = LayoutB::kColumnsInterleaved;
+  static constexpr int RowsPerTile = LayoutB::kRowsPerTile;
+  static_assert(!(MmaCore::Shape::kN % ColumnsInterleaved), "");
+  static_assert(RowsPerTile == MmaCore::Shape::kK, "");
 
-private:
-    static constexpr int ColumnsInterleaved = LayoutB::kColumnsInterleaved;
-    static constexpr int RowsPerTile = LayoutB::kRowsPerTile;
-    static_assert(!(MmaCore::Shape::kN % ColumnsInterleaved), "");
-    static_assert(RowsPerTile == MmaCore::Shape::kK, "");
+  using OriginalThreadMap = typename MmaCore::IteratorThreadMapB;
+  using OriginalWarpArrangement =
+      typename OriginalThreadMap::Detail::WarpThreadArrangement;
+  static_assert(!(OriginalWarpArrangement::kStrided % ColumnsInterleaved), "");
 
-    using OriginalThreadMap = typename MmaCore::IteratorThreadMapB;
-    using OriginalWarpArrangement = typename OriginalThreadMap::Detail::WarpThreadArrangement;
-    static_assert(!(OriginalWarpArrangement::kStrided % ColumnsInterleaved), "");
+  using GmemIteratorShape =
+      MatrixShape<MmaCore::Shape::kK * ColumnsInterleaved,
+                  MmaCore::Shape::kN / ColumnsInterleaved>;
+  using GmemThreadMapB = transform::PitchLinearWarpRakedThreadMap<
+      layout::PitchLinearShape<GmemIteratorShape::kRow,
+                               GmemIteratorShape::kColumn>,
+      OriginalThreadMap::kThreads,
+      layout::PitchLinearShape<
+          OriginalWarpArrangement::kContiguous * ColumnsInterleaved,
+          OriginalWarpArrangement::kStrided / ColumnsInterleaved>,
+      MmaCore::kAccessSizeInBits / sizeof_bits<ElementB>::value>;
 
-    using GmemIteratorShape
-        = MatrixShape<MmaCore::Shape::kK * ColumnsInterleaved, MmaCore::Shape::kN / ColumnsInterleaved>;
-    using GmemThreadMapB = transform::PitchLinearWarpRakedThreadMap<
-        layout::PitchLinearShape<GmemIteratorShape::kRow, GmemIteratorShape::kColumn>, OriginalThreadMap::kThreads,
-        layout::PitchLinearShape<OriginalWarpArrangement::kContiguous * ColumnsInterleaved,
-            OriginalWarpArrangement::kStrided / ColumnsInterleaved>,
-        MmaCore::kAccessSizeInBits / sizeof_bits<ElementB>::value>;
+ public:
+  // Define iterators over tiles from the B operand
+  using IteratorB = cutlass::transform::threadblock::PredicatedTileIterator<
+      GmemIteratorShape,
+      ElementB,
+      layout::ColumnMajor,
+      0,
+      GmemThreadMapB,
+      kAlignmentB>;
 
-public:
-    // Define iterators over tiles from the B operand
-    using IteratorB = cutlass::transform::threadblock::PredicatedTileIterator<GmemIteratorShape, ElementB,
-        layout::ColumnMajor, 0, GmemThreadMapB, kAlignmentB>;
+  // ThreadMap for scale iterator
+  static_assert((MmaCore::Shape::kN % kAlignmentScale) == 0, "");
+  using IteratorScaleThreadMap = transform::PitchLinearStripminedThreadMap<
+      layout::PitchLinearShape<MmaCore::Shape::kN, 1>,
+      MmaCore::Shape::kN / kAlignmentScale,
+      kAlignmentScale>;
 
-    // ThreadMap for scale iterator
-    static_assert((MmaCore::Shape::kN % kAlignmentScale) == 0, "");
-    using IteratorScaleThreadMap
-        = transform::PitchLinearStripminedThreadMap<layout::PitchLinearShape<MmaCore::Shape::kN, 1>,
-            MmaCore::Shape::kN / kAlignmentScale, kAlignmentScale>;
+  using ScaleIterators = DefaultScaleIteratorsPipelined<typename MmaCore::Shape,
+                                                        ElementScale,
+                                                        LayoutScale,
+                                                        OperatorInfo::QuantOp,
+                                                        kAlignmentScale>;
 
-    using ScaleIterators = DefaultScaleIteratorsPipelined<typename MmaCore::Shape, ElementScale, LayoutScale,
-        OperatorInfo::QuantOp, kAlignmentScale>;
+  // Define iterators over tiles from the scale operand
+  using IteratorScale = typename ScaleIterators::IteratorScale;
 
-    // Define iterators over tiles from the scale operand
-    using IteratorScale = typename ScaleIterators::IteratorScale;
+  using SmemIteratorScale = typename ScaleIterators::SmemIteratorScale;
 
-    using SmemIteratorScale = typename ScaleIterators::SmemIteratorScale;
+  using Converters =
+      SetConverters<IteratorB, typename MmaCore::MmaPolicy::Operator, Operator>;
 
-    using Converters = SetConverters<IteratorB, typename MmaCore::MmaPolicy::Operator, Operator>;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::DqMmaPipelined<typename MmaCore::Shape, IteratorA,
-        typename MmaCore::SmemIteratorA, IteratorB, typename MmaCore::SmemIteratorB, IteratorScale, SmemIteratorScale,
-        ElementAccumulator, layout::RowMajor, typename MmaCore::MmaPolicy, typename Converters::TransformAfterLDG,
-        typename Converters::TransformAfterLDS, OperatorInfo::QuantOp>;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = cutlass::gemm::threadblock::DqMmaPipelined<
+      typename MmaCore::Shape,
+      IteratorA,
+      typename MmaCore::SmemIteratorA,
+      IteratorB,
+      typename MmaCore::SmemIteratorB,
+      IteratorScale,
+      SmemIteratorScale,
+      ElementAccumulator,
+      layout::RowMajor,
+      typename MmaCore::MmaPolicy,
+      typename Converters::TransformAfterLDG,
+      typename Converters::TransformAfterLDS,
+      OperatorInfo::QuantOp>;
 };
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma.h
index b50d66380e..31915cf389 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,7 +27,8 @@ namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), fp16 activation & int8 weight, mma pipelined (stage=2)
+/// Specialization for row-major output (OperatorClass TensorOp), fp16
+/// activation & int8 weight, mma pipelined (stage=2)
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -49,34 +50,61 @@ template <
     typename InstructionShape,
     /// Operation performed by GEMM
     typename Operator>
-struct DefaultMma<cutlass::half_t, LayoutA, kAlignmentA, uint8_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator>
-{
+struct DefaultMma<cutlass::half_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint8_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
+  using Mma = DqMma<half_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint8_t,
+                    LayoutB,
+                    kAlignmentB,
+                    half_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    2,
+                    Operator>;
 
-    using Mma = DqMma<half_t, LayoutA, kAlignmentA, uint8_t, LayoutB, kAlignmentB, half_t, layout::RowMajor,
-        kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, 2, Operator>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), fp16 activation & int4 weight, mma pipelined (stage=2)
+/// Specialization for row-major output (OperatorClass TensorOp), fp16
+/// activation & int4 weight, mma pipelined (stage=2)
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -98,35 +126,61 @@ template <
     typename InstructionShape,
     /// Operation performed by GEMM
     typename Operator>
-struct DefaultMma<cutlass::half_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator>
-{
+struct DefaultMma<cutlass::half_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint4b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
+  using Mma = DqMma<half_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint4b_t,
+                    LayoutB,
+                    kAlignmentB,
+                    half_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    2,
+                    Operator>;
 
-    using Mma = DqMma<half_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, half_t, layout::RowMajor,
-        kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, 2, Operator>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), fp16 activation & int8 weight, mma multistage
-/// (stage>=3)
+/// Specialization for row-major output (OperatorClass TensorOp), fp16
+/// activation & int8 weight, mma multistage (stage>=3)
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -152,36 +206,64 @@ template <
     int kStages,
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
-struct DefaultMma<cutlass::half_t, LayoutA, kAlignmentA, uint8_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
-    false, SharedMemoryClear>
-{
+struct DefaultMma<cutlass::half_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint8_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  kStages,
+                  Operator,
+                  false,
+                  SharedMemoryClear> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
+  using Mma = DqMma<half_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint8_t,
+                    LayoutB,
+                    kAlignmentB,
+                    half_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    kStages,
+                    Operator,
+                    SharedMemoryClear>;
 
-    using Mma = DqMma<half_t, LayoutA, kAlignmentA, uint8_t, LayoutB, kAlignmentB, half_t, layout::RowMajor,
-        kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), fp16 activation & int4 weight, mma multistage
-/// (stage>=3)
+/// Specialization for row-major output (OperatorClass TensorOp), fp16
+/// activation & int4 weight, mma multistage (stage>=3)
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -207,37 +289,65 @@ template <
     int kStages,
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
-struct DefaultMma<cutlass::half_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
-    false, SharedMemoryClear>
-{
+struct DefaultMma<cutlass::half_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint4b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  kStages,
+                  Operator,
+                  false,
+                  SharedMemoryClear> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
+  using Mma = DqMma<half_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint4b_t,
+                    LayoutB,
+                    kAlignmentB,
+                    half_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    kStages,
+                    Operator,
+                    SharedMemoryClear>;
 
-    using Mma = DqMma<half_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, half_t, layout::RowMajor,
-        kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 #ifdef ENABLE_FP8
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), fp8 activation & int4 weight, mma multistage
-/// (stage>=3)
+/// Specialization for row-major output (OperatorClass TensorOp), fp8 activation
+/// & int4 weight, mma multistage (stage>=3)
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -263,36 +373,65 @@ template <
     int kStages,
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
-struct DefaultMma<cutlass::float_e4m3_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
-    false, SharedMemoryClear>
-{
+struct DefaultMma<cutlass::float_e4m3_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint4b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  kStages,
+                  Operator,
+                  false,
+                  SharedMemoryClear> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
+  using Mma = DqMma<cutlass::float_e4m3_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint4b_t,
+                    LayoutB,
+                    kAlignmentB,
+                    half_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    kStages,
+                    Operator,
+                    SharedMemoryClear>;
 
-    using Mma = DqMma<cutlass::float_e4m3_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, half_t,
-        layout::RowMajor, kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag,
-        ThreadblockShape, WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 #endif
 
-// fp16 x fp16 specialization on Ampere to use mma multistage for 2 stage. Helps avoid reg spills on
-// large tile when not enough shared mem is present to do 3+ stage
+// fp16 x fp16 specialization on Ampere to use mma multistage for 2 stage. Helps
+// avoid reg spills on large tile when not enough shared mem is present to do 3+
+// stage
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -318,39 +457,86 @@ template <
     bool GatherA,
     /// Gather operand B by using an index array
     bool GatherB>
-struct DefaultMma<half_t, LayoutA, kAlignmentA, half_t, LayoutB, kAlignmentB, ElementAccumulator, layout::RowMajor,
-    arch::OpClassTensorOp, arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, 2, Operator, false,
-    SharedMemoryClear, GatherA, GatherB>
-{
+struct DefaultMma<half_t,
+                  LayoutA,
+                  kAlignmentA,
+                  half_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  arch::Sm80,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator,
+                  false,
+                  SharedMemoryClear,
+                  GatherA,
+                  GatherB> {
+  // Define the MmaCore components
+  // 3 is used on purpose here to trigger components for mma multistage
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          half_t,
+                                                          LayoutA,
+                                                          half_t,
+                                                          LayoutB,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          arch::OpClassTensorOp,
+                                                          3,
+                                                          Operator>;
 
-    // Define the MmaCore components
-    // 3 is used on purpose here to trigger components for mma multistage
-    using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape, WarpShape, InstructionShape,
-        half_t, LayoutA, half_t, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, 3, Operator>;
+  // Define iterators over tiles from the A operand
+  using ThreadMapA = typename MmaCore::IteratorThreadMapA;
+  using AccessTypeA = cutlass::Array<half_t, kAlignmentA>;
+  using IteratorA =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
+          half_t,
+          LayoutA,
+          1,
+          ThreadMapA,
+          AccessTypeA,
+          GatherA>;
 
-    // Define iterators over tiles from the A operand
-    using ThreadMapA = typename MmaCore::IteratorThreadMapA;
-    using AccessTypeA = cutlass::Array<half_t, kAlignmentA>;
-    using IteratorA = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>, half_t, LayoutA, 1, ThreadMapA, AccessTypeA,
-        GatherA>;
+  // Define iterators over tiles from the B operand
+  using ThreadMapB = typename MmaCore::IteratorThreadMapB;
+  using AccessTypeB = cutlass::Array<half_t, kAlignmentB>;
+  using IteratorB =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>,
+          half_t,
+          LayoutB,
+          0,
+          ThreadMapB,
+          AccessTypeB,
+          GatherB>;
 
-    // Define iterators over tiles from the B operand
-    using ThreadMapB = typename MmaCore::IteratorThreadMapB;
-    using AccessTypeB = cutlass::Array<half_t, kAlignmentB>;
-    using IteratorB = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>, half_t, LayoutB, 0, ThreadMapB, AccessTypeB,
-        GatherB>;
-
-    // Define the threadblock-scoped multistage matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::MmaMultistage<typename MmaCore::Shape, IteratorA,
-        typename MmaCore::SmemIteratorA, MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB,
-        MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, typename MmaCore::MmaPolicy, 2>;
+  // Define the threadblock-scoped multistage matrix multiply
+  using ThreadblockMma =
+      cutlass::gemm::threadblock::MmaMultistage<typename MmaCore::Shape,
+                                                IteratorA,
+                                                typename MmaCore::SmemIteratorA,
+                                                MmaCore::kCacheOpA,
+                                                IteratorB,
+                                                typename MmaCore::SmemIteratorB,
+                                                MmaCore::kCacheOpB,
+                                                ElementAccumulator,
+                                                layout::RowMajor,
+                                                typename MmaCore::MmaPolicy,
+                                                2>;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), fbf16 activation & int2 weight, mma multistage
+/// Specialization for row-major output (OperatorClass TensorOp), fbf16
+/// activation & int2 weight, mma multistage
 
 template <
     /// Layout type for A matrix operand
@@ -373,26 +559,50 @@ template <
     typename InstructionShape,
     /// Operation performed by GEMM
     typename Operator>
-struct DefaultMma<cutlass::half_t, LayoutA, kAlignmentA, uint2b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator>
-{
-private:
-    using Mma = DefaultWint2xMma<half_t, LayoutA, kAlignmentA, uint2b_t, LayoutB, kAlignmentB,
-        ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, 2, Operator>;
+struct DefaultMma<cutlass::half_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint2b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator> {
+ private:
+  using Mma = DefaultWint2xMma<half_t,
+                               LayoutA,
+                               kAlignmentA,
+                               uint2b_t,
+                               LayoutB,
+                               kAlignmentB,
+                               ElementAccumulator,
+                               layout::RowMajor,
+                               arch::OpClassTensorOp,
+                               ArchTag,
+                               ThreadblockShape,
+                               WarpShape,
+                               InstructionShape,
+                               2,
+                               Operator>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define the threadblock-scoped multistage matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped multistage matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 template <
@@ -420,29 +630,55 @@ template <
     int kStages,
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
-struct DefaultMma<half_t, LayoutA, kAlignmentA, uint2b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
-    false, SharedMemoryClear>
-{
-private:
-    using Mma = DefaultWint2xMma<half_t, LayoutA, kAlignmentA, uint2b_t, LayoutB, kAlignmentB,
-        ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+struct DefaultMma<half_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint2b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  kStages,
+                  Operator,
+                  false,
+                  SharedMemoryClear> {
+ private:
+  using Mma = DefaultWint2xMma<half_t,
+                               LayoutA,
+                               kAlignmentA,
+                               uint2b_t,
+                               LayoutB,
+                               kAlignmentB,
+                               ElementAccumulator,
+                               layout::RowMajor,
+                               arch::OpClassTensorOp,
+                               ArchTag,
+                               ThreadblockShape,
+                               WarpShape,
+                               InstructionShape,
+                               kStages,
+                               Operator,
+                               SharedMemoryClear>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define the threadblock-scoped multistage matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped multistage matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma_bf16.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma_bf16.h
index 300261c3f0..1ff648e0bf 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma_bf16.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma_bf16.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,7 +27,8 @@ namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), bf16 activation & bf16 weight
+/// Specialization for row-major output (OperatorClass TensorOp), bf16
+/// activation & bf16 weight
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -55,40 +56,85 @@ template <
     bool GatherA,
     /// Gather operand B by using an index array
     bool GatherB>
-struct DefaultMma<bfloat16_t, LayoutA, kAlignmentA, bfloat16_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator, false,
-    SharedMemoryClear, GatherA, GatherB>
-{
+struct DefaultMma<bfloat16_t,
+                  LayoutA,
+                  kAlignmentA,
+                  bfloat16_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator,
+                  false,
+                  SharedMemoryClear,
+                  GatherA,
+                  GatherB> {
+ private:
+  // Conversions only needed pre-ampere. This will trigger mma pipeline, so we
+  // convert before STS.
+  static constexpr bool arch_has_bf16_mma =
+      ArchTag::kMinComputeCapability >= 80;
+  using MmaElementA = typename platform::
+      conditional<arch_has_bf16_mma, bfloat16_t, half_t>::type;
+  using MmaElementB = typename platform::
+      conditional<arch_has_bf16_mma, bfloat16_t, half_t>::type;
 
-private:
-    // Conversions only needed pre-ampere. This will trigger mma pipeline, so we convert before STS.
-    static constexpr bool arch_has_bf16_mma = ArchTag::kMinComputeCapability >= 80;
-    using MmaElementA = typename platform::conditional<arch_has_bf16_mma, bfloat16_t, half_t>::type;
-    using MmaElementB = typename platform::conditional<arch_has_bf16_mma, bfloat16_t, half_t>::type;
+ public:
+  // Define the MmaCore components
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          MmaElementA,
+                                                          LayoutA,
+                                                          MmaElementB,
+                                                          LayoutB,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          arch::OpClassTensorOp,
+                                                          2,
+                                                          Operator>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore =
-        typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape, WarpShape, InstructionShape, MmaElementA,
-            LayoutA, MmaElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, 2, Operator>;
+  using IteratorA = cutlass::transform::threadblock::PredicatedTileIterator<
+      cutlass::MatrixShape<MmaCore::Shape::kM, MmaCore::Shape::kK>,
+      bfloat16_t,
+      LayoutA,
+      1,
+      typename MmaCore::IteratorThreadMapA,
+      kAlignmentA,
+      GatherA>;
 
-    using IteratorA = cutlass::transform::threadblock::PredicatedTileIterator<
-        cutlass::MatrixShape<MmaCore::Shape::kM, MmaCore::Shape::kK>, bfloat16_t, LayoutA, 1,
-        typename MmaCore::IteratorThreadMapA, kAlignmentA, GatherA>;
+  // Define iterators over tiles from the B operand
+  using IteratorB = cutlass::transform::threadblock::PredicatedTileIterator<
+      cutlass::MatrixShape<MmaCore::Shape::kK, MmaCore::Shape::kN>,
+      bfloat16_t,
+      LayoutB,
+      0,
+      typename MmaCore::IteratorThreadMapB,
+      kAlignmentB,
+      GatherB>;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = cutlass::transform::threadblock::PredicatedTileIterator<
-        cutlass::MatrixShape<MmaCore::Shape::kK, MmaCore::Shape::kN>, bfloat16_t, LayoutB, 0,
-        typename MmaCore::IteratorThreadMapB, kAlignmentB, GatherB>;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::MmaPipelined<typename MmaCore::Shape, IteratorA,
-        typename MmaCore::SmemIteratorA, IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator,
-        layout::RowMajor, typename MmaCore::MmaPolicy>;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma =
+      cutlass::gemm::threadblock::MmaPipelined<typename MmaCore::Shape,
+                                               IteratorA,
+                                               typename MmaCore::SmemIteratorA,
+                                               IteratorB,
+                                               typename MmaCore::SmemIteratorB,
+                                               ElementAccumulator,
+                                               layout::RowMajor,
+                                               typename MmaCore::MmaPolicy>;
 };
 
-// bf16 x bf16 specialization on Ampere to use mma multistage for 2 stage. Helps avoid reg spills on
-// large tile when not enough shared mem is present to do 3+ stage
+// bf16 x bf16 specialization on Ampere to use mma multistage for 2 stage. Helps
+// avoid reg spills on large tile when not enough shared mem is present to do 3+
+// stage
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -114,40 +160,86 @@ template <
     bool GatherA,
     /// Gather operand B by using an index array
     bool GatherB>
-struct DefaultMma<bfloat16_t, LayoutA, kAlignmentA, bfloat16_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, 2, Operator,
-    false, SharedMemoryClear, GatherA, GatherB>
-{
+struct DefaultMma<bfloat16_t,
+                  LayoutA,
+                  kAlignmentA,
+                  bfloat16_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  arch::Sm80,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator,
+                  false,
+                  SharedMemoryClear,
+                  GatherA,
+                  GatherB> {
+  // Define the MmaCore components
+  // 3 is used on purpose here to trigger components for mma multistage
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          bfloat16_t,
+                                                          LayoutA,
+                                                          bfloat16_t,
+                                                          LayoutB,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          arch::OpClassTensorOp,
+                                                          3,
+                                                          Operator>;
 
-    // Define the MmaCore components
-    // 3 is used on purpose here to trigger components for mma multistage
-    using MmaCore =
-        typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape, WarpShape, InstructionShape, bfloat16_t,
-            LayoutA, bfloat16_t, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, 3, Operator>;
+  // Define iterators over tiles from the A operand
+  using ThreadMapA = typename MmaCore::IteratorThreadMapA;
+  using AccessTypeA = cutlass::Array<bfloat16_t, kAlignmentA>;
+  using IteratorA =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
+          bfloat16_t,
+          LayoutA,
+          1,
+          ThreadMapA,
+          AccessTypeA,
+          GatherA>;
 
-    // Define iterators over tiles from the A operand
-    using ThreadMapA = typename MmaCore::IteratorThreadMapA;
-    using AccessTypeA = cutlass::Array<bfloat16_t, kAlignmentA>;
-    using IteratorA = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>, bfloat16_t, LayoutA, 1, ThreadMapA,
-        AccessTypeA, GatherA>;
+  // Define iterators over tiles from the B operand
+  using ThreadMapB = typename MmaCore::IteratorThreadMapB;
+  using AccessTypeB = cutlass::Array<bfloat16_t, kAlignmentB>;
+  using IteratorB =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>,
+          bfloat16_t,
+          LayoutB,
+          0,
+          ThreadMapB,
+          AccessTypeB,
+          GatherB>;
 
-    // Define iterators over tiles from the B operand
-    using ThreadMapB = typename MmaCore::IteratorThreadMapB;
-    using AccessTypeB = cutlass::Array<bfloat16_t, kAlignmentB>;
-    using IteratorB = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>, bfloat16_t, LayoutB, 0, ThreadMapB,
-        AccessTypeB, GatherB>;
-
-    // Define the threadblock-scoped multistage matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::MmaMultistage<typename MmaCore::Shape, IteratorA,
-        typename MmaCore::SmemIteratorA, MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB,
-        MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, typename MmaCore::MmaPolicy, 2>;
+  // Define the threadblock-scoped multistage matrix multiply
+  using ThreadblockMma =
+      cutlass::gemm::threadblock::MmaMultistage<typename MmaCore::Shape,
+                                                IteratorA,
+                                                typename MmaCore::SmemIteratorA,
+                                                MmaCore::kCacheOpA,
+                                                IteratorB,
+                                                typename MmaCore::SmemIteratorB,
+                                                MmaCore::kCacheOpB,
+                                                ElementAccumulator,
+                                                layout::RowMajor,
+                                                typename MmaCore::MmaPolicy,
+                                                2>;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), bf16 activation & int8 weight
+/// Specialization for row-major output (OperatorClass TensorOp), bf16
+/// activation & int8 weight
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -169,34 +261,61 @@ template <
     typename InstructionShape,
     /// Operation performed by GEMM
     typename Operator>
-struct DefaultMma<cutlass::bfloat16_t, LayoutA, kAlignmentA, uint8_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator>
-{
+struct DefaultMma<cutlass::bfloat16_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint8_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
+  using Mma = DqMma<bfloat16_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint8_t,
+                    LayoutB,
+                    kAlignmentB,
+                    bfloat16_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    2,
+                    Operator>;
 
-    using Mma = DqMma<bfloat16_t, LayoutA, kAlignmentA, uint8_t, LayoutB, kAlignmentB, bfloat16_t, layout::RowMajor,
-        kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, 2, Operator>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), bf16 activation & int4 weight
+/// Specialization for row-major output (OperatorClass TensorOp), bf16
+/// activation & int4 weight
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -218,34 +337,61 @@ template <
     typename InstructionShape,
     /// Operation performed by GEMM
     typename Operator>
-struct DefaultMma<cutlass::bfloat16_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator>
-{
+struct DefaultMma<cutlass::bfloat16_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint4b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
+  using Mma = DqMma<bfloat16_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint4b_t,
+                    LayoutB,
+                    kAlignmentB,
+                    bfloat16_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    2,
+                    Operator>;
 
-    using Mma = DqMma<bfloat16_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, bfloat16_t, layout::RowMajor,
-        kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, 2, Operator>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), bf16 activation & int8 weight
+/// Specialization for row-major output (OperatorClass TensorOp), bf16
+/// activation & int8 weight
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -271,35 +417,64 @@ template <
     int kStages,
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
-struct DefaultMma<cutlass::bfloat16_t, LayoutA, kAlignmentA, uint8_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
-    false, SharedMemoryClear>
-{
+struct DefaultMma<cutlass::bfloat16_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint8_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  kStages,
+                  Operator,
+                  false,
+                  SharedMemoryClear> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
+  using Mma = DqMma<bfloat16_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint8_t,
+                    LayoutB,
+                    kAlignmentB,
+                    bfloat16_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    kStages,
+                    Operator,
+                    SharedMemoryClear>;
 
-    using Mma = DqMma<bfloat16_t, LayoutA, kAlignmentA, uint8_t, LayoutB, kAlignmentB, bfloat16_t, layout::RowMajor,
-        kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), fp16 activation & int4 weight
+/// Specialization for row-major output (OperatorClass TensorOp), fp16
+/// activation & int4 weight
 template <
     /// Layout type for A matrix operand
     typename LayoutA,
@@ -325,35 +500,64 @@ template <
     int kStages,
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
-struct DefaultMma<bfloat16_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
-    false, SharedMemoryClear>
-{
+struct DefaultMma<bfloat16_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint4b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  kStages,
+                  Operator,
+                  false,
+                  SharedMemoryClear> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
 
-private:
-    static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
+  using Mma = DqMma<bfloat16_t,
+                    LayoutA,
+                    kAlignmentA,
+                    uint4b_t,
+                    LayoutB,
+                    kAlignmentB,
+                    bfloat16_t,
+                    layout::RowMajor,
+                    kAlignmentScale,
+                    ElementAccumulator,
+                    layout::RowMajor,
+                    arch::OpClassTensorOp,
+                    ArchTag,
+                    ThreadblockShape,
+                    WarpShape,
+                    InstructionShape,
+                    kStages,
+                    Operator,
+                    SharedMemoryClear>;
 
-    using Mma = DqMma<bfloat16_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAlignmentB, bfloat16_t, layout::RowMajor,
-        kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization for row-major output (OperatorClass TensorOp), fbf16 activation & int2 weight, mma multistage
+/// Specialization for row-major output (OperatorClass TensorOp), fbf16
+/// activation & int2 weight, mma multistage
 
 template <
     /// Layout type for A matrix operand
@@ -376,26 +580,50 @@ template <
     typename InstructionShape,
     /// Operation performed by GEMM
     typename Operator>
-struct DefaultMma<cutlass::bfloat16_t, LayoutA, kAlignmentA, uint2b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator>
-{
-private:
-    using Mma = DefaultWint2xMma<bfloat16_t, LayoutA, kAlignmentA, uint2b_t, LayoutB, kAlignmentB,
-        ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, 2, Operator>;
+struct DefaultMma<cutlass::bfloat16_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint2b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  2,
+                  Operator> {
+ private:
+  using Mma = DefaultWint2xMma<bfloat16_t,
+                               LayoutA,
+                               kAlignmentA,
+                               uint2b_t,
+                               LayoutB,
+                               kAlignmentB,
+                               ElementAccumulator,
+                               layout::RowMajor,
+                               arch::OpClassTensorOp,
+                               ArchTag,
+                               ThreadblockShape,
+                               WarpShape,
+                               InstructionShape,
+                               2,
+                               Operator>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define the threadblock-scoped multistage matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped multistage matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 template <
@@ -423,29 +651,55 @@ template <
     int kStages,
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
-struct DefaultMma<bfloat16_t, LayoutA, kAlignmentA, uint2b_t, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
-    false, SharedMemoryClear>
-{
-private:
-    using Mma = DefaultWint2xMma<bfloat16_t, LayoutA, kAlignmentA, uint2b_t, LayoutB, kAlignmentB,
-        ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
-        WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+struct DefaultMma<bfloat16_t,
+                  LayoutA,
+                  kAlignmentA,
+                  uint2b_t,
+                  LayoutB,
+                  kAlignmentB,
+                  ElementAccumulator,
+                  layout::RowMajor,
+                  arch::OpClassTensorOp,
+                  ArchTag,
+                  ThreadblockShape,
+                  WarpShape,
+                  InstructionShape,
+                  kStages,
+                  Operator,
+                  false,
+                  SharedMemoryClear> {
+ private:
+  using Mma = DefaultWint2xMma<bfloat16_t,
+                               LayoutA,
+                               kAlignmentA,
+                               uint2b_t,
+                               LayoutB,
+                               kAlignmentB,
+                               ElementAccumulator,
+                               layout::RowMajor,
+                               arch::OpClassTensorOp,
+                               ArchTag,
+                               ThreadblockShape,
+                               WarpShape,
+                               InstructionShape,
+                               kStages,
+                               Operator,
+                               SharedMemoryClear>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define the threadblock-scoped multistage matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped multistage matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma_core.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma_core.h
index e2bc640bac..58fd564416 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma_core.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_mma_core.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
@@ -66,10 +67,21 @@ template <
     cutlass::arch::CacheOperation::Kind CacheOpA,
     /// Cache operation of operand B
     cutlass::arch::CacheOperation::Kind CacheOpB>
-struct DefaultMmaCore<Shape_, WarpShape_, InstructionShape_, ElementA_,
-                      layout::RowMajor, uint2b_t, layout::ColumnMajor,
-                      ElementC_, LayoutC_, arch::OpClassTensorOp, Stages,
-                      Operator_, false, CacheOpA, CacheOpB> {
+struct DefaultMmaCore<Shape_,
+                      WarpShape_,
+                      InstructionShape_,
+                      ElementA_,
+                      layout::RowMajor,
+                      uint2b_t,
+                      layout::ColumnMajor,
+                      ElementC_,
+                      LayoutC_,
+                      arch::OpClassTensorOp,
+                      Stages,
+                      Operator_,
+                      false,
+                      CacheOpA,
+                      CacheOpB> {
   using Shape = Shape_;
   using WarpShape = WarpShape_;
   using InstructionShape = InstructionShape_;
@@ -104,7 +116,8 @@ struct DefaultMmaCore<Shape_, WarpShape_, InstructionShape_, ElementA_,
 
   /// Size of a threadblock-scoped access of B
   static constexpr int kMaxThreadsForB =
-      (Shape::kK * Shape::kN * sizeof_bits<ElementB>::value) / kAccessSizeInBits;
+      (Shape::kK * Shape::kN * sizeof_bits<ElementB>::value) /
+      kAccessSizeInBits;
   static constexpr int kThreadsForB =
       kMaxThreadsForB > kThreads ? kThreads : kMaxThreadsForB;
 
@@ -129,11 +142,13 @@ struct DefaultMmaCore<Shape_, WarpShape_, InstructionShape_, ElementA_,
   //
 
   using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise<
-      sizeof_bits<ElementA>::value, Shape::kK>;
+      sizeof_bits<ElementA>::value,
+      Shape::kK>;
 
   // Shared memory layout
   using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise<
-      sizeof_bits<ElementB>::value, Shape::kK>;
+      sizeof_bits<ElementB>::value,
+      Shape::kK>;
 
   //
   // Iterators to write to shared memory
@@ -141,26 +156,34 @@ struct DefaultMmaCore<Shape_, WarpShape_, InstructionShape_, ElementA_,
 
   /// ThreadMap of iterator A
   using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap<
-      layout::PitchLinearShape<Shape::kK, Shape::kM>, kThreads,
+      layout::PitchLinearShape<Shape::kK, Shape::kM>,
+      kThreads,
       layout::PitchLinearShape<kWarpThreadArrangementContiguousA,
                                kWarpThreadArrangementStridedA>,
       kAccessSizeInBits / sizeof_bits<ElementA>::value>;
 
   /// Shared memory iterator to A operand
   using SmemIteratorA = transform::threadblock::RegularTileAccessIterator<
-      MatrixShape<Shape::kM, Shape::kK>, ElementA, SmemLayoutA, 0,
+      MatrixShape<Shape::kM, Shape::kK>,
+      ElementA,
+      SmemLayoutA,
+      0,
       IteratorThreadMapA>;
 
   /// ThreadMap of iterator B
   using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap<
-      layout::PitchLinearShape<Shape::kK, Shape::kN>, kThreadsForB,
+      layout::PitchLinearShape<Shape::kK, Shape::kN>,
+      kThreadsForB,
       layout::PitchLinearShape<kWarpThreadArrangementContiguousB,
                                kWarpThreadArrangementStridedB>,
       kAccessSizeInBits / sizeof_bits<ElementB>::value>;
 
   /// Shared memory iterator to B operand
   using SmemIteratorB = transform::threadblock::RegularTileAccessIterator<
-      MatrixShape<Shape::kK, Shape::kN>, ElementB, SmemLayoutB, 1,
+      MatrixShape<Shape::kK, Shape::kN>,
+      ElementB,
+      SmemLayoutB,
+      1,
       IteratorThreadMapB>;
 
   //
@@ -168,13 +191,23 @@ struct DefaultMmaCore<Shape_, WarpShape_, InstructionShape_, ElementA_,
   //
 
   // Define the warp-level tensor op
-  using MmaTensorOp = typename cutlass::gemm::warp::DefaultMmaTensorOp<
-      WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB,
-      ElementC, LayoutC, Operator, WarpCount::kK>::Type;
+  using MmaTensorOp =
+      typename cutlass::gemm::warp::DefaultMmaTensorOp<WarpShape,
+                                                       InstructionShape,
+                                                       ElementA,
+                                                       SmemLayoutA,
+                                                       ElementB,
+                                                       SmemLayoutB,
+                                                       ElementC,
+                                                       LayoutC,
+                                                       Operator,
+                                                       WarpCount::kK>::Type;
 
   /// Policy used to define MmaPipelined
-  using MmaPolicy = MmaPolicy<MmaTensorOp, MatrixShape<0, 0>,
-                                        MatrixShape<0, 0>, WarpCount::kK>;
+  using MmaPolicy = MmaPolicy<MmaTensorOp,
+                              MatrixShape<0, 0>,
+                              MatrixShape<0, 0>,
+                              WarpCount::kK>;
 };
 
 }  // namespace threadblock
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_wint2x_mma.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_wint2x_mma.h
index 1782330de8..e4684b610b 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_wint2x_mma.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/default_wint2x_mma.h
@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -31,47 +31,61 @@ namespace threadblock {
 
 template <typename ThreadblockShape, typename ElementT, int GroupSize>
 struct DefaultQuantParamsIterators {
-private:
-    static constexpr int kAlignment = 128 / sizeof_bits<ElementT>::value;
-    static_assert((ThreadblockShape::kN % kAlignment) == 0, "");
+ private:
+  static constexpr int kAlignment = 128 / sizeof_bits<ElementT>::value;
+  static_assert((ThreadblockShape::kN % kAlignment) == 0, "");
 
-    static constexpr int kRows =
-        (GroupSize == -1) ? 1 : (ThreadblockShape::kK + GroupSize - 1) / GroupSize;
-    static constexpr int kColumns = ThreadblockShape::kN;
+  static constexpr int kRows =
+      (GroupSize == -1) ? 1
+                        : (ThreadblockShape::kK + GroupSize - 1) / GroupSize;
+  static constexpr int kColumns = ThreadblockShape::kN;
 
-    using IteratorThreadMap = transform::PitchLinearStripminedThreadMap<
-        layout::PitchLinearShape<kColumns, kRows>,
-        kColumns / kAlignment, kAlignment>;
+  using IteratorThreadMap = transform::PitchLinearStripminedThreadMap<
+      layout::PitchLinearShape<kColumns, kRows>,
+      kColumns / kAlignment,
+      kAlignment>;
 
-public:
-    using Iterator = cutlass::transform::threadblock::PredicatedTileIterator<
-        MatrixShape<kRows, kColumns>, ElementT, layout::RowMajor, 0,
-        IteratorThreadMap, kAlignment>;
-    using SmemIterator = Iterator;
+ public:
+  using Iterator = cutlass::transform::threadblock::PredicatedTileIterator<
+      MatrixShape<kRows, kColumns>,
+      ElementT,
+      layout::RowMajor,
+      0,
+      IteratorThreadMap,
+      kAlignment>;
+  using SmemIterator = Iterator;
 };
 
 template <typename ThreadblockShape, int GroupSize>
 struct DefaultQuantParamsIterators<ThreadblockShape, uint4b_t, GroupSize> {
-private:
-    static constexpr int kAlignment = 32 / sizeof_bits<uint4b_t>::value;
-    static_assert((ThreadblockShape::kN % kAlignment) == 0, "");
+ private:
+  static constexpr int kAlignment = 32 / sizeof_bits<uint4b_t>::value;
+  static_assert((ThreadblockShape::kN % kAlignment) == 0, "");
 
-    static constexpr int kRows =
-        (GroupSize == -1) ? 1 : (ThreadblockShape::kK + 2 * GroupSize - 1) / (2 * GroupSize);
-    static constexpr int kColumns =
-        (GroupSize == -1) ? ThreadblockShape::kN : ThreadblockShape::kN * 2;
+  static constexpr int kRows =
+      (GroupSize == -1)
+          ? 1
+          : (ThreadblockShape::kK + 2 * GroupSize - 1) / (2 * GroupSize);
+  static constexpr int kColumns =
+      (GroupSize == -1) ? ThreadblockShape::kN : ThreadblockShape::kN * 2;
 
-    using IteratorThreadMap = transform::PitchLinearStripminedThreadMap<
-        layout::PitchLinearShape<kColumns, kRows>,
-        kColumns / kAlignment, kAlignment>;
+  using IteratorThreadMap = transform::PitchLinearStripminedThreadMap<
+      layout::PitchLinearShape<kColumns, kRows>,
+      kColumns / kAlignment,
+      kAlignment>;
 
-public:
-    using AccessType = cutlass::Array<uint4b_t, kAlignment>;
-    using Iterator = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        MatrixShape<kRows, kColumns>, uint4b_t, layout::RowMajor,
-        0, IteratorThreadMap, AccessType>;
+ public:
+  using AccessType = cutlass::Array<uint4b_t, kAlignment>;
+  using Iterator =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          MatrixShape<kRows, kColumns>,
+          uint4b_t,
+          layout::RowMajor,
+          0,
+          IteratorThreadMap,
+          AccessType>;
 
-    using SmemIterator = Iterator;
+  using SmemIterator = Iterator;
 };
 
 template <
@@ -142,105 +156,174 @@ template <
     typename Operator,
     /// Use zfill or predicate for out-of-bound cp.async
     SharedMemoryClearOption SharedMemoryClear>
-struct DefaultWint2xMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementAccumulator,
-    layout::RowMajor, OperatorClass, ArchTag, ThreadblockShape, WarpShape, InstructionShape,
-    kStages, Operator, SharedMemoryClear>
-{
-public:
-    static_assert(platform::is_same<ElementA, half_t>::value || platform::is_same<ElementA, bfloat16_t>::value,
-        "Element A must be fp16 or bf16");
+struct DefaultWint2xMma<ElementA,
+                        LayoutA,
+                        kAlignmentA,
+                        ElementB,
+                        LayoutB,
+                        kAlignmentB,
+                        ElementAccumulator,
+                        layout::RowMajor,
+                        OperatorClass,
+                        ArchTag,
+                        ThreadblockShape,
+                        WarpShape,
+                        InstructionShape,
+                        kStages,
+                        Operator,
+                        SharedMemoryClear> {
+ public:
+  static_assert(platform::is_same<ElementA, half_t>::value ||
+                    platform::is_same<ElementA, bfloat16_t>::value,
+                "Element A must be fp16 or bf16");
 
-    static_assert(platform::is_same<ElementB, uint2b_t>::value,
-        "Element B must be uint2b_t");
+  static_assert(platform::is_same<ElementB, uint2b_t>::value,
+                "Element B must be uint2b_t");
 
-    static_assert(platform::is_same<Operator, arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
-        "Mma multistage must dequantize after ldsm");
+  static_assert(
+      platform::is_same<Operator,
+                        arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
+      "Mma multistage must dequantize after ldsm");
 
-    using ElementSuperScale = ElementA;
-    using ElementLocalScale = uint4b_t;
-    using ElementCodeScaleZp = float;
+  using ElementSuperScale = ElementA;
+  using ElementLocalScale = uint4b_t;
+  using ElementCodeScaleZp = float;
 
-    static constexpr int kGroupSize = 64;
+  static constexpr int kGroupSize = 64;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
-        ? cutlass::arch::CacheOperation::Global
-        : cutlass::arch::CacheOperation::Always;
+  static cutlass::arch::CacheOperation::Kind const CacheOpA =
+      ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpB = ((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
-        ? cutlass::arch::CacheOperation::Global
-        : cutlass::arch::CacheOperation::Always;
+  static cutlass::arch::CacheOperation::Kind const CacheOpB =
+      ((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    // Define the MmaCore components
-    // Mma core does not depend on stages, so pass in at least 3 here to mma multistage pieces are created
-    using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape, WarpShape, InstructionShape,
-        ElementA, LayoutA, ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, OperatorClass,
-        std::max(kStages, 3), Operator, false, CacheOpA, CacheOpB>;
+  // Define the MmaCore components
+  // Mma core does not depend on stages, so pass in at least 3 here to mma
+  // multistage pieces are created
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          ElementA,
+                                                          LayoutA,
+                                                          ElementB,
+                                                          layout::ColumnMajor,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          OperatorClass,
+                                                          std::max(kStages, 3),
+                                                          Operator,
+                                                          false,
+                                                          CacheOpA,
+                                                          CacheOpB>;
 
-    // Define iterators over tiles from the A operand
-    using ThreadMapA = typename MmaCore::IteratorThreadMapA;
-    using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
-    using IteratorA = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>, ElementA, LayoutA, 1, ThreadMapA,
-        AccessTypeA>;
+  // Define iterators over tiles from the A operand
+  using ThreadMapA = typename MmaCore::IteratorThreadMapA;
+  using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
+  using IteratorA =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
+          ElementA,
+          LayoutA,
+          1,
+          ThreadMapA,
+          AccessTypeA>;
 
-private:
-    static constexpr int kColumnsInterleaved = LayoutB::kColumnsInterleaved;
-    static constexpr int kRowsPerTile = LayoutB::kRowsPerTile;
-    static_assert(!(MmaCore::Shape::kN % kColumnsInterleaved), "ThreadblockShape must be disivle by kColumnsInterleaved");
-    static_assert(kRowsPerTile == MmaCore::Shape::kK, "");
+ private:
+  static constexpr int kColumnsInterleaved = LayoutB::kColumnsInterleaved;
+  static constexpr int kRowsPerTile = LayoutB::kRowsPerTile;
+  static_assert(!(MmaCore::Shape::kN % kColumnsInterleaved),
+                "ThreadblockShape must be disivle by kColumnsInterleaved");
+  static_assert(kRowsPerTile == MmaCore::Shape::kK, "");
 
-    using ThreadMapB = typename MmaCore::IteratorThreadMapB;
-    using WarpArrangement = typename ThreadMapB::Detail::WarpThreadArrangement;
-    static_assert(!(WarpArrangement::kStrided % kColumnsInterleaved), "");
+  using ThreadMapB = typename MmaCore::IteratorThreadMapB;
+  using WarpArrangement = typename ThreadMapB::Detail::WarpThreadArrangement;
+  static_assert(!(WarpArrangement::kStrided % kColumnsInterleaved), "");
 
-    using IteratorShapeB = MatrixShape<
-        MmaCore::Shape::kK * kColumnsInterleaved, MmaCore::Shape::kN / kColumnsInterleaved>;
-    using InterleavedThreadMapB = transform::PitchLinearWarpRakedThreadMap<
-        layout::PitchLinearShape<IteratorShapeB::kRow, IteratorShapeB::kColumn>,
-        ThreadMapB::kThreads,
-        layout::PitchLinearShape<WarpArrangement::kContiguous * kColumnsInterleaved,
-            WarpArrangement::kStrided / kColumnsInterleaved>,
-        MmaCore::kAccessSizeInBits / sizeof_bits<ElementB>::value>;
+  using IteratorShapeB = MatrixShape<MmaCore::Shape::kK * kColumnsInterleaved,
+                                     MmaCore::Shape::kN / kColumnsInterleaved>;
+  using InterleavedThreadMapB = transform::PitchLinearWarpRakedThreadMap<
+      layout::PitchLinearShape<IteratorShapeB::kRow, IteratorShapeB::kColumn>,
+      ThreadMapB::kThreads,
+      layout::PitchLinearShape<WarpArrangement::kContiguous *
+                                   kColumnsInterleaved,
+                               WarpArrangement::kStrided / kColumnsInterleaved>,
+      MmaCore::kAccessSizeInBits / sizeof_bits<ElementB>::value>;
 
-public:
-    // Define iterators over tiles from the B operand
-    using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
-    using IteratorB = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        IteratorShapeB, ElementB, layout::ColumnMajor, 0, InterleavedThreadMapB,
-        AccessTypeB>;
+ public:
+  // Define iterators over tiles from the B operand
+  using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
+  using IteratorB =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          IteratorShapeB,
+          ElementB,
+          layout::ColumnMajor,
+          0,
+          InterleavedThreadMapB,
+          AccessTypeB>;
 
-private:
-    // Define iterators over tiles from extra quant params for B operand
-    using IteratorSuperScale = typename DefaultQuantParamsIterators<
-        ThreadblockShape, ElementSuperScale, -1>::Iterator;
-    using SmemIteratorSuperScale = typename DefaultQuantParamsIterators<
-        ThreadblockShape, ElementSuperScale, -1>::SmemIterator;
+ private:
+  // Define iterators over tiles from extra quant params for B operand
+  using IteratorSuperScale =
+      typename DefaultQuantParamsIterators<ThreadblockShape,
+                                           ElementSuperScale,
+                                           -1>::Iterator;
+  using SmemIteratorSuperScale =
+      typename DefaultQuantParamsIterators<ThreadblockShape,
+                                           ElementSuperScale,
+                                           -1>::SmemIterator;
 
-    using IteratorLocalScale = typename DefaultQuantParamsIterators<
-        ThreadblockShape, ElementLocalScale, kGroupSize>::Iterator;
-    using SmemIteratorLocalScale = typename DefaultQuantParamsIterators<
-        ThreadblockShape, ElementLocalScale, kGroupSize>::SmemIterator;
+  using IteratorLocalScale =
+      typename DefaultQuantParamsIterators<ThreadblockShape,
+                                           ElementLocalScale,
+                                           kGroupSize>::Iterator;
+  using SmemIteratorLocalScale =
+      typename DefaultQuantParamsIterators<ThreadblockShape,
+                                           ElementLocalScale,
+                                           kGroupSize>::SmemIterator;
 
-    using IteratorCodeScaleZp = typename DefaultQuantParamsIterators<
-        ThreadblockShape, ElementCodeScaleZp, -1>::Iterator;
-    using SmemIteratorCodeScaleZp = typename DefaultQuantParamsIterators<
-        ThreadblockShape, ElementCodeScaleZp, -1>::Iterator;
+  using IteratorCodeScaleZp =
+      typename DefaultQuantParamsIterators<ThreadblockShape,
+                                           ElementCodeScaleZp,
+                                           -1>::Iterator;
+  using SmemIteratorCodeScaleZp =
+      typename DefaultQuantParamsIterators<ThreadblockShape,
+                                           ElementCodeScaleZp,
+                                           -1>::Iterator;
 
-public:
-    using QuantParamsAccessor = Wint2ParamsAccessor<
-        ElementA, ThreadblockShape, IteratorSuperScale, SmemIteratorSuperScale,
-        IteratorLocalScale, SmemIteratorLocalScale,
-        IteratorCodeScaleZp, SmemIteratorCodeScaleZp, kStages, kGroupSize>;
+ public:
+  using QuantParamsAccessor = Wint2ParamsAccessor<ElementA,
+                                                  ThreadblockShape,
+                                                  IteratorSuperScale,
+                                                  SmemIteratorSuperScale,
+                                                  IteratorLocalScale,
+                                                  SmemIteratorLocalScale,
+                                                  IteratorCodeScaleZp,
+                                                  SmemIteratorCodeScaleZp,
+                                                  kStages,
+                                                  kGroupSize>;
 
-    // Define the threadblock-scoped multistage matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::Wint2xMmaMultistage<
-        typename MmaCore::Shape,
-        IteratorA, typename MmaCore::SmemIteratorA, MmaCore::kCacheOpA,
-        IteratorB, typename MmaCore::SmemIteratorB, MmaCore::kCacheOpB,
-        ElementAccumulator, layout::RowMajor, typename MmaCore::MmaPolicy,
-        kStages, QuantParamsAccessor, SharedMemoryClear>;
+  // Define the threadblock-scoped multistage matrix multiply
+  using ThreadblockMma = cutlass::gemm::threadblock::Wint2xMmaMultistage<
+      typename MmaCore::Shape,
+      IteratorA,
+      typename MmaCore::SmemIteratorA,
+      MmaCore::kCacheOpA,
+      IteratorB,
+      typename MmaCore::SmemIteratorB,
+      MmaCore::kCacheOpB,
+      ElementAccumulator,
+      layout::RowMajor,
+      typename MmaCore::MmaPolicy,
+      kStages,
+      QuantParamsAccessor,
+      SharedMemoryClear>;
 };
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_base.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_base.h
index 1fb7f7eb28..51e410aad3 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_base.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_base.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -47,30 +48,33 @@
 
 ////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 // SFINAE trick so I can keep the same loop code for Volta and dispatch to the
-// correct warp level mma. On volta, all data is stored to shared memory as FP16.
+// correct warp level mma. On volta, all data is stored to shared memory as
+// FP16.
 template <typename WarpMma, int kExpansionFactor = 1>
-CUTLASS_DEVICE void run_warp_mma(WarpMma& warp_mma, typename WarpMma::FragmentC& D,
-    typename WarpMma::FragmentA const& A, typename WarpMma::FragmentB const& B, typename WarpMma::FragmentC const& C,
-    int const warp_tileB_k_offset)
-{
-    warp_mma(D, A, B, C);
+CUTLASS_DEVICE void run_warp_mma(WarpMma& warp_mma,
+                                 typename WarpMma::FragmentC& D,
+                                 typename WarpMma::FragmentA const& A,
+                                 typename WarpMma::FragmentB const& B,
+                                 typename WarpMma::FragmentC const& C,
+                                 int const warp_tileB_k_offset) {
+  warp_mma(D, A, B, C);
 }
 
 template <typename WarpMma, int kExpansionFactor = WarpMma::kExpansionFactor>
-CUTLASS_DEVICE void run_warp_mma(WarpMma& warp_mma, typename WarpMma::FragmentC& D,
-    typename WarpMma::TransformedFragmentA const& A, typename WarpMma::TransformedFragmentB const& B,
-    typename WarpMma::FragmentC const& C, int const warp_tileB_k_offset)
-{
-    warp_mma(D, A, B, C, warp_tileB_k_offset);
+CUTLASS_DEVICE void run_warp_mma(
+    WarpMma& warp_mma,
+    typename WarpMma::FragmentC& D,
+    typename WarpMma::TransformedFragmentA const& A,
+    typename WarpMma::TransformedFragmentB const& B,
+    typename WarpMma::FragmentC const& C,
+    int const warp_tileB_k_offset) {
+  warp_mma(D, A, B, C, warp_tileB_k_offset);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -90,168 +94,169 @@ template <
     WeightOnlyQuantOp DequantOp,
     /// Used for partial specialization,
     typename Enable = bool>
-class DqMmaBase
-{
-public:
-    ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using Shape = Shape_;
+class DqMmaBase {
+ public:
+  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using Shape = Shape_;
 
-    ///< Policy describing tuning details
-    using Policy = Policy_;
+  ///< Policy describing tuning details
+  using Policy = Policy_;
 
-    ///< Type of the scale to be loaded
-    using ElementScale = ElementScale_;
+  ///< Type of the scale to be loaded
+  using ElementScale = ElementScale_;
 
-    static_assert(DequantOp != WeightOnlyQuantOp::UNDEFINED, "");
+  static_assert(DequantOp != WeightOnlyQuantOp::UNDEFINED, "");
 
-    // Finegrained scales get streamed in via cp.async
-    static constexpr int ScalebiasStages = isFinegrained(DequantOp) ? Stages : 1;
-    // We always have scales.
-    static constexpr int ScaleElementsPerStage = Shape::kN;
-    // We sometimes have a bias
-    static constexpr int BiasElementsPerStage = hasZero(DequantOp) ? Shape::kN : 0;
+  // Finegrained scales get streamed in via cp.async
+  static constexpr int ScalebiasStages = isFinegrained(DequantOp) ? Stages : 1;
+  // We always have scales.
+  static constexpr int ScaleElementsPerStage = Shape::kN;
+  // We sometimes have a bias
+  static constexpr int BiasElementsPerStage =
+      hasZero(DequantOp) ? Shape::kN : 0;
 
+  //
+  // Dependent types
+  //
+
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
+
+  /// Shape describing the overall GEMM computed from shared memory
+  /// by each warp.
+  using WarpGemm = typename Policy::Operator::Shape;
+
+  /// Shape describing the number of warps filling the CTA
+  using WarpCount = GemmShape<Shape::kM / WarpGemm::kM,
+                              Shape::kN / WarpGemm::kN,
+                              Shape::kK / WarpGemm::kK>;
+
+  /// Number of warp-level GEMM operations
+  static int const kWarpGemmIterations =
+      (WarpGemm::kK / Operator::Policy::MmaShape::kK);
+
+  static constexpr int kNumKIterationsPerWarpBLoad =
+      Operator::IteratorB::InstructionShape::kRow /
+      Operator::InstructionShape::kK;
+
+  static_assert(!(kWarpGemmIterations % kNumKIterationsPerWarpBLoad), "");
+  static constexpr int kWarpGemmIterationsForB =
+      kWarpGemmIterations / kNumKIterationsPerWarpBLoad;
+
+  /// Number of stages
+  static int const kStages = Stages;
+
+  /// Tensor reference to the A operand
+  using TensorRefA =
+      TensorRef<typename Operator::ElementA, typename Operator::LayoutA>;
+
+  /// Tensor reference to the B operand
+  using TensorRefB =
+      TensorRef<typename Operator::ElementB, typename Operator::LayoutB>;
+
+  //
+  // Nested structs
+  //
+
+  /// Shared storage object needed by threadblock-scoped GEMM
+  class SharedStorage {
+   public:
     //
-    // Dependent types
+    // Type definitions
     //
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+    /// Shape of the A matrix operand in shared memory
+    using ShapeA =
+        MatrixShape<Shape::kM + Policy::SmemPaddingA::kRow,
+                    Shape::kK * kStages + Policy::SmemPaddingA::kColumn>;
 
-    /// Shape describing the overall GEMM computed from shared memory
-    /// by each warp.
-    using WarpGemm = typename Policy::Operator::Shape;
+    /// Shape of the B matrix operand in shared memory
+    using ShapeB = MatrixShape<Shape::kK * kStages + Policy::SmemPaddingB::kRow,
+                               Shape::kN + Policy::SmemPaddingB::kColumn>;
 
-    /// Shape describing the number of warps filling the CTA
-    using WarpCount = GemmShape<Shape::kM / WarpGemm::kM, Shape::kN / WarpGemm::kN, Shape::kK / WarpGemm::kK>;
+    /// Shape of the shared memory buffer for the scales for the B matrix.
+    using ShapeScale = MatrixShape<ScalebiasStages, ScaleElementsPerStage>;
+    /// Shape of the shared memory buffer for the biases of the B matrix.
+    using ShapeZero = MatrixShape<ScalebiasStages, BiasElementsPerStage>;
 
-    /// Number of warp-level GEMM operations
-    static int const kWarpGemmIterations = (WarpGemm::kK / Operator::Policy::MmaShape::kK);
-
-    static constexpr int kNumKIterationsPerWarpBLoad
-        = Operator::IteratorB::InstructionShape::kRow / Operator::InstructionShape::kK;
-
-    static_assert(!(kWarpGemmIterations % kNumKIterationsPerWarpBLoad), "");
-    static constexpr int kWarpGemmIterationsForB = kWarpGemmIterations / kNumKIterationsPerWarpBLoad;
-
-    /// Number of stages
-    static int const kStages = Stages;
-
-    /// Tensor reference to the A operand
-    using TensorRefA = TensorRef<typename Operator::ElementA, typename Operator::LayoutA>;
-
-    /// Tensor reference to the B operand
-    using TensorRefB = TensorRef<typename Operator::ElementB, typename Operator::LayoutB>;
-
-    //
-    // Nested structs
-    //
-
-    /// Shared storage object needed by threadblock-scoped GEMM
-    class SharedStorage
-    {
-    public:
-        //
-        // Type definitions
-        //
-
-        /// Shape of the A matrix operand in shared memory
-        using ShapeA
-            = MatrixShape<Shape::kM + Policy::SmemPaddingA::kRow, Shape::kK * kStages + Policy::SmemPaddingA::kColumn>;
-
-        /// Shape of the B matrix operand in shared memory
-        using ShapeB
-            = MatrixShape<Shape::kK * kStages + Policy::SmemPaddingB::kRow, Shape::kN + Policy::SmemPaddingB::kColumn>;
-
-        /// Shape of the shared memory buffer for the scales for the B matrix.
-        using ShapeScale = MatrixShape<ScalebiasStages, ScaleElementsPerStage>;
-        /// Shape of the shared memory buffer for the biases of the B matrix.
-        using ShapeZero = MatrixShape<ScalebiasStages, BiasElementsPerStage>;
-
-    public:
-        //
-        // Data members
-        //
-
-        /// Buffer for A operand
-        AlignedBuffer<typename Operator::ElementA, ShapeA::kCount> operand_A;
-
-        /// Buffer for B operand
-        AlignedBuffer<typename Operator::ElementB, ShapeB::kCount> operand_B;
-
-        /// Buffer to hold scales for threadblock
-        AlignedBuffer<ElementScale, ShapeScale::kCount> operand_scale;
-
-        /// Buffer to hold scales for threadblock
-        AlignedBuffer<ElementScale, ShapeZero::kCount> operand_zero;
-
-    public:
-        //
-        // Methods
-        //
-
-        /// Returns a layout object for the A matrix
-        CUTLASS_DEVICE
-        static typename Operator::LayoutA LayoutA()
-        {
-            return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn});
-        }
-
-        /// Returns a layout object for the B matrix
-        CUTLASS_HOST_DEVICE
-        static typename Operator::LayoutB LayoutB()
-        {
-            return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn});
-        }
-
-        /// Returns a TensorRef to the A operand
-        CUTLASS_HOST_DEVICE
-        TensorRefA operand_A_ref()
-        {
-            return TensorRefA{operand_A.data(), LayoutA()};
-        }
-
-        /// Returns a TensorRef to the B operand
-        CUTLASS_HOST_DEVICE
-        TensorRefB operand_B_ref()
-        {
-            return TensorRefB{operand_B.data(), LayoutB()};
-        }
-    };
-
-protected:
+   public:
     //
     // Data members
     //
 
-    /// Iterator to load a warp-scoped tile of A operand from shared memory
-    typename Operator::IteratorA warp_tile_iterator_A_;
+    /// Buffer for A operand
+    AlignedBuffer<typename Operator::ElementA, ShapeA::kCount> operand_A;
 
-    /// Iterator to load a warp-scoped tile of B operand from shared memory
-    typename Operator::IteratorB warp_tile_iterator_B_;
+    /// Buffer for B operand
+    AlignedBuffer<typename Operator::ElementB, ShapeB::kCount> operand_B;
 
-public:
-    /// Construct from tensor references
+    /// Buffer to hold scales for threadblock
+    AlignedBuffer<ElementScale, ShapeScale::kCount> operand_scale;
+
+    /// Buffer to hold scales for threadblock
+    AlignedBuffer<ElementScale, ShapeZero::kCount> operand_zero;
+
+   public:
+    //
+    // Methods
+    //
+
+    /// Returns a layout object for the A matrix
     CUTLASS_DEVICE
-    DqMmaBase(
-        ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        SharedStorage& shared_storage,
-        ///< ID within the threadblock
-        int thread_idx,
-        ///< ID of warp
-        int warp_idx,
-        ///< ID of each thread within a warp
-        int lane_idx)
-        : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx)
-        , warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx)
-    {
+    static typename Operator::LayoutA LayoutA() {
+      return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn});
     }
+
+    /// Returns a layout object for the B matrix
+    CUTLASS_HOST_DEVICE
+    static typename Operator::LayoutB LayoutB() {
+      return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn});
+    }
+
+    /// Returns a TensorRef to the A operand
+    CUTLASS_HOST_DEVICE
+    TensorRefA operand_A_ref() {
+      return TensorRefA{operand_A.data(), LayoutA()};
+    }
+
+    /// Returns a TensorRef to the B operand
+    CUTLASS_HOST_DEVICE
+    TensorRefB operand_B_ref() {
+      return TensorRefB{operand_B.data(), LayoutB()};
+    }
+  };
+
+ protected:
+  //
+  // Data members
+  //
+
+  /// Iterator to load a warp-scoped tile of A operand from shared memory
+  typename Operator::IteratorA warp_tile_iterator_A_;
+
+  /// Iterator to load a warp-scoped tile of B operand from shared memory
+  typename Operator::IteratorB warp_tile_iterator_B_;
+
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  DqMmaBase(
+      ///< Shared storage needed for internal use by threadblock-scoped GEMM
+      SharedStorage& shared_storage,
+      ///< ID within the threadblock
+      int thread_idx,
+      ///< ID of warp
+      int warp_idx,
+      ///< ID of each thread within a warp
+      int lane_idx)
+      : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx),
+        warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {}
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage.h
index 3c4036dd8c..4a4a3137be 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -48,12 +49,9 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -102,9 +100,9 @@ template <
     typename Enable = void>
 class DqMmaMultistage;
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
 
 #include "cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h"
 #include "cutlass_extensions/gemm/threadblock/dq_mma_multistage_percol.h"
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h
index e87a51b22c..02b3b11840 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -48,12 +49,9 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -98,460 +96,587 @@ template <
     WeightOnlyQuantOp QuantOp_,
     /// Use zfill or predicate for out-of-bound cp.async
     SharedMemoryClearOption SharedMemoryClear>
-class DqMmaMultistage<Shape_, IteratorA_, SmemIteratorA_, CacheOpA, IteratorB_, SmemIteratorB_, CacheOpB,
-    IteratorScale_, SmemIteratorScale_, ElementC_, LayoutC_, Policy_, Stages, TransformBAfterLDS_, QuantOp_,
-    SharedMemoryClear, std::enable_if_t<isFinegrained(QuantOp_)>>
-    : public DqMmaBase<Shape_, Policy_, typename IteratorScale_::Element, Stages, QuantOp_>
-{
-public:
-    ///< Base class
-    using Base = DqMmaBase<Shape_, Policy_, typename IteratorScale_::Element, Stages, QuantOp_>;
-    ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using Shape = Shape_;
-    ///< Iterates over tiles of A operand in global memory
-    using IteratorA = IteratorA_;
-    ///< Iterates over tiles of B operand in global memory
-    using IteratorB = IteratorB_;
-    ///< Data type of accumulator matrix
-    using ElementC = ElementC_;
-    ///< Layout of accumulator matrix
-    using LayoutC = LayoutC_;
-    ///< Policy describing tuning details
-    using Policy = Policy_;
+class DqMmaMultistage<Shape_,
+                      IteratorA_,
+                      SmemIteratorA_,
+                      CacheOpA,
+                      IteratorB_,
+                      SmemIteratorB_,
+                      CacheOpB,
+                      IteratorScale_,
+                      SmemIteratorScale_,
+                      ElementC_,
+                      LayoutC_,
+                      Policy_,
+                      Stages,
+                      TransformBAfterLDS_,
+                      QuantOp_,
+                      SharedMemoryClear,
+                      std::enable_if_t<isFinegrained(QuantOp_)>>
+    : public DqMmaBase<Shape_,
+                       Policy_,
+                       typename IteratorScale_::Element,
+                       Stages,
+                       QuantOp_> {
+ public:
+  ///< Base class
+  using Base = DqMmaBase<Shape_,
+                         Policy_,
+                         typename IteratorScale_::Element,
+                         Stages,
+                         QuantOp_>;
+  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using Shape = Shape_;
+  ///< Iterates over tiles of A operand in global memory
+  using IteratorA = IteratorA_;
+  ///< Iterates over tiles of B operand in global memory
+  using IteratorB = IteratorB_;
+  ///< Data type of accumulator matrix
+  using ElementC = ElementC_;
+  ///< Layout of accumulator matrix
+  using LayoutC = LayoutC_;
+  ///< Policy describing tuning details
+  using Policy = Policy_;
 
-    using IteratorScale = IteratorScale_;
-    using ElementScale = typename IteratorScale::Element;
-    using LayoutScale = typename IteratorScale::Layout;
+  using IteratorScale = IteratorScale_;
+  using ElementScale = typename IteratorScale::Element;
+  using LayoutScale = typename IteratorScale::Layout;
 
-    using SmemIteratorA = SmemIteratorA_;
-    using SmemIteratorB = SmemIteratorB_;
-    using SmemIteratorScale = SmemIteratorScale_;
+  using SmemIteratorA = SmemIteratorA_;
+  using SmemIteratorB = SmemIteratorB_;
+  using SmemIteratorScale = SmemIteratorScale_;
 
-    static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
-    static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
+  static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
+  static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
 
-    using TransformBAfterLDS = TransformBAfterLDS_;
+  using TransformBAfterLDS = TransformBAfterLDS_;
 
-    static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
-    //
-    // Dependent types
-    //
+  static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
+  //
+  // Dependent types
+  //
 
-    /// Fragment of accumulator tile
-    using FragmentC = typename Policy::Operator::FragmentC;
+  /// Fragment of accumulator tile
+  using FragmentC = typename Policy::Operator::FragmentC;
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
 
-    /// Minimum architecture is Sm80 to support cp.async
-    using ArchTag = arch::Sm80;
+  /// Minimum architecture is Sm80 to support cp.async
+  using ArchTag = arch::Sm80;
 
-    using Dequantizer = warp::MmaTensorOpDequantizer<Operator, typename Base::WarpGemm, Operand::kB, ElementScale,
-        LayoutScale, 32, QuantOp>;
+  using Dequantizer = warp::MmaTensorOpDequantizer<Operator,
+                                                   typename Base::WarpGemm,
+                                                   Operand::kB,
+                                                   ElementScale,
+                                                   LayoutScale,
+                                                   32,
+                                                   QuantOp>;
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = Operator::kTransformA;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = Operator::kTransformA;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = Operator::kTransformB;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = Operator::kTransformB;
 
-    static_assert(Base::SharedStorage::ShapeScale::kRow == Stages, "");
-    static_assert(Base::SharedStorage::ShapeScale::kColumn == Shape::kN, "");
+  static_assert(Base::SharedStorage::ShapeScale::kRow == Stages, "");
+  static_assert(Base::SharedStorage::ShapeScale::kColumn == Shape::kN, "");
 
-    /// Internal structure exposed for introspection.
-    struct Detail
-    {
+  /// Internal structure exposed for introspection.
+  struct Detail {
+    static_assert(Base::kWarpGemmIterations > 1,
+                  "The pipelined structure requires at least two warp-level "
+                  "GEMM operations.");
 
-        static_assert(Base::kWarpGemmIterations > 1,
-            "The pipelined structure requires at least two warp-level "
-            "GEMM operations.");
+    /// Number of cp.async instructions to load one stage of operand A
+    static int const AsyncCopyIterationsPerStageA =
+        IteratorA::ThreadMap::Iterations::kCount;
 
-        /// Number of cp.async instructions to load one stage of operand A
-        static int const AsyncCopyIterationsPerStageA = IteratorA::ThreadMap::Iterations::kCount;
+    /// Number of cp.async instructions to load one stage of operand B
+    static int const AsyncCopyIterationsPerStageB =
+        IteratorB::ThreadMap::Iterations::kCount;
 
-        /// Number of cp.async instructions to load one stage of operand B
-        static int const AsyncCopyIterationsPerStageB = IteratorB::ThreadMap::Iterations::kCount;
+    /// Number of stages
+    static int const kStages = Stages;
 
-        /// Number of stages
-        static int const kStages = Stages;
+    /// Number of cp.async instructions to load on group of operand A
+    static int const kAccessesPerGroupA =
+        (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
 
-        /// Number of cp.async instructions to load on group of operand A
-        static int const kAccessesPerGroupA
-            = (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
+    /// Number of cp.async instructions to load on group of operand B
+    static int const kAccessesPerGroupB =
+        (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
+  };
 
-        /// Number of cp.async instructions to load on group of operand B
-        static int const kAccessesPerGroupB
-            = (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
-    };
+ private:
+  using WarpFragmentA = typename Operator::FragmentA;
+  using WarpFragmentB = typename Operator::FragmentB;
+  Dequantizer warp_dequantizer_;
 
-private:
-    using WarpFragmentA = typename Operator::FragmentA;
-    using WarpFragmentB = typename Operator::FragmentB;
-    Dequantizer warp_dequantizer_;
+  using ElementA = typename IteratorA::Element;
+  using ElementB = typename IteratorB::Element;
+  using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
 
-    using ElementA = typename IteratorA::Element;
-    using ElementB = typename IteratorB::Element;
-    using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
+  static constexpr bool RequiresTileInterleave =
+      layout::IsColumnMajorTileInterleave<
+          typename LayoutDetailsForB::Layout>::value;
+  static_assert(!RequiresTileInterleave ||
+                    (RequiresTileInterleave &&
+                     (Shape::kK == LayoutDetailsForB::ThreadblockK)),
+                "Layout K must match threadblockK");
 
-    static constexpr bool RequiresTileInterleave
-        = layout::IsColumnMajorTileInterleave<typename LayoutDetailsForB::Layout>::value;
-    static_assert(!RequiresTileInterleave || (RequiresTileInterleave && (Shape::kK == LayoutDetailsForB::ThreadblockK)),
-        "Layout K must match threadblockK");
+ private:
+  //
+  // Data members
+  //
 
-private:
-    //
-    // Data members
-    //
+  /// Iterator to write threadblock-scoped tile of A operand to shared memory
+  SmemIteratorA smem_iterator_A_;
 
-    /// Iterator to write threadblock-scoped tile of A operand to shared memory
-    SmemIteratorA smem_iterator_A_;
+  /// Iterator to write threadblock-scoped tile of B operand to shared memory
+  SmemIteratorB smem_iterator_B_;
 
-    /// Iterator to write threadblock-scoped tile of B operand to shared memory
-    SmemIteratorB smem_iterator_B_;
+  /// Iterator to write threadblock-scoped tile of scale and zero operand to
+  /// shared memory
+  SmemIteratorScale smem_iterator_scale_;
 
-    /// Iterator to write threadblock-scoped tile of scale and zero operand to shared memory
-    SmemIteratorScale smem_iterator_scale_;
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  DqMmaMultistage(
+      ///< Shared storage needed for internal use by threadblock-scoped GEMM
+      typename Base::SharedStorage& shared_storage,
+      /// The group size for quantization
+      int const group_size,
+      ///< ID within the threadblock
+      int thread_idx,
+      ///< ID of warp
+      int warp_idx,
+      ///< ID of each thread within a warp
+      int lane_idx)
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
+        warp_dequantizer_(
+            {shared_storage.operand_scale.data(), LayoutScale(Shape::kN)},
+            {shared_storage.operand_zero.data(), LayoutScale(Shape::kN)},
+            (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) /
+                Base::WarpCount::kM,
+            lane_idx),
+        smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
+        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx),
+        smem_iterator_scale_(LayoutScale(Shape::kN),
+                             shared_storage.operand_scale.data(),
+                             shared_storage.operand_zero.data(),
+                             {Base::kStages, Shape::kN},
+                             thread_idx,
+                             group_size) {
+    // Compute warp location within threadblock tile by mapping the warp_id to
+    // three coordinates:
+    //   _m: the warp's position within the threadblock along the M dimension
+    //   _n: the warp's position within the threadblock along the N dimension
+    //   _k: the warp's position within the threadblock along the K dimension
 
-public:
-    /// Construct from tensor references
-    CUTLASS_DEVICE
-    DqMmaMultistage(
-        ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        typename Base::SharedStorage& shared_storage,
-        /// The group size for quantization
-        int const group_size,
-        ///< ID within the threadblock
-        int thread_idx,
-        ///< ID of warp
-        int warp_idx,
-        ///< ID of each thread within a warp
-        int lane_idx)
-        : Base(shared_storage, thread_idx, warp_idx, lane_idx)
-        , warp_dequantizer_({shared_storage.operand_scale.data(), LayoutScale(Shape::kN)},
-              {shared_storage.operand_zero.data(), LayoutScale(Shape::kN)},
-              (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) / Base::WarpCount::kM, lane_idx)
-        , smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx)
-        , smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx)
-        , smem_iterator_scale_(LayoutScale(Shape::kN), shared_storage.operand_scale.data(),
-              shared_storage.operand_zero.data(), {Base::kStages, Shape::kN}, thread_idx, group_size)
-    {
-        // Compute warp location within threadblock tile by mapping the warp_id to
-        // three coordinates:
-        //   _m: the warp's position within the threadblock along the M dimension
-        //   _n: the warp's position within the threadblock along the N dimension
-        //   _k: the warp's position within the threadblock along the K dimension
+    int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
 
-        int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
-        int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
+    int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
 
-        int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
-        int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
+    // Add per-warp offsets in units of warp-level tiles
+    this->warp_tile_iterator_A_.add_tile_offset(
+        {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
+    this->warp_tile_iterator_B_.add_tile_offset(
+        {Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  }
 
-        // Add per-warp offsets in units of warp-level tiles
-        this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
-        this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  CUTLASS_DEVICE
+  void copy_scales_and_advance(IteratorScale& iterator_scale,
+                               int stage = -1,
+                               int k_iter = -1) {
+    static_assert(IteratorScale::Shape::kRow == 1, "Scale stride must be 1.");
+
+    typename IteratorScale::AccessType* gmem_scale_ptr =
+        iterator_scale.get_scale();
+    typename IteratorScale::AccessType* gmem_zero_ptr =
+        iterator_scale.get_zero();
+
+    typename IteratorScale::AccessType* smem_scale_ptr =
+        reinterpret_cast<typename IteratorScale::AccessType*>(
+            this->smem_iterator_scale_.get_scale());
+    typename IteratorScale::AccessType* smem_zero_ptr =
+        reinterpret_cast<typename IteratorScale::AccessType*>(
+            this->smem_iterator_scale_.get_zero());
+
+    int const kSrcBytes = sizeof_bits<typename IteratorScale::Element>::value *
+                          IteratorScale::kAlignment / 8;
+
+    cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(
+        smem_scale_ptr, gmem_scale_ptr, iterator_scale.valid());
+
+    if (gmem_zero_ptr != nullptr) {
+      cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(
+          smem_zero_ptr, gmem_zero_ptr, iterator_scale.valid());
     }
 
-    CUTLASS_DEVICE
-    void copy_scales_and_advance(IteratorScale& iterator_scale, int stage = -1, int k_iter = -1)
-    {
-        static_assert(IteratorScale::Shape::kRow == 1, "Scale stride must be 1.");
-
-        typename IteratorScale::AccessType* gmem_scale_ptr = iterator_scale.get_scale();
-        typename IteratorScale::AccessType* gmem_zero_ptr = iterator_scale.get_zero();
-
-        typename IteratorScale::AccessType* smem_scale_ptr
-            = reinterpret_cast<typename IteratorScale::AccessType*>(this->smem_iterator_scale_.get_scale());
-        typename IteratorScale::AccessType* smem_zero_ptr
-            = reinterpret_cast<typename IteratorScale::AccessType*>(this->smem_iterator_scale_.get_zero());
-
-        int const kSrcBytes = sizeof_bits<typename IteratorScale::Element>::value * IteratorScale::kAlignment / 8;
-
-        cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(smem_scale_ptr, gmem_scale_ptr, iterator_scale.valid());
-
-        if (gmem_zero_ptr != nullptr)
-        {
-            cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(smem_zero_ptr, gmem_zero_ptr, iterator_scale.valid());
+    if (iterator_scale.group_size_ == 64) {
+      iterator_scale.add_tile_offset({1, 0});
+    } else if (iterator_scale.group_size_ == 128) {
+      if constexpr (Shape::kK == 128) {
+        iterator_scale.add_tile_offset({1, 0});
+      } else if constexpr (Shape::kK == 64) {
+        if (iterator_scale.row_groupsize64_ & 0x1) {
+          iterator_scale.add_tile_offset({1, 0});
         }
-
-        if (iterator_scale.group_size_ == 64)
-        {
-            iterator_scale.add_tile_offset({1, 0});
-        }
-        else if (iterator_scale.group_size_ == 128)
-        {
-            if constexpr (Shape::kK == 128)
-            {
-                iterator_scale.add_tile_offset({1, 0});
-            }
-            else if constexpr (Shape::kK == 64)
-            {
-                if (iterator_scale.row_groupsize64_ & 0x1)
-                {
-                    iterator_scale.add_tile_offset({1, 0});
-                }
-            }
-            else
-            {
-                static_assert(Shape::kK == 0, "Unsupported k tile shape, can only be 64 or 128");
-            }
-        }
-
-        iterator_scale.row_groupsize64_++;
-
-        this->smem_iterator_scale_.add_tile_offset({1, 0});
+      } else {
+        static_assert(Shape::kK == 0,
+                      "Unsupported k tile shape, can only be 64 or 128");
+      }
     }
 
-    CUTLASS_DEVICE
-    void copy_tiles_and_advance(
-        IteratorA& iterator_A, IteratorB& iterator_B, int group_start_A = 0, int group_start_B = 0)
-    {
-        iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector);
-        this->smem_iterator_A_.set_iteration_index(group_start_A);
+    iterator_scale.row_groupsize64_++;
+
+    this->smem_iterator_scale_.add_tile_offset({1, 0});
+  }
+
+  CUTLASS_DEVICE
+  void copy_tiles_and_advance(IteratorA& iterator_A,
+                              IteratorB& iterator_B,
+                              int group_start_A = 0,
+                              int group_start_B = 0) {
+    iterator_A.set_iteration_index(group_start_A *
+                                   IteratorA::kAccessesPerVector);
+    this->smem_iterator_A_.set_iteration_index(group_start_A);
+
+    // Async Copy for operand A
+    CUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) {
+      if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                this->smem_iterator_A_.get());
+
+        int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
+                              IteratorA::ThreadMap::kElementsPerAccess /
+                              IteratorA::kAccessesPerVector / 8;
 
-        // Async Copy for operand A
         CUTLASS_PRAGMA_UNROLL
-        for (int j = 0; j < Detail::kAccessesPerGroupA; ++j)
-        {
-            if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA)
-            {
-                typename IteratorA::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorA::AccessType*>(this->smem_iterator_A_.get());
+        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
+          auto gmem_ptr = iterator_A.get();
 
-                int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value
-                    * IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector / 8;
+          if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
+                dst_ptr + v, gmem_ptr, iterator_A.valid());
+          } else {
+            cutlass::arch::cp_async<kSrcBytes, kCacheOpA>(
+                dst_ptr + v, gmem_ptr, iterator_A.valid());
+          }
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorA::kAccessesPerVector; ++v)
-                {
-                    auto gmem_ptr = iterator_A.get();
-
-                    if (SharedMemoryClear == SharedMemoryClearOption::kZfill)
-                    {
-                        cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(dst_ptr + v, gmem_ptr, iterator_A.valid());
-                    }
-                    else
-                    {
-                        cutlass::arch::cp_async<kSrcBytes, kCacheOpA>(dst_ptr + v, gmem_ptr, iterator_A.valid());
-                    }
-
-                    ++iterator_A;
-                }
-
-                ++this->smem_iterator_A_;
-            }
+          ++iterator_A;
         }
 
-        iterator_B.set_iteration_index(group_start_B * IteratorB::kAccessesPerVector);
-        this->smem_iterator_B_.set_iteration_index(group_start_B);
-
-        // Async Copy for operand B
-        CUTLASS_PRAGMA_UNROLL
-        for (int j = 0; j < Detail::kAccessesPerGroupB; ++j)
-        {
-            if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB)
-            {
-                typename IteratorB::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorB::AccessType*>(this->smem_iterator_B_.get());
-
-                int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value
-                    * IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector / 8;
-
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorB::kAccessesPerVector; ++v)
-                {
-                    auto gmem_ptr = iterator_B.get();
-
-                    if (SharedMemoryClear == SharedMemoryClearOption::kZfill)
-                    {
-                        cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(dst_ptr + v, gmem_ptr, iterator_B.valid());
-                    }
-                    else
-                    {
-                        cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(dst_ptr + v, gmem_ptr, iterator_B.valid());
-                    }
-
-                    ++iterator_B;
-                }
-                ++this->smem_iterator_B_;
-            }
-        }
+        ++this->smem_iterator_A_;
+      }
     }
 
-    /// Perform a threadblock-scoped matrix multiply-accumulate
-    CUTLASS_DEVICE
-    void operator()(
-        ///< problem size of GEMM
-        int gemm_k_iterations,
-        ///< destination accumulator tile
-        FragmentC& accum,
-        ///< iterator over A operand in global memory
-        IteratorA iterator_A,
-        ///< iterator over B operand in global memory
-        IteratorB iterator_B,
-        ///< iterator over scale operand in global memory
-        IteratorScale iterator_scale,
-        ///< initial value of accumulator
-        FragmentC const& src_accum)
-    {
+    iterator_B.set_iteration_index(group_start_B *
+                                   IteratorB::kAccessesPerVector);
+    this->smem_iterator_B_.set_iteration_index(group_start_B);
 
-        //
-        // Prologue
-        //
+    // Async Copy for operand B
+    CUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) {
+      if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                this->smem_iterator_B_.get());
 
-        TransformBAfterLDS lds_converter;
+        int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value *
+                              IteratorB::ThreadMap::kElementsPerAccess /
+                              IteratorB::kAccessesPerVector / 8;
 
-        // Issue several complete stages
         CUTLASS_PRAGMA_UNROLL
-        for (int stage = 0; stage < Base::kStages - 1; ++stage, --gemm_k_iterations)
-        {
+        for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
+          auto gmem_ptr = iterator_B.get();
 
-            iterator_A.clear_mask(gemm_k_iterations == 0);
-            iterator_B.clear_mask(gemm_k_iterations == 0);
-            iterator_scale.clear_mask(gemm_k_iterations == 0);
+          if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
+                dst_ptr + v, gmem_ptr, iterator_B.valid());
+          } else {
+            cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(
+                dst_ptr + v, gmem_ptr, iterator_B.valid());
+          }
 
-            iterator_A.set_iteration_index(0);
-            this->smem_iterator_A_.set_iteration_index(0);
+          ++iterator_B;
+        }
+        ++this->smem_iterator_B_;
+      }
+    }
+  }
 
-            // Async Copy for operand A
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j)
-            {
-                typename IteratorA::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorA::AccessType*>(this->smem_iterator_A_.get());
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  CUTLASS_DEVICE
+  void operator()(
+      ///< problem size of GEMM
+      int gemm_k_iterations,
+      ///< destination accumulator tile
+      FragmentC& accum,
+      ///< iterator over A operand in global memory
+      IteratorA iterator_A,
+      ///< iterator over B operand in global memory
+      IteratorB iterator_B,
+      ///< iterator over scale operand in global memory
+      IteratorScale iterator_scale,
+      ///< initial value of accumulator
+      FragmentC const& src_accum) {
+    //
+    // Prologue
+    //
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorA::kAccessesPerVector; ++v)
-                {
-                    int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value
-                        * IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector / 8;
+    TransformBAfterLDS lds_converter;
 
-                    cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
-                        dst_ptr + v, iterator_A.get(), iterator_A.valid());
+    // Issue several complete stages
+    CUTLASS_PRAGMA_UNROLL
+    for (int stage = 0; stage < Base::kStages - 1;
+         ++stage, --gemm_k_iterations) {
+      iterator_A.clear_mask(gemm_k_iterations == 0);
+      iterator_B.clear_mask(gemm_k_iterations == 0);
+      iterator_scale.clear_mask(gemm_k_iterations == 0);
 
-                    ++iterator_A;
-                }
+      iterator_A.set_iteration_index(0);
+      this->smem_iterator_A_.set_iteration_index(0);
 
-                ++this->smem_iterator_A_;
-            }
+      // Async Copy for operand A
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                this->smem_iterator_A_.get());
 
-            iterator_B.set_iteration_index(0);
-            this->smem_iterator_B_.set_iteration_index(0);
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
+          int const kSrcBytes =
+              sizeof_bits<typename IteratorA::Element>::value *
+              IteratorA::ThreadMap::kElementsPerAccess /
+              IteratorA::kAccessesPerVector / 8;
 
-            // Async Copy for operand B
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j)
-            {
-                typename IteratorB::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorB::AccessType*>(this->smem_iterator_B_.get());
+          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
+              dst_ptr + v, iterator_A.get(), iterator_A.valid());
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorB::kAccessesPerVector; ++v)
-                {
-                    int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value
-                        * IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector / 8;
-
-                    cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
-                        dst_ptr + v, iterator_B.get(), iterator_B.valid());
-
-                    ++iterator_B;
-                }
-
-                ++this->smem_iterator_B_;
-            }
-
-            copy_scales_and_advance(iterator_scale, stage, gemm_k_iterations);
-
-            // Move to the next stage
-            iterator_A.add_tile_offset({0, 1});
-            iterator_B.add_tile_offset({1, 0});
-
-            this->smem_iterator_A_.add_tile_offset({0, 1});
-            this->smem_iterator_B_.add_tile_offset({1, 0});
-
-            // Defines the boundary of a stage of cp.async.
-            cutlass::arch::cp_async_fence();
+          ++iterator_A;
         }
 
-        // Perform accumulation in the 'd' output operand
-        accum = src_accum;
+        ++this->smem_iterator_A_;
+      }
 
-        //
-        // Clear the remaining tiles of SMEM. This is a functional requirement for some kernels
-        // so that all accumulator elements outside the GEMM footprint are zero.
-        //
+      iterator_B.set_iteration_index(0);
+      this->smem_iterator_B_.set_iteration_index(0);
 
-        if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage)
-        {
+      // Async Copy for operand B
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                this->smem_iterator_B_.get());
 
-            /// Iterator to write threadblock-scoped tile of A operand to shared memory
-            SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
+          int const kSrcBytes =
+              sizeof_bits<typename IteratorB::Element>::value *
+              IteratorB::ThreadMap::kElementsPerAccess /
+              IteratorB::kAccessesPerVector / 8;
 
-            typename IteratorA::AccessType zero_A;
-            zero_A.clear();
+          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
+              dst_ptr + v, iterator_B.get(), iterator_B.valid());
 
-            last_smem_iterator_A.set_iteration_index(0);
-
-            // Async Copy for operand A
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j)
-            {
-
-                typename IteratorA::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorA::AccessType*>(last_smem_iterator_A.get());
-
-                *dst_ptr = zero_A;
-
-                ++last_smem_iterator_A;
-            }
-
-            /// Iterator to write threadblock-scoped tile of B operand to shared memory
-            SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_);
-            typename IteratorB::AccessType zero_B;
-
-            zero_B.clear();
-            last_smem_iterator_B.set_iteration_index(0);
-
-            // Async Copy for operand B
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j)
-            {
-
-                typename IteratorB::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorB::AccessType*>(last_smem_iterator_B.get());
-
-                *dst_ptr = zero_B;
-
-                ++last_smem_iterator_B;
-            }
+          ++iterator_B;
         }
 
-        // Wait until we have at least one committed global fetch stage. (#uncommitted = Base::kStages - 1 - #committed)
-        cutlass::arch::cp_async_wait<Base::kStages - 2>();
-        __syncthreads();
+        ++this->smem_iterator_B_;
+      }
 
-        // Pair of fragments used to overlap shared memory loads and math
-        // instructions
-        WarpFragmentA warp_frag_A[2];
-        WarpFragmentB warp_frag_B[2];
-        typename Dequantizer::FragmentScale warp_frag_scales;
-        typename Dequantizer::FragmentZero warp_frag_zeros;
+      copy_scales_and_advance(iterator_scale, stage, gemm_k_iterations);
 
-        Operator warp_mma;
+      // Move to the next stage
+      iterator_A.add_tile_offset({0, 1});
+      iterator_B.add_tile_offset({1, 0});
 
-        this->warp_tile_iterator_A_.set_kgroup_index(0);
-        this->warp_tile_iterator_B_.set_kgroup_index(0);
+      this->smem_iterator_A_.add_tile_offset({0, 1});
+      this->smem_iterator_B_.add_tile_offset({1, 0});
 
-        this->warp_tile_iterator_A_.load(warp_frag_A[0]);
-        this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+      // Defines the boundary of a stage of cp.async.
+      cutlass::arch::cp_async_fence();
+    }
 
-        warp_dequantizer_.load(warp_frag_scales, warp_frag_zeros);
+    // Perform accumulation in the 'd' output operand
+    accum = src_accum;
+
+    //
+    // Clear the remaining tiles of SMEM. This is a functional requirement for
+    // some kernels so that all accumulator elements outside the GEMM footprint
+    // are zero.
+    //
+
+    if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) {
+      /// Iterator to write threadblock-scoped tile of A operand to shared
+      /// memory
+      SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
+
+      typename IteratorA::AccessType zero_A;
+      zero_A.clear();
+
+      last_smem_iterator_A.set_iteration_index(0);
+
+      // Async Copy for operand A
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                last_smem_iterator_A.get());
+
+        *dst_ptr = zero_A;
+
+        ++last_smem_iterator_A;
+      }
+
+      /// Iterator to write threadblock-scoped tile of B operand to shared
+      /// memory
+      SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_);
+      typename IteratorB::AccessType zero_B;
+
+      zero_B.clear();
+      last_smem_iterator_B.set_iteration_index(0);
+
+      // Async Copy for operand B
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                last_smem_iterator_B.get());
+
+        *dst_ptr = zero_B;
+
+        ++last_smem_iterator_B;
+      }
+    }
+
+    // Wait until we have at least one committed global fetch stage.
+    // (#uncommitted = Base::kStages - 1 - #committed)
+    cutlass::arch::cp_async_wait<Base::kStages - 2>();
+    __syncthreads();
+
+    // Pair of fragments used to overlap shared memory loads and math
+    // instructions
+    WarpFragmentA warp_frag_A[2];
+    WarpFragmentB warp_frag_B[2];
+    typename Dequantizer::FragmentScale warp_frag_scales;
+    typename Dequantizer::FragmentZero warp_frag_zeros;
+
+    Operator warp_mma;
+
+    this->warp_tile_iterator_A_.set_kgroup_index(0);
+    this->warp_tile_iterator_B_.set_kgroup_index(0);
+
+    this->warp_tile_iterator_A_.load(warp_frag_A[0]);
+    this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+
+    warp_dequantizer_.load(warp_frag_scales, warp_frag_zeros);
+
+    // if((threadIdx.x||threadIdx.y||threadIdx.z)==0){
+    //     uint32_t* frag_b_reg_ptr =
+    //     reinterpret_cast<uint32_t*>(&warp_frag_B[0]); printf("####
+    //     warp_frag_b_load [0] bid:%d-%d-%d,"
+    //                 " frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
+    //             blockIdx.x,blockIdx.y,blockIdx.z,
+    //             frag_b_reg_ptr[0],
+    //             frag_b_reg_ptr[1],
+    //             frag_b_reg_ptr[2],
+    //             frag_b_reg_ptr[3],
+    //             frag_b_reg_ptr[4],
+    //             frag_b_reg_ptr[5],
+    //             frag_b_reg_ptr[6],
+    //             frag_b_reg_ptr[7]
+    //             );
+    // }
+
+    ++this->warp_tile_iterator_A_;
+    ++this->warp_tile_iterator_B_;
+    warp_dequantizer_.add_pointer_offset(Shape::kN);
+
+    iterator_A.clear_mask(gemm_k_iterations == 0);
+    iterator_B.clear_mask(gemm_k_iterations == 0);
+    iterator_scale.clear_mask(gemm_k_iterations == 0);
+
+    int smem_write_stage_idx = Base::kStages - 1;
+    int smem_read_stage_idx = 0;
+
+    //
+    // Mainloop
+    //
+
+    CUTLASS_GEMM_LOOP
+    for (; gemm_k_iterations > (-Base::kStages + 1);) {
+      //
+      // Loop over GEMM K dimension
+      //
+
+      // Computes a warp-level GEMM on data held in shared memory
+      // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate
+      CUTLASS_PRAGMA_UNROLL
+      for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
+           ++warp_mma_k) {
+        // Load warp-level tiles from shared memory, wrapping to k offset if
+        // this is the last group as the case may be.
+
+        this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) %
+                                                     Base::kWarpGemmIterations);
+        this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
+        ++this->warp_tile_iterator_A_;
+
+        int const warp_tileB_k_compute_offset =
+            warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
+        int const warp_tileB_k_load_offset =
+            warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
+        if (warp_tileB_k_compute_offset ==
+            Base::kNumKIterationsPerWarpBLoad - 1) {
+          this->warp_tile_iterator_B_.set_kgroup_index(
+              (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
+          this->warp_tile_iterator_B_.load(
+              warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
+          ++this->warp_tile_iterator_B_;
+        }
+
+        typename TransformBAfterLDS::result_type converted_frag_B =
+            lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
+        warp_dequantizer_.dequantize(
+            converted_frag_B, warp_frag_scales, warp_frag_zeros);
+
+        using FragmentOperandB =
+            cutlass::Array<ElementA, Operator::FragmentB::kElements>;
+        constexpr cutlass::FloatRoundStyle RoundStyle =
+            cutlass::FloatRoundStyle::round_to_nearest;
+        constexpr int ConversionVectorWidth =
+            TransformBAfterLDS::result_type::kElements;
+        static_assert(ConversionVectorWidth == FragmentOperandB::kElements);
+
+        using Converter = cutlass::NumericArrayConverter<ElementA,
+                                                         ElementScale,
+                                                         ConversionVectorWidth,
+                                                         RoundStyle>;
+
+        FragmentOperandB converted_frag_B_operand =
+            Converter::convert(converted_frag_B);
 
         // if((threadIdx.x||threadIdx.y||threadIdx.z)==0){
-        //     uint32_t* frag_b_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_B[0]);
-        //     printf("#### warp_frag_b_load [0] bid:%d-%d-%d,"
-        //                 " frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
+        //     uint32_t* frag_b_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&warp_frag_B[(warp_tileB_k_load_offset)
+        //     % 2]); uint32_t* converted_frag_B_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&converted_frag_B); printf("####
+        //     after lds_converter bid:%d-%d-%d"
+        //             " frag_b_reg_ptr[%d]:%x-%x-%x-%x-%x-%x-%x-%x"
+        //             " converted_frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
         //             blockIdx.x,blockIdx.y,blockIdx.z,
+        //             ((warp_tileB_k_load_offset) % 2),
         //             frag_b_reg_ptr[0],
         //             frag_b_reg_ptr[1],
         //             frag_b_reg_ptr[2],
@@ -559,195 +684,124 @@ public:
         //             frag_b_reg_ptr[4],
         //             frag_b_reg_ptr[5],
         //             frag_b_reg_ptr[6],
-        //             frag_b_reg_ptr[7]
-        //             );
+        //             frag_b_reg_ptr[7],
+        //             converted_frag_B_reg_ptr[0],
+        //             converted_frag_B_reg_ptr[1],
+        //             converted_frag_B_reg_ptr[2],
+        //             converted_frag_B_reg_ptr[3],
+        //             converted_frag_B_reg_ptr[4],
+        //             converted_frag_B_reg_ptr[5],
+        //             converted_frag_B_reg_ptr[6],
+        //             converted_frag_B_reg_ptr[7]
+        //     );
         // }
 
-        ++this->warp_tile_iterator_A_;
-        ++this->warp_tile_iterator_B_;
-        warp_dequantizer_.add_pointer_offset(Shape::kN);
+        run_warp_mma(warp_mma,
+                     accum,
+                     warp_frag_A[warp_mma_k % 2],
+                     converted_frag_B_operand,
+                     accum,
+                     warp_tileB_k_compute_offset);
 
-        iterator_A.clear_mask(gemm_k_iterations == 0);
-        iterator_B.clear_mask(gemm_k_iterations == 0);
-        iterator_scale.clear_mask(gemm_k_iterations == 0);
+        // Issue global->shared copies for the this stage
+        if (warp_mma_k < Base::kWarpGemmIterations - 1) {
+          int group_start_iteration_A, group_start_iteration_B;
 
-        int smem_write_stage_idx = Base::kStages - 1;
-        int smem_read_stage_idx = 0;
+          group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
+          group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
 
-        //
-        // Mainloop
-        //
+          copy_tiles_and_advance(iterator_A,
+                                 iterator_B,
+                                 group_start_iteration_A,
+                                 group_start_iteration_B);
 
-        CUTLASS_GEMM_LOOP
-        for (; gemm_k_iterations > (-Base::kStages + 1);)
-        {
-            //
-            // Loop over GEMM K dimension
-            //
-
-            // Computes a warp-level GEMM on data held in shared memory
-            // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate
-            CUTLASS_PRAGMA_UNROLL
-            for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k)
-            {
-
-                // Load warp-level tiles from shared memory, wrapping to k offset if
-                // this is the last group as the case may be.
-
-                this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
-                this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
-                ++this->warp_tile_iterator_A_;
-
-                int const warp_tileB_k_compute_offset = warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
-                int const warp_tileB_k_load_offset = warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
-                if (warp_tileB_k_compute_offset == Base::kNumKIterationsPerWarpBLoad - 1)
-                {
-                    this->warp_tile_iterator_B_.set_kgroup_index(
-                        (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
-                    this->warp_tile_iterator_B_.load(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
-                    ++this->warp_tile_iterator_B_;
-                }
-
-                typename TransformBAfterLDS::result_type converted_frag_B
-                    = lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
-                warp_dequantizer_.dequantize(converted_frag_B, warp_frag_scales, warp_frag_zeros);
-
-                using FragmentOperandB = cutlass::Array<ElementA, Operator::FragmentB::kElements>;
-                constexpr cutlass::FloatRoundStyle RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
-                constexpr int ConversionVectorWidth = TransformBAfterLDS::result_type::kElements;
-                static_assert(ConversionVectorWidth == FragmentOperandB::kElements);
-
-                using Converter
-                    = cutlass::NumericArrayConverter<ElementA, ElementScale, ConversionVectorWidth, RoundStyle>;
-
-                FragmentOperandB converted_frag_B_operand = Converter::convert(converted_frag_B);
-
-
-                // if((threadIdx.x||threadIdx.y||threadIdx.z)==0){
-                //     uint32_t* frag_b_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_B[(warp_tileB_k_load_offset) % 2]);
-                //     uint32_t* converted_frag_B_reg_ptr = reinterpret_cast<uint32_t*>(&converted_frag_B);
-                //     printf("#### after lds_converter bid:%d-%d-%d"
-                //             " frag_b_reg_ptr[%d]:%x-%x-%x-%x-%x-%x-%x-%x"
-                //             " converted_frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
-                //             blockIdx.x,blockIdx.y,blockIdx.z,
-                //             ((warp_tileB_k_load_offset) % 2),
-                //             frag_b_reg_ptr[0],
-                //             frag_b_reg_ptr[1],
-                //             frag_b_reg_ptr[2],
-                //             frag_b_reg_ptr[3],
-                //             frag_b_reg_ptr[4],
-                //             frag_b_reg_ptr[5],
-                //             frag_b_reg_ptr[6],
-                //             frag_b_reg_ptr[7],
-                //             converted_frag_B_reg_ptr[0],
-                //             converted_frag_B_reg_ptr[1],
-                //             converted_frag_B_reg_ptr[2],
-                //             converted_frag_B_reg_ptr[3],
-                //             converted_frag_B_reg_ptr[4],
-                //             converted_frag_B_reg_ptr[5],
-                //             converted_frag_B_reg_ptr[6],
-                //             converted_frag_B_reg_ptr[7]
-                //     );
-                // }
-
-                run_warp_mma(warp_mma, accum, warp_frag_A[warp_mma_k % 2], converted_frag_B_operand, accum,
-                    warp_tileB_k_compute_offset);
-
-                // Issue global->shared copies for the this stage
-                if (warp_mma_k < Base::kWarpGemmIterations - 1)
-                {
-                    int group_start_iteration_A, group_start_iteration_B;
-
-                    group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
-                    group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
-
-                    copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
-
-                    // This is the first group of a given stage, so we issue the loads for the B scales immediately.
-                    if (group_start_iteration_B == 0)
-                    {
-                        copy_scales_and_advance(iterator_scale);
-                    }
-                }
-
-                if (warp_mma_k + 2 == Base::kWarpGemmIterations)
-                {
-                    int group_start_iteration_A, group_start_iteration_B;
-                    group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
-                    group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
-
-                    copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
-
-                    // Inserts a memory fence between stages of cp.async instructions.
-                    cutlass::arch::cp_async_fence();
-
-                    // Wait until we have at least one committed global fetch stage. (#uncommitted = Base::kStages - 1 -
-                    // #committed)
-                    arch::cp_async_wait<Base::kStages - 2>();
-                    __syncthreads();
-
-                    // Move to the next stage
-                    iterator_A.add_tile_offset({0, 1});
-                    iterator_B.add_tile_offset({1, 0});
-
-                    this->smem_iterator_A_.add_tile_offset({0, 1});
-                    this->smem_iterator_B_.add_tile_offset({1, 0});
-
-                    // Add negative offsets to return iterators to the 'start' of the
-                    // circular buffer in shared memory
-                    if (smem_write_stage_idx == (Base::kStages - 1))
-                    {
-                        this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
-                        this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
-                        this->smem_iterator_scale_.add_tile_offset({-Base::kStages, 0});
-                        smem_write_stage_idx = 0;
-                    }
-                    else
-                    {
-                        ++smem_write_stage_idx;
-                    }
-
-                    if (smem_read_stage_idx == (Base::kStages - 1))
-                    {
-                        this->warp_tile_iterator_A_.add_tile_offset(
-                            {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
-                        this->warp_tile_iterator_B_.add_tile_offset(
-                            {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0});
-                        warp_dequantizer_.add_pointer_offset(-Base::kStages * Shape::kN);
-                        smem_read_stage_idx = 0;
-                    }
-                    else
-                    {
-                        ++smem_read_stage_idx;
-                    }
-
-                    --gemm_k_iterations;
-                    iterator_A.clear_mask(gemm_k_iterations == 0);
-                    iterator_B.clear_mask(gemm_k_iterations == 0);
-                    iterator_scale.clear_mask(gemm_k_iterations == 0);
-                }
-            }
-
-            // Load the scale needed for the next tile iteration.
-            warp_dequantizer_.load(warp_frag_scales, warp_frag_zeros);
-            // Update internal pointer to set of scales in shared memory.
-            warp_dequantizer_.add_pointer_offset(Shape::kN);
+          // This is the first group of a given stage, so we issue the loads for
+          // the B scales immediately.
+          if (group_start_iteration_B == 0) {
+            copy_scales_and_advance(iterator_scale);
+          }
         }
 
-        if (SharedMemoryClear == SharedMemoryClearOption::kZfill)
-        {
-            // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop
-            cutlass::arch::cp_async_fence();
-            cutlass::arch::cp_async_wait<0>();
-            __syncthreads();
+        if (warp_mma_k + 2 == Base::kWarpGemmIterations) {
+          int group_start_iteration_A, group_start_iteration_B;
+          group_start_iteration_A =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
+          group_start_iteration_B =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
+
+          copy_tiles_and_advance(iterator_A,
+                                 iterator_B,
+                                 group_start_iteration_A,
+                                 group_start_iteration_B);
+
+          // Inserts a memory fence between stages of cp.async instructions.
+          cutlass::arch::cp_async_fence();
+
+          // Wait until we have at least one committed global fetch stage.
+          // (#uncommitted = Base::kStages - 1 - #committed)
+          arch::cp_async_wait<Base::kStages - 2>();
+          __syncthreads();
+
+          // Move to the next stage
+          iterator_A.add_tile_offset({0, 1});
+          iterator_B.add_tile_offset({1, 0});
+
+          this->smem_iterator_A_.add_tile_offset({0, 1});
+          this->smem_iterator_B_.add_tile_offset({1, 0});
+
+          // Add negative offsets to return iterators to the 'start' of the
+          // circular buffer in shared memory
+          if (smem_write_stage_idx == (Base::kStages - 1)) {
+            this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
+            this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
+            this->smem_iterator_scale_.add_tile_offset({-Base::kStages, 0});
+            smem_write_stage_idx = 0;
+          } else {
+            ++smem_write_stage_idx;
+          }
+
+          if (smem_read_stage_idx == (Base::kStages - 1)) {
+            this->warp_tile_iterator_A_.add_tile_offset(
+                {0,
+                 -Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterations});
+            this->warp_tile_iterator_B_.add_tile_offset(
+                {-Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterationsForB,
+                 0});
+            warp_dequantizer_.add_pointer_offset(-Base::kStages * Shape::kN);
+            smem_read_stage_idx = 0;
+          } else {
+            ++smem_read_stage_idx;
+          }
+
+          --gemm_k_iterations;
+          iterator_A.clear_mask(gemm_k_iterations == 0);
+          iterator_B.clear_mask(gemm_k_iterations == 0);
+          iterator_scale.clear_mask(gemm_k_iterations == 0);
         }
+      }
+
+      // Load the scale needed for the next tile iteration.
+      warp_dequantizer_.load(warp_frag_scales, warp_frag_zeros);
+      // Update internal pointer to set of scales in shared memory.
+      warp_dequantizer_.add_pointer_offset(Shape::kN);
     }
+
+    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+      // commit and drain all pending and predicated LDGSTS pnz from the GEMM
+      // mainloop
+      cutlass::arch::cp_async_fence();
+      cutlass::arch::cp_async_wait<0>();
+      __syncthreads();
+    }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage_percol.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage_percol.h
index 83efdc5cb0..6371da633c 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage_percol.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_multistage_percol.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -48,12 +49,9 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -98,550 +96,605 @@ template <
     WeightOnlyQuantOp QuantOp_,
     /// Use zfill or predicate for out-of-bound cp.async
     SharedMemoryClearOption SharedMemoryClear>
-class DqMmaMultistage<Shape_, IteratorA_, SmemIteratorA_, CacheOpA, IteratorB_, SmemIteratorB_, CacheOpB,
-    IteratorScale_, SmemIteratorScale_, ElementC_, LayoutC_, Policy_, Stages, TransformBAfterLDS_, QuantOp_,
-    SharedMemoryClear, std::enable_if_t<!isFinegrained(QuantOp_)>>
-    : public DqMmaBase<Shape_, Policy_, typename IteratorScale_::Element, Stages, QuantOp_>
-{
-public:
-    ///< Base class
-    using Base = DqMmaBase<Shape_, Policy_, typename IteratorScale_::Element, Stages, QuantOp_>;
-    ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using Shape = Shape_;
-    ///< Iterates over tiles of A operand in global memory
-    using IteratorA = IteratorA_;
-    ///< Iterates over tiles of B operand in global memory
-    using IteratorB = IteratorB_;
-    ///< Data type of accumulator matrix
-    using ElementC = ElementC_;
-    ///< Layout of accumulator matrix
-    using LayoutC = LayoutC_;
-    ///< Policy describing tuning details
-    using Policy = Policy_;
+class DqMmaMultistage<Shape_,
+                      IteratorA_,
+                      SmemIteratorA_,
+                      CacheOpA,
+                      IteratorB_,
+                      SmemIteratorB_,
+                      CacheOpB,
+                      IteratorScale_,
+                      SmemIteratorScale_,
+                      ElementC_,
+                      LayoutC_,
+                      Policy_,
+                      Stages,
+                      TransformBAfterLDS_,
+                      QuantOp_,
+                      SharedMemoryClear,
+                      std::enable_if_t<!isFinegrained(QuantOp_)>>
+    : public DqMmaBase<Shape_,
+                       Policy_,
+                       typename IteratorScale_::Element,
+                       Stages,
+                       QuantOp_> {
+ public:
+  ///< Base class
+  using Base = DqMmaBase<Shape_,
+                         Policy_,
+                         typename IteratorScale_::Element,
+                         Stages,
+                         QuantOp_>;
+  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using Shape = Shape_;
+  ///< Iterates over tiles of A operand in global memory
+  using IteratorA = IteratorA_;
+  ///< Iterates over tiles of B operand in global memory
+  using IteratorB = IteratorB_;
+  ///< Data type of accumulator matrix
+  using ElementC = ElementC_;
+  ///< Layout of accumulator matrix
+  using LayoutC = LayoutC_;
+  ///< Policy describing tuning details
+  using Policy = Policy_;
 
-    using IteratorScale = IteratorScale_;
-    using ElementScale = typename IteratorScale::Element;
-    using LayoutScale = typename IteratorScale::Layout;
+  using IteratorScale = IteratorScale_;
+  using ElementScale = typename IteratorScale::Element;
+  using LayoutScale = typename IteratorScale::Layout;
 
-    using SmemIteratorA = SmemIteratorA_;
-    using SmemIteratorB = SmemIteratorB_;
-    using SmemIteratorScale = SmemIteratorScale_;
+  using SmemIteratorA = SmemIteratorA_;
+  using SmemIteratorB = SmemIteratorB_;
+  using SmemIteratorScale = SmemIteratorScale_;
 
-    static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
-    static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
+  static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
+  static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
 
-    using TransformBAfterLDS = TransformBAfterLDS_;
+  using TransformBAfterLDS = TransformBAfterLDS_;
 
-    static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
+  static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
 
-    //
-    // Dependent types
-    //
+  //
+  // Dependent types
+  //
 
-    /// Fragment of operand Scale loaded from global memory;
-    using FragmentScale = typename IteratorScale::Fragment;
+  /// Fragment of operand Scale loaded from global memory;
+  using FragmentScale = typename IteratorScale::Fragment;
 
-    /// Fragment of accumulator tile
-    using FragmentC = typename Policy::Operator::FragmentC;
+  /// Fragment of accumulator tile
+  using FragmentC = typename Policy::Operator::FragmentC;
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
 
-    /// Minimum architecture is Sm80 to support cp.async
-    using ArchTag = arch::Sm80;
+  /// Minimum architecture is Sm80 to support cp.async
+  using ArchTag = arch::Sm80;
 
-    using Dequantizer = warp::MmaTensorOpDequantizer<Operator, typename Base::WarpGemm, Operand::kB, ElementScale,
-        LayoutScale, 32, QuantOp>;
+  using Dequantizer = warp::MmaTensorOpDequantizer<Operator,
+                                                   typename Base::WarpGemm,
+                                                   Operand::kB,
+                                                   ElementScale,
+                                                   LayoutScale,
+                                                   32,
+                                                   QuantOp>;
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = Operator::kTransformA;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = Operator::kTransformA;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = Operator::kTransformB;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = Operator::kTransformB;
 
-    /// Internal structure exposed for introspection.
-    struct Detail
-    {
+  /// Internal structure exposed for introspection.
+  struct Detail {
+    static_assert(Base::kWarpGemmIterations > 1,
+                  "The pipelined structure requires at least two warp-level "
+                  "GEMM operations.");
 
-        static_assert(Base::kWarpGemmIterations > 1,
-            "The pipelined structure requires at least two warp-level "
-            "GEMM operations.");
+    /// Number of cp.async instructions to load one stage of operand A
+    static int const AsyncCopyIterationsPerStageA =
+        IteratorA::ThreadMap::Iterations::kCount;
 
-        /// Number of cp.async instructions to load one stage of operand A
-        static int const AsyncCopyIterationsPerStageA = IteratorA::ThreadMap::Iterations::kCount;
+    /// Number of cp.async instructions to load one stage of operand B
+    static int const AsyncCopyIterationsPerStageB =
+        IteratorB::ThreadMap::Iterations::kCount;
 
-        /// Number of cp.async instructions to load one stage of operand B
-        static int const AsyncCopyIterationsPerStageB = IteratorB::ThreadMap::Iterations::kCount;
+    /// Number of stages
+    static int const kStages = Stages;
 
-        /// Number of stages
-        static int const kStages = Stages;
+    /// Number of cp.async instructions to load on group of operand A
+    static int const kAccessesPerGroupA =
+        (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
 
-        /// Number of cp.async instructions to load on group of operand A
-        static int const kAccessesPerGroupA
-            = (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
+    /// Number of cp.async instructions to load on group of operand B
+    static int const kAccessesPerGroupB =
+        (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
+  };
 
-        /// Number of cp.async instructions to load on group of operand B
-        static int const kAccessesPerGroupB
-            = (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
-    };
+ private:
+  using WarpFragmentA = typename Operator::FragmentA;
+  using WarpFragmentB = typename Operator::FragmentB;
+  Dequantizer warp_dequantizer_;
 
-private:
-    using WarpFragmentA = typename Operator::FragmentA;
-    using WarpFragmentB = typename Operator::FragmentB;
-    Dequantizer warp_dequantizer_;
+  using ElementA = typename IteratorA::Element;
+  using ElementB = typename IteratorB::Element;
+  using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
 
-    using ElementA = typename IteratorA::Element;
-    using ElementB = typename IteratorB::Element;
-    using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
+  static constexpr bool RequiresTileInterleave =
+      layout::IsColumnMajorTileInterleave<
+          typename LayoutDetailsForB::Layout>::value;
+  static_assert(!RequiresTileInterleave ||
+                    (RequiresTileInterleave &&
+                     (Shape::kK == LayoutDetailsForB::ThreadblockK)),
+                "Layout K must match threadblockK");
 
-    static constexpr bool RequiresTileInterleave
-        = layout::IsColumnMajorTileInterleave<typename LayoutDetailsForB::Layout>::value;
-    static_assert(!RequiresTileInterleave || (RequiresTileInterleave && (Shape::kK == LayoutDetailsForB::ThreadblockK)),
-        "Layout K must match threadblockK");
+ private:
+  //
+  // Data members
+  //
 
-private:
-    //
-    // Data members
-    //
+  /// Iterator to write threadblock-scoped tile of A operand to shared memory
+  SmemIteratorA smem_iterator_A_;
 
-    /// Iterator to write threadblock-scoped tile of A operand to shared memory
-    SmemIteratorA smem_iterator_A_;
+  /// Iterator to write threadblock-scoped tile of B operand to shared memory
+  SmemIteratorB smem_iterator_B_;
 
-    /// Iterator to write threadblock-scoped tile of B operand to shared memory
-    SmemIteratorB smem_iterator_B_;
+  /// Iterator to write threadblock-scoped tile of scale operand to shared
+  /// memory
+  SmemIteratorScale smem_iterator_scale_;
 
-    /// Iterator to write threadblock-scoped tile of scale operand to shared memory
-    SmemIteratorScale smem_iterator_scale_;
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  DqMmaMultistage(
+      ///< Shared storage needed for internal use by threadblock-scoped GEMM
+      typename Base::SharedStorage& shared_storage,
+      ///< Group size for quantization. Not used by this main loop since it
+      ///< assumes per-column
+      int const group_size,
+      ///< ID within the threadblock
+      int thread_idx,
+      ///< ID of warp
+      int warp_idx,
+      ///< ID of each thread within a warp
+      int lane_idx)
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
+        warp_dequantizer_(
+            {shared_storage.operand_scale.data(), LayoutScale(Shape::kN)},
+            (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) /
+                Base::WarpCount::kM,
+            lane_idx),
+        smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
+        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx),
+        smem_iterator_scale_(LayoutScale(Shape::kN),
+                             shared_storage.operand_scale.data(),
+                             {1, Shape::kN},
+                             thread_idx) {
+    // Compute warp location within threadblock tile by mapping the warp_id to
+    // three coordinates:
+    //   _m: the warp's position within the threadblock along the M dimension
+    //   _n: the warp's position within the threadblock along the N dimension
+    //   _k: the warp's position within the threadblock along the K dimension
 
-public:
-    /// Construct from tensor references
-    CUTLASS_DEVICE
-    DqMmaMultistage(
-        ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        typename Base::SharedStorage& shared_storage,
-        ///< Group size for quantization. Not used by this main loop since it assumes per-column
-        int const group_size,
-        ///< ID within the threadblock
-        int thread_idx,
-        ///< ID of warp
-        int warp_idx,
-        ///< ID of each thread within a warp
-        int lane_idx)
-        : Base(shared_storage, thread_idx, warp_idx, lane_idx)
-        , warp_dequantizer_({shared_storage.operand_scale.data(), LayoutScale(Shape::kN)},
-              (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) / Base::WarpCount::kM, lane_idx)
-        , smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx)
-        , smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx)
-        , smem_iterator_scale_(LayoutScale(Shape::kN), shared_storage.operand_scale.data(), {1, Shape::kN}, thread_idx)
-    {
-        // Compute warp location within threadblock tile by mapping the warp_id to
-        // three coordinates:
-        //   _m: the warp's position within the threadblock along the M dimension
-        //   _n: the warp's position within the threadblock along the N dimension
-        //   _k: the warp's position within the threadblock along the K dimension
+    int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
 
-        int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
-        int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
+    int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
 
-        int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
-        int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
+    // Add per-warp offsets in units of warp-level tiles
+    this->warp_tile_iterator_A_.add_tile_offset(
+        {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
+    this->warp_tile_iterator_B_.add_tile_offset(
+        {Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  }
 
-        // Add per-warp offsets in units of warp-level tiles
-        this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
-        this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  CUTLASS_DEVICE
+  void copy_tiles_and_advance(IteratorA& iterator_A,
+                              IteratorB& iterator_B,
+                              int group_start_A = 0,
+                              int group_start_B = 0) {
+    iterator_A.set_iteration_index(group_start_A *
+                                   IteratorA::kAccessesPerVector);
+    this->smem_iterator_A_.set_iteration_index(group_start_A);
+
+    // Async Copy for operand A
+    CUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) {
+      if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                this->smem_iterator_A_.get());
+
+        int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
+                              IteratorA::ThreadMap::kElementsPerAccess /
+                              IteratorA::kAccessesPerVector / 8;
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
+          auto gmem_ptr = iterator_A.get();
+
+          if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
+                dst_ptr + v, gmem_ptr, iterator_A.valid());
+          } else {
+            cutlass::arch::cp_async<kSrcBytes, kCacheOpA>(
+                dst_ptr + v, gmem_ptr, iterator_A.valid());
+          }
+
+          ++iterator_A;
+        }
+
+        ++this->smem_iterator_A_;
+      }
     }
 
-    CUTLASS_DEVICE
-    void copy_tiles_and_advance(
-        IteratorA& iterator_A, IteratorB& iterator_B, int group_start_A = 0, int group_start_B = 0)
-    {
-        iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector);
-        this->smem_iterator_A_.set_iteration_index(group_start_A);
+    iterator_B.set_iteration_index(group_start_B *
+                                   IteratorB::kAccessesPerVector);
+    this->smem_iterator_B_.set_iteration_index(group_start_B);
+
+    // Async Copy for operand B
+    CUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) {
+      if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                this->smem_iterator_B_.get());
+
+        int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value *
+                              IteratorB::ThreadMap::kElementsPerAccess /
+                              IteratorB::kAccessesPerVector / 8;
 
-        // Async Copy for operand A
         CUTLASS_PRAGMA_UNROLL
-        for (int j = 0; j < Detail::kAccessesPerGroupA; ++j)
-        {
-            if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA)
-            {
-                typename IteratorA::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorA::AccessType*>(this->smem_iterator_A_.get());
+        for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
+          auto gmem_ptr = iterator_B.get();
 
-                int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value
-                    * IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector / 8;
+          if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
+                dst_ptr + v, gmem_ptr, iterator_B.valid());
+          } else {
+            cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(
+                dst_ptr + v, gmem_ptr, iterator_B.valid());
+          }
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorA::kAccessesPerVector; ++v)
-                {
-                    auto gmem_ptr = iterator_A.get();
+          ++iterator_B;
+        }
+        ++this->smem_iterator_B_;
+      }
+    }
+  }
 
-                    if (SharedMemoryClear == SharedMemoryClearOption::kZfill)
-                    {
-                        cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(dst_ptr + v, gmem_ptr, iterator_A.valid());
-                    }
-                    else
-                    {
-                        cutlass::arch::cp_async<kSrcBytes, kCacheOpA>(dst_ptr + v, gmem_ptr, iterator_A.valid());
-                    }
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  CUTLASS_DEVICE
+  void operator()(
+      ///< problem size of GEMM
+      int gemm_k_iterations,
+      ///< destination accumulator tile
+      FragmentC& accum,
+      ///< iterator over A operand in global memory
+      IteratorA iterator_A,
+      ///< iterator over B operand in global memory
+      IteratorB iterator_B,
+      ///< iterator over scale operand in global memory
+      IteratorScale iterator_scale,
+      ///< initial value of accumulator
+      FragmentC const& src_accum) {
+    //
+    // Prologue
+    //
 
-                    ++iterator_A;
-                }
+    TransformBAfterLDS lds_converter;
 
-                ++this->smem_iterator_A_;
-            }
+    // NOTE - switch to ldg.sts
+    // Issue this first, so cp.async.commit_group will commit this load as well.
+    // Note: we do not commit here and this load will commit in the same group
+    // as
+    //       the first load of A.
+    FragmentScale tb_frag_scales;
+    tb_frag_scales.clear();
+    iterator_scale.load(tb_frag_scales);
+    this->smem_iterator_scale_.store(tb_frag_scales);
+
+    // Issue several complete stages
+    CUTLASS_PRAGMA_UNROLL
+    for (int stage = 0; stage < Base::kStages - 1;
+         ++stage, --gemm_k_iterations) {
+      iterator_A.clear_mask(gemm_k_iterations == 0);
+      iterator_B.clear_mask(gemm_k_iterations == 0);
+
+      iterator_A.set_iteration_index(0);
+      this->smem_iterator_A_.set_iteration_index(0);
+
+      // Async Copy for operand A
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                this->smem_iterator_A_.get());
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
+          int const kSrcBytes =
+              sizeof_bits<typename IteratorA::Element>::value *
+              IteratorA::ThreadMap::kElementsPerAccess /
+              IteratorA::kAccessesPerVector / 8;
+
+          int src_bytes = (iterator_A.valid() ? kSrcBytes : 0);
+
+          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
+              dst_ptr + v, iterator_A.get(), iterator_A.valid());
+
+          ++iterator_A;
         }
 
-        iterator_B.set_iteration_index(group_start_B * IteratorB::kAccessesPerVector);
-        this->smem_iterator_B_.set_iteration_index(group_start_B);
+        ++this->smem_iterator_A_;
+      }
+
+      iterator_B.set_iteration_index(0);
+      this->smem_iterator_B_.set_iteration_index(0);
+
+      // Async Copy for operand B
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                this->smem_iterator_B_.get());
 
-        // Async Copy for operand B
         CUTLASS_PRAGMA_UNROLL
-        for (int j = 0; j < Detail::kAccessesPerGroupB; ++j)
-        {
-            if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB)
-            {
-                typename IteratorB::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorB::AccessType*>(this->smem_iterator_B_.get());
+        for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
+          int const kSrcBytes =
+              sizeof_bits<typename IteratorB::Element>::value *
+              IteratorB::ThreadMap::kElementsPerAccess /
+              IteratorB::kAccessesPerVector / 8;
 
-                int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value
-                    * IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector / 8;
+          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
+              dst_ptr + v, iterator_B.get(), iterator_B.valid());
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorB::kAccessesPerVector; ++v)
-                {
-                    auto gmem_ptr = iterator_B.get();
-
-                    if (SharedMemoryClear == SharedMemoryClearOption::kZfill)
-                    {
-                        cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(dst_ptr + v, gmem_ptr, iterator_B.valid());
-                    }
-                    else
-                    {
-                        cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(dst_ptr + v, gmem_ptr, iterator_B.valid());
-                    }
-
-                    ++iterator_B;
-                }
-                ++this->smem_iterator_B_;
-            }
+          ++iterator_B;
         }
+
+        ++this->smem_iterator_B_;
+      }
+
+      // Move to the next stage
+      iterator_A.add_tile_offset({0, 1});
+      iterator_B.add_tile_offset({1, 0});
+
+      this->smem_iterator_A_.add_tile_offset({0, 1});
+      this->smem_iterator_B_.add_tile_offset({1, 0});
+
+      // Defines the boundary of a stage of cp.async.
+      cutlass::arch::cp_async_fence();
     }
 
-    /// Perform a threadblock-scoped matrix multiply-accumulate
-    CUTLASS_DEVICE
-    void operator()(
-        ///< problem size of GEMM
-        int gemm_k_iterations,
-        ///< destination accumulator tile
-        FragmentC& accum,
-        ///< iterator over A operand in global memory
-        IteratorA iterator_A,
-        ///< iterator over B operand in global memory
-        IteratorB iterator_B,
-        ///< iterator over scale operand in global memory
-        IteratorScale iterator_scale,
-        ///< initial value of accumulator
-        FragmentC const& src_accum)
-    {
+    // Perform accumulation in the 'd' output operand
+    accum = src_accum;
 
-        //
-        // Prologue
-        //
+    //
+    // Clear the remaining tiles of SMEM. This is a functional requirement for
+    // some kernels so that all accumulator elements outside the GEMM footprint
+    // are zero.
+    //
 
-        TransformBAfterLDS lds_converter;
+    if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) {
+      /// Iterator to write threadblock-scoped tile of A operand to shared
+      /// memory
+      SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
 
-        // NOTE - switch to ldg.sts
-        // Issue this first, so cp.async.commit_group will commit this load as well.
-        // Note: we do not commit here and this load will commit in the same group as
-        //       the first load of A.
-        FragmentScale tb_frag_scales;
-        tb_frag_scales.clear();
-        iterator_scale.load(tb_frag_scales);
-        this->smem_iterator_scale_.store(tb_frag_scales);
+      typename IteratorA::AccessType zero_A;
+      zero_A.clear();
 
-        // Issue several complete stages
-        CUTLASS_PRAGMA_UNROLL
-        for (int stage = 0; stage < Base::kStages - 1; ++stage, --gemm_k_iterations)
-        {
+      last_smem_iterator_A.set_iteration_index(0);
 
-            iterator_A.clear_mask(gemm_k_iterations == 0);
-            iterator_B.clear_mask(gemm_k_iterations == 0);
+      // Async Copy for operand A
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                last_smem_iterator_A.get());
 
-            iterator_A.set_iteration_index(0);
-            this->smem_iterator_A_.set_iteration_index(0);
+        *dst_ptr = zero_A;
 
-            // Async Copy for operand A
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j)
-            {
-                typename IteratorA::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorA::AccessType*>(this->smem_iterator_A_.get());
+        ++last_smem_iterator_A;
+      }
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorA::kAccessesPerVector; ++v)
-                {
-                    int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value
-                        * IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector / 8;
+      /// Iterator to write threadblock-scoped tile of B operand to shared
+      /// memory
+      SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_);
+      typename IteratorB::AccessType zero_B;
 
-                    int src_bytes = (iterator_A.valid() ? kSrcBytes : 0);
+      zero_B.clear();
+      last_smem_iterator_B.set_iteration_index(0);
 
-                    cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
-                        dst_ptr + v, iterator_A.get(), iterator_A.valid());
+      // Async Copy for operand B
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                last_smem_iterator_B.get());
 
-                    ++iterator_A;
-                }
+        *dst_ptr = zero_B;
 
-                ++this->smem_iterator_A_;
-            }
+        ++last_smem_iterator_B;
+      }
+    }
 
-            iterator_B.set_iteration_index(0);
-            this->smem_iterator_B_.set_iteration_index(0);
+    // Wait until we have at least one committed global fetch stage.
+    // (#uncommitted = Base::kStages - 1 - #committed)
+    cutlass::arch::cp_async_wait<Base::kStages - 2>();
+    __syncthreads();
 
-            // Async Copy for operand B
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j)
-            {
-                typename IteratorB::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorB::AccessType*>(this->smem_iterator_B_.get());
+    // Pair of fragments used to overlap shared memory loads and math
+    // instructions
+    WarpFragmentA warp_frag_A[2];
+    WarpFragmentB warp_frag_B[2];
+    typename Dequantizer::FragmentScale warp_frag_scales;
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorB::kAccessesPerVector; ++v)
-                {
-                    int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value
-                        * IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector / 8;
+    Operator warp_mma;
 
-                    cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
-                        dst_ptr + v, iterator_B.get(), iterator_B.valid());
+    this->warp_tile_iterator_A_.set_kgroup_index(0);
+    this->warp_tile_iterator_B_.set_kgroup_index(0);
 
-                    ++iterator_B;
-                }
+    this->warp_tile_iterator_A_.load(warp_frag_A[0]);
+    this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+    warp_dequantizer_.load(warp_frag_scales);
 
-                ++this->smem_iterator_B_;
-            }
+    ++this->warp_tile_iterator_A_;
+    ++this->warp_tile_iterator_B_;
 
-            // Move to the next stage
-            iterator_A.add_tile_offset({0, 1});
-            iterator_B.add_tile_offset({1, 0});
+    iterator_A.clear_mask(gemm_k_iterations == 0);
+    iterator_B.clear_mask(gemm_k_iterations == 0);
 
-            this->smem_iterator_A_.add_tile_offset({0, 1});
-            this->smem_iterator_B_.add_tile_offset({1, 0});
+    int smem_write_stage_idx = Base::kStages - 1;
+    int smem_read_stage_idx = 0;
 
-            // Defines the boundary of a stage of cp.async.
-            cutlass::arch::cp_async_fence();
-        }
+    //
+    // Mainloop
+    //
 
-        // Perform accumulation in the 'd' output operand
-        accum = src_accum;
+    CUTLASS_GEMM_LOOP
+    for (; gemm_k_iterations > (-Base::kStages + 1);) {
+      //
+      // Loop over GEMM K dimension
+      //
 
-        //
-        // Clear the remaining tiles of SMEM. This is a functional requirement for some kernels
-        // so that all accumulator elements outside the GEMM footprint are zero.
-        //
-
-        if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage)
-        {
-
-            /// Iterator to write threadblock-scoped tile of A operand to shared memory
-            SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
-
-            typename IteratorA::AccessType zero_A;
-            zero_A.clear();
-
-            last_smem_iterator_A.set_iteration_index(0);
-
-            // Async Copy for operand A
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j)
-            {
-
-                typename IteratorA::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorA::AccessType*>(last_smem_iterator_A.get());
-
-                *dst_ptr = zero_A;
-
-                ++last_smem_iterator_A;
-            }
-
-            /// Iterator to write threadblock-scoped tile of B operand to shared memory
-            SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_);
-            typename IteratorB::AccessType zero_B;
-
-            zero_B.clear();
-            last_smem_iterator_B.set_iteration_index(0);
-
-            // Async Copy for operand B
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j)
-            {
-
-                typename IteratorB::AccessType* dst_ptr
-                    = reinterpret_cast<typename IteratorB::AccessType*>(last_smem_iterator_B.get());
-
-                *dst_ptr = zero_B;
-
-                ++last_smem_iterator_B;
-            }
-        }
-
-        // Wait until we have at least one committed global fetch stage. (#uncommitted = Base::kStages - 1 - #committed)
-        cutlass::arch::cp_async_wait<Base::kStages - 2>();
-        __syncthreads();
-
-        // Pair of fragments used to overlap shared memory loads and math
-        // instructions
-        WarpFragmentA warp_frag_A[2];
-        WarpFragmentB warp_frag_B[2];
-        typename Dequantizer::FragmentScale warp_frag_scales;
-
-        Operator warp_mma;
-
-        this->warp_tile_iterator_A_.set_kgroup_index(0);
-        this->warp_tile_iterator_B_.set_kgroup_index(0);
-
-        this->warp_tile_iterator_A_.load(warp_frag_A[0]);
-        this->warp_tile_iterator_B_.load(warp_frag_B[0]);
-        warp_dequantizer_.load(warp_frag_scales);
+      // Computes a warp-level GEMM on data held in shared memory
+      // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate
+      CUTLASS_PRAGMA_UNROLL
+      for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
+           ++warp_mma_k) {
+        // Load warp-level tiles from shared memory, wrapping to k offset if
+        // this is the last group as the case may be.
 
+        this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) %
+                                                     Base::kWarpGemmIterations);
+        this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
         ++this->warp_tile_iterator_A_;
-        ++this->warp_tile_iterator_B_;
 
-        iterator_A.clear_mask(gemm_k_iterations == 0);
-        iterator_B.clear_mask(gemm_k_iterations == 0);
-
-        int smem_write_stage_idx = Base::kStages - 1;
-        int smem_read_stage_idx = 0;
-
-        //
-        // Mainloop
-        //
-
-        CUTLASS_GEMM_LOOP
-        for (; gemm_k_iterations > (-Base::kStages + 1);)
-        {
-            //
-            // Loop over GEMM K dimension
-            //
-
-            // Computes a warp-level GEMM on data held in shared memory
-            // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate
-            CUTLASS_PRAGMA_UNROLL
-            for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k)
-            {
-
-                // Load warp-level tiles from shared memory, wrapping to k offset if
-                // this is the last group as the case may be.
-
-                this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
-                this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
-                ++this->warp_tile_iterator_A_;
-
-                int const warp_tileB_k_compute_offset = warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
-                int const warp_tileB_k_load_offset = warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
-                if (warp_tileB_k_compute_offset == Base::kNumKIterationsPerWarpBLoad - 1)
-                {
-                    this->warp_tile_iterator_B_.set_kgroup_index(
-                        (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
-                    this->warp_tile_iterator_B_.load(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
-                    ++this->warp_tile_iterator_B_;
-                }
-
-                typename TransformBAfterLDS::result_type converted_frag_B
-                    = lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
-                warp_dequantizer_.dequantize(converted_frag_B, warp_frag_scales);
-
-                using FragmentOperandB = cutlass::Array<ElementA, Operator::FragmentB::kElements>;
-                constexpr cutlass::FloatRoundStyle RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
-                constexpr int ConversionVectorWidth = TransformBAfterLDS::result_type::kElements;
-                static_assert(ConversionVectorWidth == FragmentOperandB::kElements);
-
-                using Converter
-                    = cutlass::NumericArrayConverter<ElementA, ElementScale, ConversionVectorWidth, RoundStyle>;
-
-                FragmentOperandB converted_frag_B_operand = Converter::convert(converted_frag_B);
-                run_warp_mma(warp_mma, accum, warp_frag_A[warp_mma_k % 2], converted_frag_B_operand, accum,
-                    warp_tileB_k_compute_offset);
-
-                // Issue global->shared copies for the this stage
-                if (warp_mma_k < Base::kWarpGemmIterations - 1)
-                {
-                    int group_start_iteration_A, group_start_iteration_B;
-
-                    group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
-                    group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
-
-                    copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
-                }
-
-                if (warp_mma_k + 2 == Base::kWarpGemmIterations)
-                {
-                    int group_start_iteration_A, group_start_iteration_B;
-                    group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
-                    group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
-
-                    copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
-
-                    // Inserts a memory fence between stages of cp.async instructions.
-                    cutlass::arch::cp_async_fence();
-
-                    // Wait until we have at least one committed global fetch stage. (#uncommitted = Base::kStages - 1 -
-                    // #committed)
-                    arch::cp_async_wait<Base::kStages - 2>();
-                    __syncthreads();
-
-                    // Move to the next stage
-                    iterator_A.add_tile_offset({0, 1});
-                    iterator_B.add_tile_offset({1, 0});
-
-                    this->smem_iterator_A_.add_tile_offset({0, 1});
-                    this->smem_iterator_B_.add_tile_offset({1, 0});
-
-                    // Add negative offsets to return iterators to the 'start' of the
-                    // circular buffer in shared memory
-                    if (smem_write_stage_idx == (Base::kStages - 1))
-                    {
-                        this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
-                        this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
-                        smem_write_stage_idx = 0;
-                    }
-                    else
-                    {
-                        ++smem_write_stage_idx;
-                    }
-
-                    if (smem_read_stage_idx == (Base::kStages - 1))
-                    {
-                        this->warp_tile_iterator_A_.add_tile_offset(
-                            {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
-                        this->warp_tile_iterator_B_.add_tile_offset(
-                            {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0});
-                        smem_read_stage_idx = 0;
-                    }
-                    else
-                    {
-                        ++smem_read_stage_idx;
-                    }
-
-                    --gemm_k_iterations;
-                    iterator_A.clear_mask(gemm_k_iterations == 0);
-                    iterator_B.clear_mask(gemm_k_iterations == 0);
-                }
-            }
+        int const warp_tileB_k_compute_offset =
+            warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
+        int const warp_tileB_k_load_offset =
+            warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
+        if (warp_tileB_k_compute_offset ==
+            Base::kNumKIterationsPerWarpBLoad - 1) {
+          this->warp_tile_iterator_B_.set_kgroup_index(
+              (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
+          this->warp_tile_iterator_B_.load(
+              warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
+          ++this->warp_tile_iterator_B_;
         }
 
-        if (SharedMemoryClear == SharedMemoryClearOption::kZfill)
-        {
-            // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop
-            cutlass::arch::cp_async_fence();
-            cutlass::arch::cp_async_wait<0>();
-            __syncthreads();
+        typename TransformBAfterLDS::result_type converted_frag_B =
+            lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
+        warp_dequantizer_.dequantize(converted_frag_B, warp_frag_scales);
+
+        using FragmentOperandB =
+            cutlass::Array<ElementA, Operator::FragmentB::kElements>;
+        constexpr cutlass::FloatRoundStyle RoundStyle =
+            cutlass::FloatRoundStyle::round_to_nearest;
+        constexpr int ConversionVectorWidth =
+            TransformBAfterLDS::result_type::kElements;
+        static_assert(ConversionVectorWidth == FragmentOperandB::kElements);
+
+        using Converter = cutlass::NumericArrayConverter<ElementA,
+                                                         ElementScale,
+                                                         ConversionVectorWidth,
+                                                         RoundStyle>;
+
+        FragmentOperandB converted_frag_B_operand =
+            Converter::convert(converted_frag_B);
+        run_warp_mma(warp_mma,
+                     accum,
+                     warp_frag_A[warp_mma_k % 2],
+                     converted_frag_B_operand,
+                     accum,
+                     warp_tileB_k_compute_offset);
+
+        // Issue global->shared copies for the this stage
+        if (warp_mma_k < Base::kWarpGemmIterations - 1) {
+          int group_start_iteration_A, group_start_iteration_B;
+
+          group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
+          group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
+
+          copy_tiles_and_advance(iterator_A,
+                                 iterator_B,
+                                 group_start_iteration_A,
+                                 group_start_iteration_B);
         }
+
+        if (warp_mma_k + 2 == Base::kWarpGemmIterations) {
+          int group_start_iteration_A, group_start_iteration_B;
+          group_start_iteration_A =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
+          group_start_iteration_B =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
+
+          copy_tiles_and_advance(iterator_A,
+                                 iterator_B,
+                                 group_start_iteration_A,
+                                 group_start_iteration_B);
+
+          // Inserts a memory fence between stages of cp.async instructions.
+          cutlass::arch::cp_async_fence();
+
+          // Wait until we have at least one committed global fetch stage.
+          // (#uncommitted = Base::kStages - 1 - #committed)
+          arch::cp_async_wait<Base::kStages - 2>();
+          __syncthreads();
+
+          // Move to the next stage
+          iterator_A.add_tile_offset({0, 1});
+          iterator_B.add_tile_offset({1, 0});
+
+          this->smem_iterator_A_.add_tile_offset({0, 1});
+          this->smem_iterator_B_.add_tile_offset({1, 0});
+
+          // Add negative offsets to return iterators to the 'start' of the
+          // circular buffer in shared memory
+          if (smem_write_stage_idx == (Base::kStages - 1)) {
+            this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
+            this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
+            smem_write_stage_idx = 0;
+          } else {
+            ++smem_write_stage_idx;
+          }
+
+          if (smem_read_stage_idx == (Base::kStages - 1)) {
+            this->warp_tile_iterator_A_.add_tile_offset(
+                {0,
+                 -Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterations});
+            this->warp_tile_iterator_B_.add_tile_offset(
+                {-Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterationsForB,
+                 0});
+            smem_read_stage_idx = 0;
+          } else {
+            ++smem_read_stage_idx;
+          }
+
+          --gemm_k_iterations;
+          iterator_A.clear_mask(gemm_k_iterations == 0);
+          iterator_B.clear_mask(gemm_k_iterations == 0);
+        }
+      }
     }
+
+    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+      // commit and drain all pending and predicated LDGSTS pnz from the GEMM
+      // mainloop
+      cutlass::arch::cp_async_fence();
+      cutlass::arch::cp_async_wait<0>();
+      __syncthreads();
+    }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined.h
index bd3e38971b..dd7e8ae4b7 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -53,27 +54,27 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions.
+/// Structure to compute the matrix product targeting CUDA cores and SIMT math
+/// instructions.
 template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Iterates over tiles of A operand in global memory
-    //  (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator)
+    //  (concept: ReadableTileIterator | ForwardTileIterator |
+    //  MaskedTileIterator)
     typename IteratorA_,
     /// Iterates over tiles of A operand in shared memory
     /// (concept: WriteableTileIterator | RandomAccessTileIterator)
     typename SmemIteratorA_,
     /// Iterates over tiles of B operand in global memory
-    //  (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator)
+    //  (concept: ReadableTileIterator | ForwardTileIterator |
+    //  MaskedTileIterator)
     typename IteratorB_,
     /// Iterates over tiles of B operand in shared memory
     /// (concept: WriteableTileIterator | RandomAccessTileIterator)
@@ -98,9 +99,9 @@ template <
     typename Enable = void>
 class DqMmaPipelined;
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
 
 #include "cutlass_extensions/gemm/threadblock/dq_mma_pipelined_finegrained.h"
 #include "cutlass_extensions/gemm/threadblock/dq_mma_pipelined_percol.h"
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_finegrained.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_finegrained.h
index 50bdd0d85b..01fe06329b 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_finegrained.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_finegrained.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -53,27 +54,27 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions.
+/// Structure to compute the matrix product targeting CUDA cores and SIMT math
+/// instructions.
 template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Iterates over tiles of A operand in global memory
-    //  (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator)
+    //  (concept: ReadableTileIterator | ForwardTileIterator |
+    //  MaskedTileIterator)
     typename IteratorA_,
     /// Iterates over tiles of A operand in shared memory
     /// (concept: WriteableTileIterator | RandomAccessTileIterator)
     typename SmemIteratorA_,
     /// Iterates over tiles of B operand in global memory
-    //  (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator)
+    //  (concept: ReadableTileIterator | ForwardTileIterator |
+    //  MaskedTileIterator)
     typename IteratorB_,
     /// Iterates over tiles of B operand in shared memory
     /// (concept: WriteableTileIterator | RandomAccessTileIterator)
@@ -94,393 +95,442 @@ template <
     typename TransformBAfterLDS_,
     /// The quantization operator being used
     WeightOnlyQuantOp QuantOp_>
-class DqMmaPipelined<Shape_, IteratorA_, SmemIteratorA_, IteratorB_, SmemIteratorB_, IteratorScale_, SmemIteratorScale_,
-    ElementC_, LayoutC_, Policy_, TransformBAfterLDG_, TransformBAfterLDS_, QuantOp_,
-    std::enable_if_t<isFinegrained(QuantOp_)>>
-    : public DqMmaBase<Shape_, Policy_, typename SmemIteratorScale_::Element, 2, QuantOp_>
-{
-public:
-    ///< Base class
-    using Base = DqMmaBase<Shape_, Policy_, typename SmemIteratorScale_::Element, 2, QuantOp_>;
+class DqMmaPipelined<Shape_,
+                     IteratorA_,
+                     SmemIteratorA_,
+                     IteratorB_,
+                     SmemIteratorB_,
+                     IteratorScale_,
+                     SmemIteratorScale_,
+                     ElementC_,
+                     LayoutC_,
+                     Policy_,
+                     TransformBAfterLDG_,
+                     TransformBAfterLDS_,
+                     QuantOp_,
+                     std::enable_if_t<isFinegrained(QuantOp_)>>
+    : public DqMmaBase<Shape_,
+                       Policy_,
+                       typename SmemIteratorScale_::Element,
+                       2,
+                       QuantOp_> {
+ public:
+  ///< Base class
+  using Base = DqMmaBase<Shape_,
+                         Policy_,
+                         typename SmemIteratorScale_::Element,
+                         2,
+                         QuantOp_>;
 
-    using Shape = Shape_;         ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory
-    using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory
-    using ElementC = ElementC_;   ///< Data type of accumulator matrix
-    using LayoutC = LayoutC_;     ///< Layout of accumulator matrix
-    using Policy = Policy_;       ///< Policy describing tuning details
+  using Shape =
+      Shape_;  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using IteratorA =
+      IteratorA_;  ///< Iterates over tiles of A operand in global memory
+  using IteratorB =
+      IteratorB_;  ///< Iterates over tiles of B operand in global memory
+  using ElementC = ElementC_;  ///< Data type of accumulator matrix
+  using LayoutC = LayoutC_;    ///< Layout of accumulator matrix
+  using Policy = Policy_;      ///< Policy describing tuning details
 
-    using IteratorScale = IteratorScale_;
-    using ElementScale = typename IteratorScale::Element;
-    using LayoutScale = typename IteratorScale::Layout;
+  using IteratorScale = IteratorScale_;
+  using ElementScale = typename IteratorScale::Element;
+  using LayoutScale = typename IteratorScale::Layout;
 
-    using SmemIteratorA = SmemIteratorA_;
-    using SmemIteratorB = SmemIteratorB_;
-    using SmemIteratorScale = SmemIteratorScale_;
+  using SmemIteratorA = SmemIteratorA_;
+  using SmemIteratorB = SmemIteratorB_;
+  using SmemIteratorScale = SmemIteratorScale_;
 
-    using TransformBAfterLDG = TransformBAfterLDG_;
-    using TransformBAfterLDS = TransformBAfterLDS_;
+  using TransformBAfterLDG = TransformBAfterLDG_;
+  using TransformBAfterLDS = TransformBAfterLDS_;
 
-    static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
+  static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
 
-    //
-    // Dependent types
-    //
+  //
+  // Dependent types
+  //
 
-    /// Fragment of operand A loaded from global memory
-    using FragmentA = typename IteratorA::Fragment;
+  /// Fragment of operand A loaded from global memory
+  using FragmentA = typename IteratorA::Fragment;
 
-    /// Fragment of operand B loaded from global memory
-    using FragmentB = typename IteratorB::Fragment;
+  /// Fragment of operand B loaded from global memory
+  using FragmentB = typename IteratorB::Fragment;
 
-    /// Fragment of operand Scale loaded from global memory;
-    using FragmentScale = typename IteratorScale::Fragment;
+  /// Fragment of operand Scale loaded from global memory;
+  using FragmentScale = typename IteratorScale::Fragment;
 
-    /// Fragment of accumulator tile
-    using FragmentC = typename Policy::Operator::FragmentC;
+  /// Fragment of accumulator tile
+  using FragmentC = typename Policy::Operator::FragmentC;
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
 
-    /// Obtain the arch tag from the warp-level operator
-    using ArchTag = typename Policy::Operator::ArchTag;
+  /// Obtain the arch tag from the warp-level operator
+  using ArchTag = typename Policy::Operator::ArchTag;
 
-    using Dequantizer = warp::MmaTensorOpDequantizer<Operator, typename Base::WarpGemm, Operand::kB,
-        typename SmemIteratorScale::Element, LayoutScale, 32, QuantOp>;
+  using Dequantizer =
+      warp::MmaTensorOpDequantizer<Operator,
+                                   typename Base::WarpGemm,
+                                   Operand::kB,
+                                   typename SmemIteratorScale::Element,
+                                   LayoutScale,
+                                   32,
+                                   QuantOp>;
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = Operator::kTransformA;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = Operator::kTransformA;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = Operator::kTransformB;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = Operator::kTransformB;
 
-    // staticaly assert kStages for DqMmaPipelined is two (Double-buffered pipeline)
-    static_assert((Base::kStages == 2), "DqMmaPipelined requires kStages set to value 2");
+  // staticaly assert kStages for DqMmaPipelined is two (Double-buffered
+  // pipeline)
+  static_assert((Base::kStages == 2),
+                "DqMmaPipelined requires kStages set to value 2");
 
-    static_assert(Base::SharedStorage::ShapeScale::kRow == Base::kStages, "");
-    static_assert(Base::SharedStorage::ShapeScale::kColumn == Shape::kN, "");
+  static_assert(Base::SharedStorage::ShapeScale::kRow == Base::kStages, "");
+  static_assert(Base::SharedStorage::ShapeScale::kColumn == Shape::kN, "");
 
-private:
-    using WarpFragmentA = typename Operator::FragmentA;
-    using WarpFragmentB = typename Operator::FragmentB;
-    Dequantizer warp_dequantizer_;
+ private:
+  using WarpFragmentA = typename Operator::FragmentA;
+  using WarpFragmentB = typename Operator::FragmentB;
+  Dequantizer warp_dequantizer_;
 
-    using WarpFragmentScale = typename Dequantizer::FragmentScale;
-    using WarpFragmentZero = typename Dequantizer::FragmentZero;
+  using WarpFragmentScale = typename Dequantizer::FragmentScale;
+  using WarpFragmentZero = typename Dequantizer::FragmentZero;
 
-    using ElementA = typename IteratorA::Element;
-    using ElementB = typename IteratorB::Element;
-    using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
+  using ElementA = typename IteratorA::Element;
+  using ElementB = typename IteratorB::Element;
+  using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
 
-    static constexpr bool RequiresTileInterleave
-        = layout::IsColumnMajorTileInterleave<typename LayoutDetailsForB::Layout>::value;
-    static_assert(!RequiresTileInterleave || (RequiresTileInterleave && (Shape::kK == LayoutDetailsForB::ThreadblockK)),
-        "Layout K must match threadblockK");
+  static constexpr bool RequiresTileInterleave =
+      layout::IsColumnMajorTileInterleave<
+          typename LayoutDetailsForB::Layout>::value;
+  static_assert(!RequiresTileInterleave ||
+                    (RequiresTileInterleave &&
+                     (Shape::kK == LayoutDetailsForB::ThreadblockK)),
+                "Layout K must match threadblockK");
 
-protected:
-    /// Iterator to write threadblock-scoped tile of A operand to shared memory
-    SmemIteratorA smem_iterator_A_;
+ protected:
+  /// Iterator to write threadblock-scoped tile of A operand to shared memory
+  SmemIteratorA smem_iterator_A_;
 
-    /// Iterator to write threadblock-scoped tile of B operand to shared memory
-    SmemIteratorB smem_iterator_B_;
+  /// Iterator to write threadblock-scoped tile of B operand to shared memory
+  SmemIteratorB smem_iterator_B_;
 
-    /// Iterator to write threadblock-scoped tile of scale and zero operand to shared memory
-    SmemIteratorScale smem_iterator_scale_;
+  /// Iterator to write threadblock-scoped tile of scale and zero operand to
+  /// shared memory
+  SmemIteratorScale smem_iterator_scale_;
 
-public:
-    /// Construct from tensor references
-    CUTLASS_DEVICE
-    DqMmaPipelined(typename Base::SharedStorage&
-                       shared_storage, ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        int const group_size,          ///< The group size for quantization
-        int thread_idx,                ///< ID within the threadblock
-        int warp_idx,                  ///< ID of warp
-        int lane_idx                   ///< ID of each thread within a warp
-        )
-        : Base(shared_storage, thread_idx, warp_idx, lane_idx)
-        , warp_dequantizer_({shared_storage.operand_scale.data(), LayoutScale(Shape::kN)},
-              {shared_storage.operand_zero.data(), LayoutScale(Shape::kN)},
-              (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) / Base::WarpCount::kM, lane_idx)
-        , smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx)
-        , smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx)
-        , smem_iterator_scale_(LayoutScale(Shape::kN), shared_storage.operand_scale.data(),
-              shared_storage.operand_zero.data(), {Base::kStages, Shape::kN}, thread_idx, group_size)
-    {
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  DqMmaPipelined(typename Base::SharedStorage&
+                     shared_storage,    ///< Shared storage needed for internal
+                                        ///< use by threadblock-scoped GEMM
+                 int const group_size,  ///< The group size for quantization
+                 int thread_idx,        ///< ID within the threadblock
+                 int warp_idx,          ///< ID of warp
+                 int lane_idx           ///< ID of each thread within a warp
+                 )
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
+        warp_dequantizer_(
+            {shared_storage.operand_scale.data(), LayoutScale(Shape::kN)},
+            {shared_storage.operand_zero.data(), LayoutScale(Shape::kN)},
+            (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) /
+                Base::WarpCount::kM,
+            lane_idx),
+        smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
+        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx),
+        smem_iterator_scale_(LayoutScale(Shape::kN),
+                             shared_storage.operand_scale.data(),
+                             shared_storage.operand_zero.data(),
+                             {Base::kStages, Shape::kN},
+                             thread_idx,
+                             group_size) {
+    // Compute warp location within threadblock tile by mapping the warp_id to
+    // three coordinates:
+    //   _m: the warp's position within the threadblock along the M dimension
+    //   _n: the warp's position within the threadblock along the N dimension
+    //   _k: the warp's position within the threadblock along the K dimension
 
-        // Compute warp location within threadblock tile by mapping the warp_id to
-        // three coordinates:
-        //   _m: the warp's position within the threadblock along the M dimension
-        //   _n: the warp's position within the threadblock along the N dimension
-        //   _k: the warp's position within the threadblock along the K dimension
+    int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
 
-        int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
-        int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
+    int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
 
-        int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
-        int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
+    // Add per-warp offsets in units of warp-level tiles
+    this->warp_tile_iterator_A_.add_tile_offset(
+        {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
+    this->warp_tile_iterator_B_.add_tile_offset(
+        {Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  }
 
-        // Add per-warp offsets in units of warp-level tiles
-        this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
-        this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  CUTLASS_DEVICE
+  void copy_scales_and_advance(IteratorScale& iterator_scale) {
+    using TransformScale =
+        NumericArrayConverter<typename SmemIteratorScale::Element,
+                              typename FragmentScale::Element,
+                              FragmentScale::kElements>;
+
+    FragmentScale tb_frag_scales;
+    FragmentScale tb_frag_zeros;
+    tb_frag_scales.clear();
+    tb_frag_zeros.clear();
+
+    TransformScale transformScale;
+
+    using FragmentElement = typename FragmentScale::Element;
+
+    auto gmem_scale_ptr = iterator_scale.get_scale();
+    auto gmem_zero_ptr = iterator_scale.get_zero();
+
+    arch::global_load<FragmentScale, sizeof(FragmentScale)>(
+        tb_frag_scales, gmem_scale_ptr, iterator_scale.valid());
+
+    if (gmem_zero_ptr != nullptr) {
+      arch::global_load<FragmentScale, sizeof(FragmentScale)>(
+          tb_frag_zeros, gmem_zero_ptr, iterator_scale.valid());
     }
 
-    CUTLASS_DEVICE
-    void copy_scales_and_advance(IteratorScale& iterator_scale)
-    {
-        using TransformScale = NumericArrayConverter<typename SmemIteratorScale::Element,
-            typename FragmentScale::Element, FragmentScale::kElements>;
+    typename TransformScale::result_type tb_frag_scales_fp16 =
+        transformScale(tb_frag_scales);
+    typename TransformScale::result_type tb_frag_zeros_fp16;
+    if (gmem_zero_ptr != nullptr)
+      tb_frag_zeros_fp16 = transformScale(tb_frag_zeros);
 
-        FragmentScale tb_frag_scales;
-        FragmentScale tb_frag_zeros;
-        tb_frag_scales.clear();
-        tb_frag_zeros.clear();
+    auto frag_scale_ptr_fp16 =
+        reinterpret_cast<typename SmemIteratorScale::Element*>(
+            &tb_frag_scales_fp16);
+    auto frag_zero_ptr_fp16 =
+        reinterpret_cast<typename SmemIteratorScale::Element*>(
+            &tb_frag_zeros_fp16);
+    auto smem_scale_ptr = this->smem_iterator_scale_.get_scale();
+    auto smem_zero_ptr = this->smem_iterator_scale_.get_zero();
 
-        TransformScale transformScale;
+    if (iterator_scale.valid()) {
+      auto smem_offset = cast_smem_ptr_to_uint(smem_scale_ptr);
+      arch::shared_store<sizeof(FragmentScale)>(smem_offset,
+                                                frag_scale_ptr_fp16);
 
-        using FragmentElement = typename FragmentScale::Element;
-
-        auto gmem_scale_ptr = iterator_scale.get_scale();
-        auto gmem_zero_ptr = iterator_scale.get_zero();
-
-        arch::global_load<FragmentScale, sizeof(FragmentScale)>(tb_frag_scales, gmem_scale_ptr, iterator_scale.valid());
-
-        if (gmem_zero_ptr != nullptr)
-        {
-            arch::global_load<FragmentScale, sizeof(FragmentScale)>(
-                tb_frag_zeros, gmem_zero_ptr, iterator_scale.valid());
-        }
-
-        typename TransformScale::result_type tb_frag_scales_fp16 = transformScale(tb_frag_scales);
-        typename TransformScale::result_type tb_frag_zeros_fp16;
-        if (gmem_zero_ptr != nullptr)
-            tb_frag_zeros_fp16 = transformScale(tb_frag_zeros);
-
-        auto frag_scale_ptr_fp16 = reinterpret_cast<typename SmemIteratorScale::Element*>(&tb_frag_scales_fp16);
-        auto frag_zero_ptr_fp16 = reinterpret_cast<typename SmemIteratorScale::Element*>(&tb_frag_zeros_fp16);
-        auto smem_scale_ptr = this->smem_iterator_scale_.get_scale();
-        auto smem_zero_ptr = this->smem_iterator_scale_.get_zero();
-
-        if (iterator_scale.valid())
-        {
-            auto smem_offset = cast_smem_ptr_to_uint(smem_scale_ptr);
-            arch::shared_store<sizeof(FragmentScale)>(smem_offset, frag_scale_ptr_fp16);
-
-            if (gmem_zero_ptr != nullptr)
-            {
-                smem_offset = cast_smem_ptr_to_uint(smem_zero_ptr);
-                arch::shared_store<sizeof(FragmentScale)>(smem_offset, frag_zero_ptr_fp16);
-            }
-        }
-
-        if (iterator_scale.group_size_ == 64)
-        {
-            iterator_scale.add_tile_offset({1, 0});
-        }
-        else if (iterator_scale.group_size_ == 128)
-        {
-            if constexpr (Shape::kK == 128)
-            {
-                iterator_scale.add_tile_offset({1, 0});
-            }
-            else if constexpr (Shape::kK == 64)
-            {
-                if (iterator_scale.row_groupsize64_ & 0x1)
-                {
-                    iterator_scale.add_tile_offset({1, 0});
-                }
-            }
-            else
-            {
-                static_assert(Shape::kK == 0, "Unsupported k tile shape, can only be 64 or 128");
-            }
-        }
-
-        iterator_scale.row_groupsize64_++;
-
-        this->smem_iterator_scale_.add_tile_offset({1, 0});
+      if (gmem_zero_ptr != nullptr) {
+        smem_offset = cast_smem_ptr_to_uint(smem_zero_ptr);
+        arch::shared_store<sizeof(FragmentScale)>(smem_offset,
+                                                  frag_zero_ptr_fp16);
+      }
     }
 
-    /// Perform a threadblock-scoped matrix multiply-accumulate
-    CUTLASS_DEVICE
-    void operator()(int gemm_k_iterations, ///< number of iterations of the mainloop
-        FragmentC& accum,                  ///< destination accumulator tile
-        IteratorA iterator_A,              ///< iterator over A operand in global memory
-        IteratorB iterator_B,              ///< iterator over B operand in global memory
-        IteratorScale iterator_scale,      ///< iterator over scale operand in global memory
-        FragmentC const& src_accum)
-    {                                      ///< source accumulator tile
+    if (iterator_scale.group_size_ == 64) {
+      iterator_scale.add_tile_offset({1, 0});
+    } else if (iterator_scale.group_size_ == 128) {
+      if constexpr (Shape::kK == 128) {
+        iterator_scale.add_tile_offset({1, 0});
+      } else if constexpr (Shape::kK == 64) {
+        if (iterator_scale.row_groupsize64_ & 0x1) {
+          iterator_scale.add_tile_offset({1, 0});
+        }
+      } else {
+        static_assert(Shape::kK == 0,
+                      "Unsupported k tile shape, can only be 64 or 128");
+      }
+    }
 
-        //
-        // Prologue
-        //
-        TransformBAfterLDG ldg_converter;
-        TransformBAfterLDS lds_converter;
+    iterator_scale.row_groupsize64_++;
 
-        using TransformA
-            = NumericArrayConverter<typename WarpFragmentA::Element, typename FragmentA::Element, FragmentA::kElements>;
+    this->smem_iterator_scale_.add_tile_offset({1, 0});
+  }
 
-        // These transforms are mainly to handle when we have bfloat activations and weights in GMEM and want
-        // to issue HMMA on architectures older than Ampere. We will convert to FP16 before STS.
-        TransformA transformA;
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  CUTLASS_DEVICE
+  void operator()(
+      int gemm_k_iterations,  ///< number of iterations of the mainloop
+      FragmentC& accum,       ///< destination accumulator tile
+      IteratorA iterator_A,   ///< iterator over A operand in global memory
+      IteratorB iterator_B,   ///< iterator over B operand in global memory
+      IteratorScale
+          iterator_scale,  ///< iterator over scale operand in global memory
+      FragmentC const& src_accum) {  ///< source accumulator tile
 
-        // Perform accumulation in the 'd' output operand
-        accum = src_accum;
+    //
+    // Prologue
+    //
+    TransformBAfterLDG ldg_converter;
+    TransformBAfterLDS lds_converter;
 
-        FragmentA tb_frag_A;
-        FragmentB tb_frag_B;
+    using TransformA = NumericArrayConverter<typename WarpFragmentA::Element,
+                                             typename FragmentA::Element,
+                                             FragmentA::kElements>;
 
-        tb_frag_A.clear();
-        tb_frag_B.clear();
+    // These transforms are mainly to handle when we have bfloat activations and
+    // weights in GMEM and want to issue HMMA on architectures older than
+    // Ampere. We will convert to FP16 before STS.
+    TransformA transformA;
 
-        // The last kblock is loaded in the prolog
-        iterator_A.load(tb_frag_A);
-        iterator_B.load(tb_frag_B);
+    // Perform accumulation in the 'd' output operand
+    accum = src_accum;
 
-        ++iterator_A;
-        ++iterator_B;
+    FragmentA tb_frag_A;
+    FragmentB tb_frag_B;
 
-        this->smem_iterator_A_.store(transformA(tb_frag_A));
-        this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
+    tb_frag_A.clear();
+    tb_frag_B.clear();
 
-        ++this->smem_iterator_A_;
-        ++this->smem_iterator_B_;
+    // The last kblock is loaded in the prolog
+    iterator_A.load(tb_frag_A);
+    iterator_B.load(tb_frag_B);
 
-        copy_scales_and_advance(iterator_scale);
+    ++iterator_A;
+    ++iterator_B;
 
-        __syncthreads();
+    this->smem_iterator_A_.store(transformA(tb_frag_A));
+    this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
 
-        // Pair of fragments used to overlap shared memory loads and math instructions
-        WarpFragmentA warp_frag_A[2];
-        WarpFragmentB warp_frag_B[2];
-        WarpFragmentScale warp_frag_scales;
-        WarpFragmentZero warp_frag_zero;
+    ++this->smem_iterator_A_;
+    ++this->smem_iterator_B_;
 
-        this->warp_tile_iterator_A_.set_kgroup_index(0);
-        this->warp_tile_iterator_B_.set_kgroup_index(0);
+    copy_scales_and_advance(iterator_scale);
 
-        this->warp_tile_iterator_A_.load(warp_frag_A[0]);
-        this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+    __syncthreads();
 
-        warp_dequantizer_.load(warp_frag_scales, warp_frag_zero);
+    // Pair of fragments used to overlap shared memory loads and math
+    // instructions
+    WarpFragmentA warp_frag_A[2];
+    WarpFragmentB warp_frag_B[2];
+    WarpFragmentScale warp_frag_scales;
+    WarpFragmentZero warp_frag_zero;
 
+    this->warp_tile_iterator_A_.set_kgroup_index(0);
+    this->warp_tile_iterator_B_.set_kgroup_index(0);
+
+    this->warp_tile_iterator_A_.load(warp_frag_A[0]);
+    this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+
+    warp_dequantizer_.load(warp_frag_scales, warp_frag_zero);
+
+    ++this->warp_tile_iterator_A_;
+    ++this->warp_tile_iterator_B_;
+    warp_dequantizer_.add_pointer_offset(Shape::kN);
+
+    Operator warp_mma;
+
+    int smem_write_stage_idx = 1;
+
+    // Avoid reading out of bounds
+    iterator_A.clear_mask(gemm_k_iterations <= 1);
+    iterator_B.clear_mask(gemm_k_iterations <= 1);
+    iterator_scale.clear_mask(gemm_k_iterations <= 1);
+
+    // Issue loads during the first warp-level matrix multiply-add *AFTER*
+    // issuing shared memory loads (which have the tighest latency requirement).
+
+    //
+    // Mainloop
+    //
+
+    // Note: The main loop does not support Base::kWarpGemmIterations == 2.
+    CUTLASS_GEMM_LOOP
+    for (; gemm_k_iterations > 0; --gemm_k_iterations) {
+      //
+      // Loop over GEMM K dimension
+      //
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
+           ++warp_mma_k) {
+        // Load warp-level tiles from shared memory, wrapping to k offset if
+        // this is the last group as the case may be.
+
+        if (warp_mma_k == Base::kWarpGemmIterations - 1) {
+          // Write fragments to shared memory
+          this->smem_iterator_A_.store(transformA(tb_frag_A));
+
+          this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
+
+          __syncthreads();
+
+          ++this->smem_iterator_A_;
+          ++this->smem_iterator_B_;
+
+          // Add negative offsets to return iterators to the 'start' of the
+          // circular buffer in shared memory
+          if (smem_write_stage_idx == 1) {
+            this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
+            this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
+            this->smem_iterator_scale_.add_tile_offset({-Base::kStages, 0});
+          } else {
+            this->warp_tile_iterator_A_.add_tile_offset(
+                {0,
+                 -Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterations});
+            this->warp_tile_iterator_B_.add_tile_offset(
+                {-Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterationsForB,
+                 0});
+            warp_dequantizer_.add_pointer_offset(-Base::kStages * Shape::kN);
+          }
+
+          smem_write_stage_idx ^= 1;
+        }
+
+        this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) %
+                                                     Base::kWarpGemmIterations);
+        this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
         ++this->warp_tile_iterator_A_;
-        ++this->warp_tile_iterator_B_;
-        warp_dequantizer_.add_pointer_offset(Shape::kN);
 
-        Operator warp_mma;
-
-        int smem_write_stage_idx = 1;
-
-        // Avoid reading out of bounds
-        iterator_A.clear_mask(gemm_k_iterations <= 1);
-        iterator_B.clear_mask(gemm_k_iterations <= 1);
-        iterator_scale.clear_mask(gemm_k_iterations <= 1);
-
-        // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing
-        // shared memory loads (which have the tighest latency requirement).
-
-        //
-        // Mainloop
-        //
-
-        // Note: The main loop does not support Base::kWarpGemmIterations == 2.
-        CUTLASS_GEMM_LOOP
-        for (; gemm_k_iterations > 0; --gemm_k_iterations)
-        {
-            //
-            // Loop over GEMM K dimension
-            //
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k)
-            {
-
-                // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group
-                // as the case may be.
-
-                if (warp_mma_k == Base::kWarpGemmIterations - 1)
-                {
-
-                    // Write fragments to shared memory
-                    this->smem_iterator_A_.store(transformA(tb_frag_A));
-
-                    this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
-
-                    __syncthreads();
-
-                    ++this->smem_iterator_A_;
-                    ++this->smem_iterator_B_;
-
-                    // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory
-                    if (smem_write_stage_idx == 1)
-                    {
-                        this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
-                        this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
-                        this->smem_iterator_scale_.add_tile_offset({-Base::kStages, 0});
-                    }
-                    else
-                    {
-                        this->warp_tile_iterator_A_.add_tile_offset(
-                            {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
-                        this->warp_tile_iterator_B_.add_tile_offset(
-                            {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0});
-                        warp_dequantizer_.add_pointer_offset(-Base::kStages * Shape::kN);
-                    }
-
-                    smem_write_stage_idx ^= 1;
-                }
-
-                this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
-                this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
-                ++this->warp_tile_iterator_A_;
-
-                int const warp_tileB_k_compute_offset = warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
-                int const warp_tileB_k_load_offset = warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
-                // We are just about to finish computing on a fragment of B, so initiate the load for the next fragment.
-                if (warp_tileB_k_compute_offset == Base::kNumKIterationsPerWarpBLoad - 1)
-                {
-                    this->warp_tile_iterator_B_.set_kgroup_index(
-                        (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
-                    this->warp_tile_iterator_B_.load(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
-                    ++this->warp_tile_iterator_B_;
-                }
-
-                if (warp_mma_k == 0)
-                {
-
-                    iterator_A.load(tb_frag_A);
-                    iterator_B.load(tb_frag_B);
-
-                    ++iterator_A;
-                    ++iterator_B;
-
-                    copy_scales_and_advance(iterator_scale);
-
-                    // Avoid reading out of bounds if this was the last loop iteration
-                    iterator_A.clear_mask(gemm_k_iterations <= 2);
-                    iterator_B.clear_mask(gemm_k_iterations <= 2);
-                    iterator_scale.clear_mask(gemm_k_iterations <= 2);
-                }
-
-                typename TransformBAfterLDS::result_type converted_frag_B
-                    = lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
-                warp_dequantizer_.dequantize(converted_frag_B, warp_frag_scales, warp_frag_zero);
-                run_warp_mma(
-                    warp_mma, accum, warp_frag_A[warp_mma_k % 2], converted_frag_B, accum, warp_tileB_k_compute_offset);
-            }
-
-            // Load the scales needed for the next tile iteration
-            warp_dequantizer_.load(warp_frag_scales, warp_frag_zero);
-            // Update internal pointer to the set of scales in shared memory
-            warp_dequantizer_.add_pointer_offset(Shape::kN);
+        int const warp_tileB_k_compute_offset =
+            warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
+        int const warp_tileB_k_load_offset =
+            warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
+        // We are just about to finish computing on a fragment of B, so initiate
+        // the load for the next fragment.
+        if (warp_tileB_k_compute_offset ==
+            Base::kNumKIterationsPerWarpBLoad - 1) {
+          this->warp_tile_iterator_B_.set_kgroup_index(
+              (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
+          this->warp_tile_iterator_B_.load(
+              warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
+          ++this->warp_tile_iterator_B_;
         }
+
+        if (warp_mma_k == 0) {
+          iterator_A.load(tb_frag_A);
+          iterator_B.load(tb_frag_B);
+
+          ++iterator_A;
+          ++iterator_B;
+
+          copy_scales_and_advance(iterator_scale);
+
+          // Avoid reading out of bounds if this was the last loop iteration
+          iterator_A.clear_mask(gemm_k_iterations <= 2);
+          iterator_B.clear_mask(gemm_k_iterations <= 2);
+          iterator_scale.clear_mask(gemm_k_iterations <= 2);
+        }
+
+        typename TransformBAfterLDS::result_type converted_frag_B =
+            lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
+        warp_dequantizer_.dequantize(
+            converted_frag_B, warp_frag_scales, warp_frag_zero);
+        run_warp_mma(warp_mma,
+                     accum,
+                     warp_frag_A[warp_mma_k % 2],
+                     converted_frag_B,
+                     accum,
+                     warp_tileB_k_compute_offset);
+      }
+
+      // Load the scales needed for the next tile iteration
+      warp_dequantizer_.load(warp_frag_scales, warp_frag_zero);
+      // Update internal pointer to the set of scales in shared memory
+      warp_dequantizer_.add_pointer_offset(Shape::kN);
     }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_percol.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_percol.h
index 316ea9f80a..e6f512edfe 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_percol.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_percol.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -53,27 +54,27 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions.
+/// Structure to compute the matrix product targeting CUDA cores and SIMT math
+/// instructions.
 template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Iterates over tiles of A operand in global memory
-    //  (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator)
+    //  (concept: ReadableTileIterator | ForwardTileIterator |
+    //  MaskedTileIterator)
     typename IteratorA_,
     /// Iterates over tiles of A operand in shared memory
     /// (concept: WriteableTileIterator | RandomAccessTileIterator)
     typename SmemIteratorA_,
     /// Iterates over tiles of B operand in global memory
-    //  (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator)
+    //  (concept: ReadableTileIterator | ForwardTileIterator |
+    //  MaskedTileIterator)
     typename IteratorB_,
     /// Iterates over tiles of B operand in shared memory
     /// (concept: WriteableTileIterator | RandomAccessTileIterator)
@@ -94,306 +95,361 @@ template <
     typename TransformBAfterLDS_,
     /// The quantization operator being used
     WeightOnlyQuantOp QuantOp_>
-class DqMmaPipelined<Shape_, IteratorA_, SmemIteratorA_, IteratorB_, SmemIteratorB_, IteratorScale_, SmemIteratorScale_,
-    ElementC_, LayoutC_, Policy_, TransformBAfterLDG_, TransformBAfterLDS_, QuantOp_,
-    std::enable_if_t<!isFinegrained(QuantOp_)>>
-    : public DqMmaBase<Shape_, Policy_, typename SmemIteratorScale_::Element, 2, QuantOp_>
-{
-public:
-    ///< Base class
-    using Base = DqMmaBase<Shape_, Policy_, typename SmemIteratorScale_::Element, 2, QuantOp_>;
+class DqMmaPipelined<Shape_,
+                     IteratorA_,
+                     SmemIteratorA_,
+                     IteratorB_,
+                     SmemIteratorB_,
+                     IteratorScale_,
+                     SmemIteratorScale_,
+                     ElementC_,
+                     LayoutC_,
+                     Policy_,
+                     TransformBAfterLDG_,
+                     TransformBAfterLDS_,
+                     QuantOp_,
+                     std::enable_if_t<!isFinegrained(QuantOp_)>>
+    : public DqMmaBase<Shape_,
+                       Policy_,
+                       typename SmemIteratorScale_::Element,
+                       2,
+                       QuantOp_> {
+ public:
+  ///< Base class
+  using Base = DqMmaBase<Shape_,
+                         Policy_,
+                         typename SmemIteratorScale_::Element,
+                         2,
+                         QuantOp_>;
 
-    using Shape = Shape_;         ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory
-    using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory
-    using ElementC = ElementC_;   ///< Data type of accumulator matrix
-    using LayoutC = LayoutC_;     ///< Layout of accumulator matrix
-    using Policy = Policy_;       ///< Policy describing tuning details
+  using Shape =
+      Shape_;  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using IteratorA =
+      IteratorA_;  ///< Iterates over tiles of A operand in global memory
+  using IteratorB =
+      IteratorB_;  ///< Iterates over tiles of B operand in global memory
+  using ElementC = ElementC_;  ///< Data type of accumulator matrix
+  using LayoutC = LayoutC_;    ///< Layout of accumulator matrix
+  using Policy = Policy_;      ///< Policy describing tuning details
 
-    using IteratorScale = IteratorScale_;
-    using ElementScale = typename IteratorScale::Element;
-    using LayoutScale = typename IteratorScale::Layout;
+  using IteratorScale = IteratorScale_;
+  using ElementScale = typename IteratorScale::Element;
+  using LayoutScale = typename IteratorScale::Layout;
 
-    using SmemIteratorA = SmemIteratorA_;
-    using SmemIteratorB = SmemIteratorB_;
-    using SmemIteratorScale = SmemIteratorScale_;
+  using SmemIteratorA = SmemIteratorA_;
+  using SmemIteratorB = SmemIteratorB_;
+  using SmemIteratorScale = SmemIteratorScale_;
 
-    using TransformBAfterLDG = TransformBAfterLDG_;
-    using TransformBAfterLDS = TransformBAfterLDS_;
+  using TransformBAfterLDG = TransformBAfterLDG_;
+  using TransformBAfterLDS = TransformBAfterLDS_;
 
-    static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
+  static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
+
+  //
+  // Dependent types
+  //
+
+  /// Fragment of operand A loaded from global memory
+  using FragmentA = typename IteratorA::Fragment;
+
+  /// Fragment of operand B loaded from global memory
+  using FragmentB = typename IteratorB::Fragment;
+
+  /// Fragment of operand Scale loaded from global memory;
+  using FragmentScale = typename IteratorScale::Fragment;
+
+  /// Fragment of accumulator tile
+  using FragmentC = typename Policy::Operator::FragmentC;
+
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
+
+  /// Obtain the arch tag from the warp-level operator
+  using ArchTag = typename Policy::Operator::ArchTag;
+
+  using Dequantizer = warp::MmaTensorOpDequantizer<
+      Operator,
+      typename Base::WarpGemm,
+      Operand::kB,
+      typename SmemIteratorScale::Fragment::Element,
+      LayoutScale,
+      32,
+      QuantOp>;
+
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = Operator::kTransformA;
+
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = Operator::kTransformB;
+
+  // staticaly assert kStages for DqMmaPipelined is two (Double-buffered
+  // pipeline)
+  static_assert((Base::kStages == 2),
+                "DqMmaPipelined requires kStages set to value 2");
+
+ private:
+  using WarpFragmentA = typename Operator::FragmentA;
+  using WarpFragmentB = typename Operator::FragmentB;
+  Dequantizer warp_dequantizer_;
+
+  using ElementA = typename IteratorA::Element;
+  using ElementB = typename IteratorB::Element;
+  using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
+
+  static constexpr bool RequiresTileInterleave =
+      layout::IsColumnMajorTileInterleave<
+          typename LayoutDetailsForB::Layout>::value;
+  static_assert(!RequiresTileInterleave ||
+                    (RequiresTileInterleave &&
+                     (Shape::kK == LayoutDetailsForB::ThreadblockK)),
+                "Layout K must match threadblockK");
+
+ protected:
+  /// Iterator to write threadblock-scoped tile of A operand to shared memory
+  SmemIteratorA smem_iterator_A_;
+
+  /// Iterator to write threadblock-scoped tile of B operand to shared memory
+  SmemIteratorB smem_iterator_B_;
+
+  /// Iterator to write threadblock-scoped tile of scale operand to shared
+  /// memory
+  SmemIteratorScale smem_iterator_scale_;
+
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  DqMmaPipelined(
+      typename Base::SharedStorage&
+          shared_storage,  ///< Shared storage needed for internal use by
+                           ///< threadblock-scoped GEMM
+      int const
+          group_size,  ///< Will not be used, just to adapt to finegrained
+                       ///< modifications and make the compilation successful.
+                       ///< Because DqMmaPipelined is only enabled for sm<80, so
+                       ///< even if this argument is not added, it does not
+                       ///< affect compilation for sm>=80.
+      int thread_idx,  ///< ID within the threadblock
+      int warp_idx,    ///< ID of warp
+      int lane_idx     ///< ID of each thread within a warp
+      )
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
+        warp_dequantizer_(
+            {shared_storage.operand_scale.data(), LayoutScale(Shape::kN)},
+            (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) /
+                Base::WarpCount::kM,
+            lane_idx),
+        smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
+        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx),
+        smem_iterator_scale_(LayoutScale(Shape::kN),
+                             shared_storage.operand_scale.data(),
+                             {1, Shape::kN},
+                             thread_idx) {
+    // Compute warp location within threadblock tile by mapping the warp_id to
+    // three coordinates:
+    //   _m: the warp's position within the threadblock along the M dimension
+    //   _n: the warp's position within the threadblock along the N dimension
+    //   _k: the warp's position within the threadblock along the K dimension
+
+    int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
+
+    int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
+    int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
+
+    // Add per-warp offsets in units of warp-level tiles
+    this->warp_tile_iterator_A_.add_tile_offset(
+        {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
+    this->warp_tile_iterator_B_.add_tile_offset(
+        {Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  }
+
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  CUTLASS_DEVICE
+  void operator()(
+      int gemm_k_iterations,  ///< number of iterations of the mainloop
+      FragmentC& accum,       ///< destination accumulator tile
+      IteratorA iterator_A,   ///< iterator over A operand in global memory
+      IteratorB iterator_B,   ///< iterator over B operand in global memory
+      IteratorScale
+          iterator_scale,  ///< iterator over scale operand in global memory
+      FragmentC const& src_accum) {  ///< source accumulator tile
 
     //
-    // Dependent types
+    // Prologue
+    //
+    TransformBAfterLDG ldg_converter;
+    TransformBAfterLDS lds_converter;
+
+    using TransformA = NumericArrayConverter<typename WarpFragmentA::Element,
+                                             typename FragmentA::Element,
+                                             FragmentA::kElements>;
+
+    using TransformScale =
+        NumericArrayConverter<typename SmemIteratorScale::Fragment::Element,
+                              typename FragmentScale::Element,
+                              FragmentScale::kElements>;
+
+    // These transforms are mainly to handle when we have bfloat activations and
+    // weights in GMEM and want to issue HMMA on architectures older than
+    // Ampere. We will convert to FP16 before STS.
+    TransformA transformA;
+    TransformScale transformScale;
+
+    // Perform accumulation in the 'd' output operand
+    accum = src_accum;
+
+    FragmentA tb_frag_A;
+    FragmentB tb_frag_B;
+    FragmentScale tb_frag_scales;
+
+    using WarpFragmentScale = typename Dequantizer::FragmentScale;
+    WarpFragmentScale warp_frag_scales;
+
+    tb_frag_A.clear();
+    tb_frag_B.clear();
+    tb_frag_scales.clear();
+
+    // The last kblock is loaded in the prolog
+    iterator_A.load(tb_frag_A);
+    iterator_B.load(tb_frag_B);
+    iterator_scale.load(tb_frag_scales);
+
+    ++iterator_A;
+    ++iterator_B;
+
+    this->smem_iterator_A_.store(transformA(tb_frag_A));
+    this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
+    this->smem_iterator_scale_.store(transformScale(tb_frag_scales));
+
+    ++this->smem_iterator_A_;
+    ++this->smem_iterator_B_;
+
+    __syncthreads();
+
+    warp_dequantizer_.load(warp_frag_scales);
+
+    // Pair of fragments used to overlap shared memory loads and math
+    // instructions
+    WarpFragmentA warp_frag_A[2];
+    WarpFragmentB warp_frag_B[2];
+
+    this->warp_tile_iterator_A_.set_kgroup_index(0);
+    this->warp_tile_iterator_B_.set_kgroup_index(0);
+
+    this->warp_tile_iterator_A_.load(warp_frag_A[0]);
+    this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+
+    ++this->warp_tile_iterator_A_;
+    ++this->warp_tile_iterator_B_;
+
+    Operator warp_mma;
+
+    int smem_write_stage_idx = 1;
+
+    // Avoid reading out of bounds
+    iterator_A.clear_mask(gemm_k_iterations <= 1);
+    iterator_B.clear_mask(gemm_k_iterations <= 1);
+
+    // Issue loads during the first warp-level matrix multiply-add *AFTER*
+    // issuing shared memory loads (which have the tighest latency requirement).
+
+    //
+    // Mainloop
     //
 
-    /// Fragment of operand A loaded from global memory
-    using FragmentA = typename IteratorA::Fragment;
+    // Note: The main loop does not support Base::kWarpGemmIterations == 2.
+    CUTLASS_GEMM_LOOP
+    for (; gemm_k_iterations > 0; --gemm_k_iterations) {
+      //
+      // Loop over GEMM K dimension
+      //
 
-    /// Fragment of operand B loaded from global memory
-    using FragmentB = typename IteratorB::Fragment;
+      CUTLASS_PRAGMA_UNROLL
+      for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
+           ++warp_mma_k) {
+        // Load warp-level tiles from shared memory, wrapping to k offset if
+        // this is the last group as the case may be.
 
-    /// Fragment of operand Scale loaded from global memory;
-    using FragmentScale = typename IteratorScale::Fragment;
+        if (warp_mma_k == Base::kWarpGemmIterations - 1) {
+          // Write fragments to shared memory
+          this->smem_iterator_A_.store(transformA(tb_frag_A));
 
-    /// Fragment of accumulator tile
-    using FragmentC = typename Policy::Operator::FragmentC;
+          this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+          __syncthreads();
 
-    /// Obtain the arch tag from the warp-level operator
-    using ArchTag = typename Policy::Operator::ArchTag;
+          ++this->smem_iterator_A_;
+          ++this->smem_iterator_B_;
 
-    using Dequantizer = warp::MmaTensorOpDequantizer<Operator, typename Base::WarpGemm, Operand::kB,
-        typename SmemIteratorScale::Fragment::Element, LayoutScale, 32, QuantOp>;
+          // Add negative offsets to return iterators to the 'start' of the
+          // circular buffer in shared memory
+          if (smem_write_stage_idx == 1) {
+            this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
+            this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
+          } else {
+            this->warp_tile_iterator_A_.add_tile_offset(
+                {0,
+                 -Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterations});
+            this->warp_tile_iterator_B_.add_tile_offset(
+                {-Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterationsForB,
+                 0});
+          }
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = Operator::kTransformA;
-
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = Operator::kTransformB;
-
-    // staticaly assert kStages for DqMmaPipelined is two (Double-buffered pipeline)
-    static_assert((Base::kStages == 2), "DqMmaPipelined requires kStages set to value 2");
-
-private:
-    using WarpFragmentA = typename Operator::FragmentA;
-    using WarpFragmentB = typename Operator::FragmentB;
-    Dequantizer warp_dequantizer_;
-
-    using ElementA = typename IteratorA::Element;
-    using ElementB = typename IteratorB::Element;
-    using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
-
-    static constexpr bool RequiresTileInterleave
-        = layout::IsColumnMajorTileInterleave<typename LayoutDetailsForB::Layout>::value;
-    static_assert(!RequiresTileInterleave || (RequiresTileInterleave && (Shape::kK == LayoutDetailsForB::ThreadblockK)),
-        "Layout K must match threadblockK");
-
-protected:
-    /// Iterator to write threadblock-scoped tile of A operand to shared memory
-    SmemIteratorA smem_iterator_A_;
-
-    /// Iterator to write threadblock-scoped tile of B operand to shared memory
-    SmemIteratorB smem_iterator_B_;
-
-    /// Iterator to write threadblock-scoped tile of scale operand to shared memory
-    SmemIteratorScale smem_iterator_scale_;
-
-public:
-    /// Construct from tensor references
-    CUTLASS_DEVICE
-    DqMmaPipelined(typename Base::SharedStorage&
-                       shared_storage, ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        int const group_size, ///< Will not be used, just to adapt to finegrained modifications and make the compilation
-                              ///< successful. Because DqMmaPipelined is only enabled for sm<80, so even if this
-                              ///< argument is not added, it does not affect compilation for sm>=80.
-        int thread_idx,       ///< ID within the threadblock
-        int warp_idx,         ///< ID of warp
-        int lane_idx          ///< ID of each thread within a warp
-        )
-        : Base(shared_storage, thread_idx, warp_idx, lane_idx)
-        , warp_dequantizer_({shared_storage.operand_scale.data(), LayoutScale(Shape::kN)},
-              (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) / Base::WarpCount::kM, lane_idx)
-        , smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx)
-        , smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx)
-        , smem_iterator_scale_(LayoutScale(Shape::kN), shared_storage.operand_scale.data(), {1, Shape::kN}, thread_idx)
-    {
-
-        // Compute warp location within threadblock tile by mapping the warp_id to
-        // three coordinates:
-        //   _m: the warp's position within the threadblock along the M dimension
-        //   _n: the warp's position within the threadblock along the N dimension
-        //   _k: the warp's position within the threadblock along the K dimension
-
-        int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
-        int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
-
-        int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
-        int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
-
-        // Add per-warp offsets in units of warp-level tiles
-        this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
-        this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
-    }
-
-    /// Perform a threadblock-scoped matrix multiply-accumulate
-    CUTLASS_DEVICE
-    void operator()(int gemm_k_iterations, ///< number of iterations of the mainloop
-        FragmentC& accum,                  ///< destination accumulator tile
-        IteratorA iterator_A,              ///< iterator over A operand in global memory
-        IteratorB iterator_B,              ///< iterator over B operand in global memory
-        IteratorScale iterator_scale,      ///< iterator over scale operand in global memory
-        FragmentC const& src_accum)
-    {                                      ///< source accumulator tile
-
-        //
-        // Prologue
-        //
-        TransformBAfterLDG ldg_converter;
-        TransformBAfterLDS lds_converter;
-
-        using TransformA
-            = NumericArrayConverter<typename WarpFragmentA::Element, typename FragmentA::Element, FragmentA::kElements>;
-
-        using TransformScale = NumericArrayConverter<typename SmemIteratorScale::Fragment::Element,
-            typename FragmentScale::Element, FragmentScale::kElements>;
-
-        // These transforms are mainly to handle when we have bfloat activations and weights in GMEM and want
-        // to issue HMMA on architectures older than Ampere. We will convert to FP16 before STS.
-        TransformA transformA;
-        TransformScale transformScale;
-
-        // Perform accumulation in the 'd' output operand
-        accum = src_accum;
-
-        FragmentA tb_frag_A;
-        FragmentB tb_frag_B;
-        FragmentScale tb_frag_scales;
-
-        using WarpFragmentScale = typename Dequantizer::FragmentScale;
-        WarpFragmentScale warp_frag_scales;
-
-        tb_frag_A.clear();
-        tb_frag_B.clear();
-        tb_frag_scales.clear();
-
-        // The last kblock is loaded in the prolog
-        iterator_A.load(tb_frag_A);
-        iterator_B.load(tb_frag_B);
-        iterator_scale.load(tb_frag_scales);
-
-        ++iterator_A;
-        ++iterator_B;
-
-        this->smem_iterator_A_.store(transformA(tb_frag_A));
-        this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
-        this->smem_iterator_scale_.store(transformScale(tb_frag_scales));
-
-        ++this->smem_iterator_A_;
-        ++this->smem_iterator_B_;
-
-        __syncthreads();
-
-        warp_dequantizer_.load(warp_frag_scales);
-
-        // Pair of fragments used to overlap shared memory loads and math instructions
-        WarpFragmentA warp_frag_A[2];
-        WarpFragmentB warp_frag_B[2];
-
-        this->warp_tile_iterator_A_.set_kgroup_index(0);
-        this->warp_tile_iterator_B_.set_kgroup_index(0);
-
-        this->warp_tile_iterator_A_.load(warp_frag_A[0]);
-        this->warp_tile_iterator_B_.load(warp_frag_B[0]);
-
-        ++this->warp_tile_iterator_A_;
-        ++this->warp_tile_iterator_B_;
-
-        Operator warp_mma;
-
-        int smem_write_stage_idx = 1;
-
-        // Avoid reading out of bounds
-        iterator_A.clear_mask(gemm_k_iterations <= 1);
-        iterator_B.clear_mask(gemm_k_iterations <= 1);
-
-        // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing
-        // shared memory loads (which have the tighest latency requirement).
-
-        //
-        // Mainloop
-        //
-
-        // Note: The main loop does not support Base::kWarpGemmIterations == 2.
-        CUTLASS_GEMM_LOOP
-        for (; gemm_k_iterations > 0; --gemm_k_iterations)
-        {
-            //
-            // Loop over GEMM K dimension
-            //
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k)
-            {
-
-                // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group
-                // as the case may be.
-
-                if (warp_mma_k == Base::kWarpGemmIterations - 1)
-                {
-
-                    // Write fragments to shared memory
-                    this->smem_iterator_A_.store(transformA(tb_frag_A));
-
-                    this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
-
-                    __syncthreads();
-
-                    ++this->smem_iterator_A_;
-                    ++this->smem_iterator_B_;
-
-                    // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory
-                    if (smem_write_stage_idx == 1)
-                    {
-                        this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
-                        this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
-                    }
-                    else
-                    {
-                        this->warp_tile_iterator_A_.add_tile_offset(
-                            {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
-                        this->warp_tile_iterator_B_.add_tile_offset(
-                            {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0});
-                    }
-
-                    smem_write_stage_idx ^= 1;
-                }
-
-                this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
-                this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
-                ++this->warp_tile_iterator_A_;
-
-                int const warp_tileB_k_compute_offset = warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
-                int const warp_tileB_k_load_offset = warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
-                // We are just about to finish computing on a fragment of B, so initiate the load for the next fragment.
-                if (warp_tileB_k_compute_offset == Base::kNumKIterationsPerWarpBLoad - 1)
-                {
-                    this->warp_tile_iterator_B_.set_kgroup_index(
-                        (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
-                    this->warp_tile_iterator_B_.load(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
-                    ++this->warp_tile_iterator_B_;
-                }
-
-                if (warp_mma_k == 0)
-                {
-
-                    iterator_A.load(tb_frag_A);
-                    iterator_B.load(tb_frag_B);
-
-                    ++iterator_A;
-                    ++iterator_B;
-
-                    // Avoid reading out of bounds if this was the last loop iteration
-                    iterator_A.clear_mask(gemm_k_iterations <= 2);
-                    iterator_B.clear_mask(gemm_k_iterations <= 2);
-                }
-
-                typename TransformBAfterLDS::result_type converted_frag_B
-                    = lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
-                warp_dequantizer_.dequantize(converted_frag_B, warp_frag_scales);
-                run_warp_mma(
-                    warp_mma, accum, warp_frag_A[warp_mma_k % 2], converted_frag_B, accum, warp_tileB_k_compute_offset);
-            }
+          smem_write_stage_idx ^= 1;
         }
+
+        this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) %
+                                                     Base::kWarpGemmIterations);
+        this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
+        ++this->warp_tile_iterator_A_;
+
+        int const warp_tileB_k_compute_offset =
+            warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
+        int const warp_tileB_k_load_offset =
+            warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
+        // We are just about to finish computing on a fragment of B, so initiate
+        // the load for the next fragment.
+        if (warp_tileB_k_compute_offset ==
+            Base::kNumKIterationsPerWarpBLoad - 1) {
+          this->warp_tile_iterator_B_.set_kgroup_index(
+              (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
+          this->warp_tile_iterator_B_.load(
+              warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
+          ++this->warp_tile_iterator_B_;
+        }
+
+        if (warp_mma_k == 0) {
+          iterator_A.load(tb_frag_A);
+          iterator_B.load(tb_frag_B);
+
+          ++iterator_A;
+          ++iterator_B;
+
+          // Avoid reading out of bounds if this was the last loop iteration
+          iterator_A.clear_mask(gemm_k_iterations <= 2);
+          iterator_B.clear_mask(gemm_k_iterations <= 2);
+        }
+
+        typename TransformBAfterLDS::result_type converted_frag_B =
+            lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
+        warp_dequantizer_.dequantize(converted_frag_B, warp_frag_scales);
+        run_warp_mma(warp_mma,
+                     accum,
+                     warp_frag_A[warp_mma_k % 2],
+                     converted_frag_B,
+                     accum,
+                     warp_tileB_k_compute_offset);
+      }
     }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_mma_base.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_mma_base.h
index bc91e724bb..9b118add9a 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_mma_base.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_mma_base.h
@@ -66,7 +66,7 @@ template <
     /// Size of extra quantized params
     typename QuantParamsShape>
 class Wint2xMmaBase {
-public:
+ public:
   ///< Size of the Gemm problem - concept: gemm::GemmShape<>
   using Shape = Shape_;
 
@@ -85,9 +85,9 @@ public:
   using WarpGemm = typename Policy::Operator::Shape;
 
   /// Shape describing the number of warps filling the CTA
-  using WarpCount =
-      GemmShape<Shape::kM / WarpGemm::kM, Shape::kN / WarpGemm::kN,
-                Shape::kK / WarpGemm::kK>;
+  using WarpCount = GemmShape<Shape::kM / WarpGemm::kM,
+                              Shape::kN / WarpGemm::kN,
+                              Shape::kK / WarpGemm::kK>;
 
   /// Number of warp-level GEMM operations
   static int const kWarpGemmIterations =
@@ -95,7 +95,8 @@ public:
 
   /// Number of warp-level GEMM operations per load for B
   static constexpr int kWarpGemmIterationsPerLoadForB =
-      Operator::IteratorB::InstructionShape::kRow / Operator::InstructionShape::kK;
+      Operator::IteratorB::InstructionShape::kRow /
+      Operator::InstructionShape::kK;
   static_assert(!(kWarpGemmIterations % kWarpGemmIterationsPerLoadForB), "");
 
   static constexpr int kWarpLoadIterationsForB =
@@ -125,7 +126,7 @@ public:
 
   /// Shared storage object needed by threadblock-scoped GEMM
   class SharedStorage {
-  public:
+   public:
     //
     // Type definitions
     //
@@ -142,7 +143,7 @@ public:
     /// Shape of all quant params in shared memory
     using QuantParamsShapeB = QuantParamsShape;
 
-  public:
+   public:
     //
     // Data members
     //
@@ -156,7 +157,7 @@ public:
     /// Buffer for extra quant params of B operand
     AlignedBuffer<uint8_t, QuantParamsShapeB::kCount> operand_quant_params_B;
 
-  public:
+   public:
     //
     // Methods
     //
@@ -186,7 +187,7 @@ public:
     }
   };
 
-protected:
+ protected:
   //
   // Data members
   //
@@ -197,7 +198,7 @@ protected:
   /// Iterator to load a warp-scoped tile of B operand from shared memory
   typename Operator::IteratorB warp_tile_iterator_B_;
 
-public:
+ public:
   /// Construct from tensor references
   CUTLASS_DEVICE
   Wint2xMmaBase(
@@ -215,8 +216,8 @@ public:
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace threadblock
-} // namespace gemm
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_mma_multistage.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_mma_multistage.h
index dd26cf68ea..245a89c152 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_mma_multistage.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_mma_multistage.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -91,11 +92,17 @@ template <
     typename QuantParamsAccessor_,
     /// Use zfill or predicate for out-of-bound cp.async
     SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone>
-class Wint2xMmaMultistage :
-  public Wint2xMmaBase<Shape_, Policy_, Stages, typename QuantParamsAccessor_::QuantParamsShape> {
-public:
+class Wint2xMmaMultistage
+    : public Wint2xMmaBase<Shape_,
+                           Policy_,
+                           Stages,
+                           typename QuantParamsAccessor_::QuantParamsShape> {
+ public:
   ///< Base class
-  using Base = Wint2xMmaBase<Shape_, Policy_, Stages, typename QuantParamsAccessor_::QuantParamsShape>;
+  using Base = Wint2xMmaBase<Shape_,
+                             Policy_,
+                             Stages,
+                             typename QuantParamsAccessor_::QuantParamsShape>;
   ///< Size of the Gemm problem - concept: gemm::GemmShape<>
   using Shape = Shape_;
   ///< Iterates over tiles of A operand in global memory
@@ -133,17 +140,19 @@ public:
   /// Minimum architecture is Sm80 to support cp.async
   using ArchTag = arch::Sm80;
 
-  //using LayoutScale = typename QuantParamsAccessor::IteratorSuperScale::Layout;
+  // using LayoutScale = typename
+  // QuantParamsAccessor::IteratorSuperScale::Layout;
   using LayoutScale = layout::RowMajor;
   using WarpTransformedFragmentB = typename Operator::TransformedFragmentB;
-  using WarpDequantizer =
-      warp::MmaTensorOpWin2xDequantizer<Operator,
-                                        typename Base::WarpGemm,
-                                        Operand::kB,
-                                        typename WarpTransformedFragmentB::Element,
-                                        LayoutScale,
-                                        QuantParamsAccessor::kGroupSize>;
-  static_assert(sizeof(WarpDequantizer) > 0, "WarpDequantizer template instantiation failed");
+  using WarpDequantizer = warp::MmaTensorOpWin2xDequantizer<
+      Operator,
+      typename Base::WarpGemm,
+      Operand::kB,
+      typename WarpTransformedFragmentB::Element,
+      LayoutScale,
+      QuantParamsAccessor::kGroupSize>;
+  static_assert(sizeof(WarpDequantizer) > 0,
+                "WarpDequantizer template instantiation failed");
 
   /// Complex transform on A operand
   static ComplexTransform const kTransformA = Operator::kTransformA;
@@ -153,7 +162,6 @@ public:
 
   /// Internal structure exposed for introspection.
   struct Detail {
-
     /// Number of cp.async instructions to load one stage of operand A
     static int const AsyncCopyIterationsPerStageA =
         IteratorA::ThreadMap::Iterations::kCount;
@@ -167,24 +175,25 @@ public:
 
     /// Number of cp.async instructions to load on group of operand A
     static int const kAccessesPerGroupA =
-        (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
+        (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
 
     /// Number of cp.async instructions to load on group of operand B
     static int const kAccessesPerGroupB =
-        (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
+        (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
 
-    // Optional staged-accumulation (e.g., tf32x3 kernels) for improved numerical
-    // accuracy, where each mainloop iteration first accumulates into a temporary
-    // set of freshly-cleared accumulators, which are subsequently added to the
-    // final accumulator set.
-    static bool const kStagedAccumulation = arch::detail::UseStagedAccumulation<Operator>::value;
+    // Optional staged-accumulation (e.g., tf32x3 kernels) for improved
+    // numerical accuracy, where each mainloop iteration first accumulates into
+    // a temporary set of freshly-cleared accumulators, which are subsequently
+    // added to the final accumulator set.
+    static bool const kStagedAccumulation =
+        arch::detail::UseStagedAccumulation<Operator>::value;
   };
 
  private:
-
   // Structure encapsulating pipeline state live from one iteration to the next
   struct PipeState {
-
     using WarpLoadedFragmentA = typename Operator::FragmentA;
     using WarpLoadedFragmentB = typename Operator::FragmentB;
     using WarpTransformedFragmentA = typename Operator::TransformedFragmentA;
@@ -197,10 +206,12 @@ public:
     /// Temporary accumulator to facilitate staged-accumulation
     FragmentC tmp_accum_;
 
-    /// Pair of A fragments used to overlap shared memory loads and math instructions
+    /// Pair of A fragments used to overlap shared memory loads and math
+    /// instructions
     WarpTransformedFragmentA warp_frag_A_[2];
 
-    /// Pair of B fragments used to overlap shared memory loads and math instructions
+    /// Pair of B fragments used to overlap shared memory loads and math
+    /// instructions
     WarpLoadedFragmentB warp_loaded_frag_B_;
     WarpTransformedFragmentB warp_frag_B_[2];
 
@@ -218,12 +229,14 @@ public:
   using LayoutDetailsForB = kernel::LayoutDetailsB<ElementA, ElementB, ArchTag>;
 
   static constexpr bool IsTileInterleaveLayout =
-      layout::IsColumnMajorTileInterleave<typename LayoutDetailsForB::Layout>::value;
-  static_assert(!IsTileInterleaveLayout || (IsTileInterleaveLayout && (Shape::kK == LayoutDetailsForB::ThreadblockK)),
-      "Layout K must match threadblockK");
+      layout::IsColumnMajorTileInterleave<
+          typename LayoutDetailsForB::Layout>::value;
+  static_assert(!IsTileInterleaveLayout ||
+                    (IsTileInterleaveLayout &&
+                     (Shape::kK == LayoutDetailsForB::ThreadblockK)),
+                "Layout K must match threadblockK");
 
  private:
-
   //
   // Data members
   //
@@ -249,8 +262,7 @@ public:
   /// Shared memory read stage index
   int smem_read_stage_idx_;
 
-public:
-
+ public:
   /// Construct from tensor references
   CUTLASS_DEVICE
   Wint2xMmaMultistage(
@@ -261,19 +273,24 @@ public:
       ///< ID of warp
       int warp_idx,
       ///< ID of each thread within a warp
-      int lane_idx
-  ) : Base(shared_storage, thread_idx, warp_idx, lane_idx),
-      smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
-      smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx),
-      quant_params_accessor_B_(shared_storage.operand_quant_params_B.data(), thread_idx, warp_idx, lane_idx),
-      warp_dequantizer_(quant_params_accessor_B_.super_scale_ref(),
-                        quant_params_accessor_B_.local_scale_ref(),
-                        quant_params_accessor_B_.code_scale_ref(),
-                        quant_params_accessor_B_.code_zp_ref(),
-                        (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) / Base::WarpCount::kM, lane_idx),
-      smem_write_stage_idx_(0),
-      smem_read_stage_idx_(0)
-  {
+      int lane_idx)
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
+        smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
+        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx),
+        quant_params_accessor_B_(shared_storage.operand_quant_params_B.data(),
+                                 thread_idx,
+                                 warp_idx,
+                                 lane_idx),
+        warp_dequantizer_(
+            quant_params_accessor_B_.super_scale_ref(),
+            quant_params_accessor_B_.local_scale_ref(),
+            quant_params_accessor_B_.code_scale_ref(),
+            quant_params_accessor_B_.code_zp_ref(),
+            (warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN)) /
+                Base::WarpCount::kM,
+            lane_idx),
+        smem_write_stage_idx_(0),
+        smem_read_stage_idx_(0) {
     // Compute warp location within threadblock tile by mapping the warp_id to
     // three coordinates:
     //   _m: the warp's position within the threadblock along the M dimension
@@ -295,22 +312,26 @@ public:
 
   /// Advance shared memory read-iterators to the next stage
   CUTLASS_DEVICE
-  void advance_smem_read_stage()
-  {
+  void advance_smem_read_stage() {
     ++smem_read_stage_idx_;
 
     if (smem_read_stage_idx_ == Base::kStages) {
       // Wrap back around to the 'start' of the circular buffer in shared memory
-      this->warp_tile_iterator_A_.add_tile_offset({0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
-      this->warp_tile_iterator_B_.add_tile_offset({-Base::kStages * Policy::kPartitionsK * Base::kWarpLoadIterationsForB, 0});
+      this->warp_tile_iterator_A_.add_tile_offset(
+          {0,
+           -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
+      this->warp_tile_iterator_B_.add_tile_offset(
+          {-Base::kStages * Policy::kPartitionsK *
+               Base::kWarpLoadIterationsForB,
+           0});
       smem_read_stage_idx_ = 0;
     }
   }
 
-  /// Advance global memory read-iterators and shared memory write-iterators to the stage
+  /// Advance global memory read-iterators and shared memory write-iterators to
+  /// the stage
   CUTLASS_DEVICE
-  void advance_smem_write_stage(IteratorA &iterator_A, IteratorB &iterator_B)
-  {
+  void advance_smem_write_stage(IteratorA &iterator_A, IteratorB &iterator_B) {
     // Advance global iterators
     iterator_A.add_tile_offset({0, 1});
     iterator_B.add_tile_offset({1, 0});
@@ -395,7 +416,9 @@ public:
         CUTLASS_PRAGMA_UNROLL
         for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
           auto gmem_ptr = iterator_B.get();
-          bool is_valid = (threadIdx.x < IteratorB::ThreadMap::kThreads) ? iterator_B.valid() : false;
+          bool is_valid = (threadIdx.x < IteratorB::ThreadMap::kThreads)
+                              ? iterator_B.valid()
+                              : false;
 
           if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
             cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
@@ -429,10 +452,9 @@ public:
       for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
         auto gmem_ptr = iterator_A.get();
 
-        int const kSrcBytes =
-            sizeof_bits<typename IteratorA::Element>::value *
-            IteratorA::ThreadMap::kElementsPerAccess /
-            IteratorA::kAccessesPerVector / 8;
+        int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
+                              IteratorA::ThreadMap::kElementsPerAccess /
+                              IteratorA::kAccessesPerVector / 8;
 
         cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
             dst_ptr + v, iterator_A.get(), iterator_A.valid());
@@ -464,10 +486,9 @@ public:
       for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
         auto gmem_ptr = iterator_B.get();
 
-        int const kSrcBytes =
-            sizeof_bits<typename IteratorB::Element>::value *
-            IteratorB::ThreadMap::kElementsPerAccess /
-            IteratorB::kAccessesPerVector / 8;
+        int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value *
+                              IteratorB::ThreadMap::kElementsPerAccess /
+                              IteratorB::kAccessesPerVector / 8;
 
         cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
             dst_ptr + v, iterator_B.get(), iterator_B.valid());
@@ -480,18 +501,22 @@ public:
   }
 
   /// GEMM prologue.  Bootstrap the global->shared memory pipeline by fetching
-  /// the global fragments needed by the first kStages-1 threadblock mainloop iterations
+  /// the global fragments needed by the first kStages-1 threadblock mainloop
+  /// iterations
   CUTLASS_DEVICE
-  void prologue(
-    IteratorA &iterator_A,      ///< [in|out] iterator over A operand in global memory
-    IteratorB &iterator_B,      ///< [in|out] iterator over B operand in global memory
-    QuantArguments &mma_quant_args, ///< iterators for extra quant params for B
-    int &gemm_k_iterations)     ///< [in|out] number of threadblock mainloop iterations remaining
+  void prologue(IteratorA &iterator_A,  ///< [in|out] iterator over A operand in
+                                        ///< global memory
+                IteratorB &iterator_B,  ///< [in|out] iterator over B operand in
+                                        ///< global memory
+                QuantArguments &
+                    mma_quant_args,  ///< iterators for extra quant params for B
+                int &gemm_k_iterations)  ///< [in|out] number of threadblock
+                                         ///< mainloop iterations remaining
   {
     // Issue several complete stages
     CUTLASS_PRAGMA_UNROLL
-    for (int stage = 0; stage < Base::kStages - 1; ++stage, --gemm_k_iterations) {
-
+    for (int stage = 0; stage < Base::kStages - 1;
+         ++stage, --gemm_k_iterations) {
       // Disable global fetching if done with global fetch iterations
       iterator_A.clear_mask(gemm_k_iterations == 0);
       iterator_B.clear_mask(gemm_k_iterations == 0);
@@ -502,11 +527,14 @@ public:
       // Async copy zipped B to shared memory.
       copy_tiles_and_advance_per_stage_B(iterator_B);
 
-      // Async copy other quantized params to shared memory, local_scale, code_scale, code_zp, super_scale.
+      // Async copy other quantized params to shared memory, local_scale,
+      // code_scale, code_zp, super_scale.
       if (stage == 0) {
-        quant_params_accessor_B_.copy_tiles_and_advance_per_stage<true>(mma_quant_args, stage);
+        quant_params_accessor_B_.copy_tiles_and_advance_per_stage<true>(
+            mma_quant_args, stage);
       } else {
-        quant_params_accessor_B_.copy_tiles_and_advance_per_stage<false>(mma_quant_args, stage);
+        quant_params_accessor_B_.copy_tiles_and_advance_per_stage<false>(
+            mma_quant_args, stage);
       }
 
       // Move to the next write stage
@@ -517,11 +545,12 @@ public:
       cutlass::arch::cp_async_fence();
     }
 
-    // Optionally clear the remaining stages of SMEM. This is a functional requirement for
-    // some kernels so that all accumulator elements outside the GEMM footprint are zero.
+    // Optionally clear the remaining stages of SMEM. This is a functional
+    // requirement for some kernels so that all accumulator elements outside the
+    // GEMM footprint are zero.
     if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) {
-
-      /// Iterator to write threadblock-scoped tile of A operand to shared memory
+      /// Iterator to write threadblock-scoped tile of A operand to shared
+      /// memory
       SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
       typename IteratorA::AccessType zero_A;
 
@@ -531,7 +560,6 @@ public:
       // Async Copy for operand A
       CUTLASS_PRAGMA_UNROLL
       for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
-
         typename IteratorA::AccessType *dst_ptr =
             reinterpret_cast<typename IteratorA::AccessType *>(
                 last_smem_iterator_A.get());
@@ -545,7 +573,8 @@ public:
         return;
       }
 
-      /// Iterator to write threadblock-scoped tile of B operand to shared memory
+      /// Iterator to write threadblock-scoped tile of B operand to shared
+      /// memory
       SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_);
       typename IteratorB::AccessType zero_B;
 
@@ -555,7 +584,6 @@ public:
       // Async Copy for operand B
       CUTLASS_PRAGMA_UNROLL
       for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
-
         typename IteratorB::AccessType *dst_ptr =
             reinterpret_cast<typename IteratorB::AccessType *>(
                 last_smem_iterator_B.get());
@@ -569,9 +597,9 @@ public:
 
   /// Wait until we have at least one completed global fetch stage
   CUTLASS_DEVICE
-  void gmem_wait()
-  {
-    // Wait until we have at least one committed global fetch stage. (#uncommitted = Base::kStages - 1 - #committed)
+  void gmem_wait() {
+    // Wait until we have at least one committed global fetch stage.
+    // (#uncommitted = Base::kStages - 1 - #committed)
     cutlass::arch::cp_async_wait<Base::kStages - 2>();
     __syncthreads();
   }
@@ -579,25 +607,31 @@ public:
   /// Perform a threadblock mainloop iteration of matrix multiply-accumulate
   CUTLASS_DEVICE
   void mac_loop_iter(
-    PipeState &pipe_state,          ///< [in|out] loop-carried pipeline state
-    FragmentC &accum,               ///< [in|out] destination accumulator tile
-    IteratorA &iterator_A,          ///< [in|out] iterator over A operand in global memory
-    IteratorB &iterator_B,          ///< [in|out] iterator over B operand in global memory
-    QuantArguments &mma_quant_args, ///< iterators for extra quant params for B
-    int &gemm_k_iterations,         ///< [in|out] number of threadblock mainloop iterations remaining
-    int stage)
-  {
+      PipeState &pipe_state,  ///< [in|out] loop-carried pipeline state
+      FragmentC &accum,       ///< [in|out] destination accumulator tile
+      IteratorA
+          &iterator_A,  ///< [in|out] iterator over A operand in global memory
+      IteratorB
+          &iterator_B,  ///< [in|out] iterator over B operand in global memory
+      QuantArguments
+          &mma_quant_args,     ///< iterators for extra quant params for B
+      int &gemm_k_iterations,  ///< [in|out] number of threadblock mainloop
+                               ///< iterations remaining
+      int stage) {
     const int mma_stage = stage - Base::kStages + 1;
 
     // Unroll the warp-level MMA tiles of a threadblock's mainloop iteration
     CUTLASS_PRAGMA_UNROLL
-    for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) {
-
-      int warp_k_compute_offset_B = warp_mma_k % Base::kWarpGemmIterationsPerLoadForB;
+    for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
+         ++warp_mma_k) {
+      int warp_k_compute_offset_B =
+          warp_mma_k % Base::kWarpGemmIterationsPerLoadForB;
 
       if (warp_k_compute_offset_B == Base::kWarpGemmIterationsPerLoadForB - 1) {
         // Load the next warp-tile's B fragment from shared memory
-        this->warp_tile_iterator_B_.set_kgroup_index(((warp_mma_k + 1) % Base::kWarpGemmIterations) / Base::kWarpLoadIterationsForB);
+        this->warp_tile_iterator_B_.set_kgroup_index(
+            ((warp_mma_k + 1) % Base::kWarpGemmIterations) /
+            Base::kWarpLoadIterationsForB);
         this->warp_tile_iterator_B_.load(pipe_state.warp_loaded_frag_B_);
         ++this->warp_tile_iterator_B_;
       }
@@ -608,28 +642,31 @@ public:
       }
 
       // Load the next warp-tile's A fragment from shared memory
-      this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
-      this->warp_tile_iterator_A_.load(pipe_state.warp_frag_A_[(warp_mma_k + 1) % 2]);
+      this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) %
+                                                   Base::kWarpGemmIterations);
+      this->warp_tile_iterator_A_.load(
+          pipe_state.warp_frag_A_[(warp_mma_k + 1) % 2]);
       ++this->warp_tile_iterator_A_;
 
       // dequantizes next warp-tile
-      warp_dequantizer_.dequantize(pipe_state.warp_frag_local_scale_,
-                                   pipe_state.warp_frag_code_scale_,
-                                   pipe_state.warp_frag_code_zp_,
-                                   pipe_state.warp_frag_super_scale_,
-                                   pipe_state.warp_loaded_frag_B_,
-                                   pipe_state.warp_frag_B_[(warp_mma_k + 1) % 2],
-                                   ((warp_mma_k == Base::kWarpGemmIterations - 1) ? (mma_stage + 1) : mma_stage) * Shape::kK,
-                                   (warp_mma_k + 1) % Base::kWarpGemmIterationsPerLoadForB);
+      warp_dequantizer_.dequantize(
+          pipe_state.warp_frag_local_scale_,
+          pipe_state.warp_frag_code_scale_,
+          pipe_state.warp_frag_code_zp_,
+          pipe_state.warp_frag_super_scale_,
+          pipe_state.warp_loaded_frag_B_,
+          pipe_state.warp_frag_B_[(warp_mma_k + 1) % 2],
+          ((warp_mma_k == Base::kWarpGemmIterations - 1) ? (mma_stage + 1)
+                                                         : mma_stage) *
+              Shape::kK,
+          (warp_mma_k + 1) % Base::kWarpGemmIterationsPerLoadForB);
 
       // Execute the current warp-tile of MMA operations
       if constexpr (Detail::kStagedAccumulation) {
-        warp_mma_(
-          pipe_state.tmp_accum_,
-          pipe_state.warp_frag_A_[warp_mma_k % 2],
-          pipe_state.warp_frag_B_[warp_mma_k % 2],
-          pipe_state.tmp_accum_
-        );
+        warp_mma_(pipe_state.tmp_accum_,
+                  pipe_state.warp_frag_A_[warp_mma_k % 2],
+                  pipe_state.warp_frag_B_[warp_mma_k % 2],
+                  pipe_state.tmp_accum_);
 
         if (warp_mma_k == 0) {
           plus<FragmentC> plus_accum;
@@ -637,11 +674,10 @@ public:
           pipe_state.tmp_accum_.clear();
         }
       } else {
-        warp_mma_(
-          accum,
-          pipe_state.warp_frag_A_[warp_mma_k % 2],
-          pipe_state.warp_frag_B_[warp_mma_k % 2],
-          accum);
+        warp_mma_(accum,
+                  pipe_state.warp_frag_A_[warp_mma_k % 2],
+                  pipe_state.warp_frag_B_[warp_mma_k % 2],
+                  accum);
       }
 
       // Except for the last warp-tile, all warp-tiles issue their share of
@@ -654,22 +690,28 @@ public:
         copy_tiles_and_advance_B(iterator_B, group_start_iteration_B);
 
         if (warp_mma_k == 0) {
-          quant_params_accessor_B_.copy_tiles_and_advance_per_stage<false>(mma_quant_args, stage);
+          quant_params_accessor_B_.copy_tiles_and_advance_per_stage<false>(
+              mma_quant_args, stage);
         }
       }
 
       // The second-to-last warp-tile also:
-      //   - performs the last warp-tile's share of global->shared fragment copies
+      //   - performs the last warp-tile's share of global->shared fragment
+      //   copies
       //   - moves to the next global fetch stage
       if (warp_mma_k + 2 == Base::kWarpGemmIterations) {
         // Performs the last warp-tile's share of global->shared fragment copies
-        if constexpr (Detail::AsyncCopyIterationsPerStageA >= Base::kWarpGemmIterations) {
-          int group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
+        if constexpr (Detail::AsyncCopyIterationsPerStageA >=
+                      Base::kWarpGemmIterations) {
+          int group_start_iteration_A =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
           copy_tiles_and_advance_A(iterator_A, group_start_iteration_A);
         }
 
-        if constexpr (Detail::AsyncCopyIterationsPerStageB >= Base::kWarpGemmIterations) {
-          int group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
+        if constexpr (Detail::AsyncCopyIterationsPerStageB >=
+                      Base::kWarpGemmIterations) {
+          int group_start_iteration_B =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
           copy_tiles_and_advance_B(iterator_B, group_start_iteration_B);
         }
 
@@ -691,7 +733,8 @@ public:
         --gemm_k_iterations;
         iterator_A.clear_mask(gemm_k_iterations == 0);
         iterator_B.clear_mask(gemm_k_iterations == 0);
-        quant_params_accessor_B_.clear_mask(mma_quant_args, gemm_k_iterations == 0);
+        quant_params_accessor_B_.clear_mask(mma_quant_args,
+                                            gemm_k_iterations == 0);
       }
     }
   }
@@ -700,12 +743,13 @@ public:
   /// multiply-accumulate.  Assumes prologue has been initiated.
   CUTLASS_DEVICE
   void gemm_iters(
-      int gemm_k_iterations,        ///< number of threadblock mainloop iterations
-      FragmentC &accum,             ///< [in|out] accumulator tile
-      IteratorA &iterator_A,        ///< [in|out] iterator over A operand in global memory
-      IteratorB &iterator_B,        ///< [in|out] iterator over B operand in global memory
-      QuantArguments &mma_quant_args)
-  {
+      int gemm_k_iterations,  ///< number of threadblock mainloop iterations
+      FragmentC &accum,       ///< [in|out] accumulator tile
+      IteratorA
+          &iterator_A,  ///< [in|out] iterator over A operand in global memory
+      IteratorB
+          &iterator_B,  ///< [in|out] iterator over B operand in global memory
+      QuantArguments &mma_quant_args) {
     PipeState pipe_state;
 
     // Disable global fetching if done with global fetch iterations
@@ -748,14 +792,13 @@ public:
     // Mainloop
     CUTLASS_GEMM_LOOP
     for (; gemm_k_iterations > (-Base::kStages + 1);) {
-      mac_loop_iter(
-        pipe_state,
-        accum,
-        iterator_A,
-        iterator_B,
-        mma_quant_args,
-        gemm_k_iterations,
-        stage);
+      mac_loop_iter(pipe_state,
+                    accum,
+                    iterator_A,
+                    iterator_B,
+                    mma_quant_args,
+                    gemm_k_iterations,
+                    stage);
       stage += 1;
     }
 
@@ -764,7 +807,8 @@ public:
       accum = plus_accum(accum, pipe_state.tmp_accum_);
     }
 
-    // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
+    // Commit and drain all pending and predicated cp.async pnz from the GEMM
+    // mainloop
     cutlass::arch::cp_async_fence();
     cutlass::arch::cp_async_wait<0>();
     __syncthreads();
@@ -772,15 +816,16 @@ public:
 
   /// Prepares the class for another prologue.
   CUTLASS_DEVICE
-  void wind_down()
-  {
-    // Catch-up the smem-read iterator to the smem-write iterator (so this class can be reused for another tile's prologue)
+  void wind_down() {
+// Catch-up the smem-read iterator to the smem-write iterator (so this class can
+// be reused for another tile's prologue)
 
-    // First, increment remaining warp tiles to get to the next full stage.  (Ideally we would
-    // just decrement one tile, but not all iterators implement --() decrement.)
-    #pragma unroll
-    for (int warp_mma_k = 1; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k)
-    {
+// First, increment remaining warp tiles to get to the next full stage. (Ideally
+// we would just decrement one tile, but not all iterators implement --()
+// decrement.)
+#pragma unroll
+    for (int warp_mma_k = 1; warp_mma_k < Base::kWarpGemmIterations;
+         ++warp_mma_k) {
       this->warp_tile_iterator_A_.set_kgroup_index(warp_mma_k);
       this->warp_tile_iterator_B_.set_kgroup_index(warp_mma_k);
 
@@ -789,22 +834,24 @@ public:
     }
     smem_read_stage_idx_++;
 
-    // Then wrap back two full stages (one for the tile advancing we just did, and one to catch the write iterators)
-    static const int kStageIters = Policy::kPartitionsK * Base::kWarpGemmIterations;
-    if (smem_read_stage_idx_ > 1)
-    {
+    // Then wrap back two full stages (one for the tile advancing we just did,
+    // and one to catch the write iterators)
+    static const int kStageIters =
+        Policy::kPartitionsK * Base::kWarpGemmIterations;
+    if (smem_read_stage_idx_ > 1) {
       this->warp_tile_iterator_A_.add_tile_offset({0, (-2 * kStageIters)});
       this->warp_tile_iterator_B_.add_tile_offset({(-2 * kStageIters), 0});
-    }
-    else
-    {
-      this->warp_tile_iterator_A_.add_tile_offset({0, ((Base::kStages - 2) * kStageIters)});
-      this->warp_tile_iterator_B_.add_tile_offset({((Base::kStages - 2) * kStageIters), 0});
+    } else {
+      this->warp_tile_iterator_A_.add_tile_offset(
+          {0, ((Base::kStages - 2) * kStageIters)});
+      this->warp_tile_iterator_B_.add_tile_offset(
+          {((Base::kStages - 2) * kStageIters), 0});
     }
     smem_read_stage_idx_ = smem_write_stage_idx_;
   }
 
-  /// Perform a threadblock-scoped matrix multiply-accumulate, pre-load B to shared memory.
+  /// Perform a threadblock-scoped matrix multiply-accumulate, pre-load B to
+  /// shared memory.
   CUTLASS_DEVICE
   void operator()(
       ///< problem size of GEMM
@@ -819,8 +866,8 @@ public:
       QuantArguments mma_quant_args,
       ///< initial value of accumulator
       FragmentC const &src_accum) {
-
-    // Prologue (start fetching iterations of global fragments into shared memory)
+    // Prologue (start fetching iterations of global fragments into shared
+    // memory)
     prologue(iterator_A, iterator_B, mma_quant_args, gemm_k_iterations);
 
     // Wait until we have at least one completed global fetch stage
@@ -830,7 +877,8 @@ public:
     accum = src_accum;
 
     // Perform the MAC-iterations
-    gemm_iters(gemm_k_iterations, accum, iterator_A, iterator_B, mma_quant_args);
+    gemm_iters(
+        gemm_k_iterations, accum, iterator_A, iterator_B, mma_quant_args);
   }
 };
 
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_params_accessor.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_params_accessor.h
index c6eb2750c8..2409087cbc 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_params_accessor.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_params_accessor.h
@@ -46,9 +46,10 @@ template <
     /// Group size for quantization
     int GroupSize_>
 class Wint2ParamsAccessor {
-public:
-  static_assert(platform::is_same<T, half_t>::value || platform::is_same<T, bfloat16_t>::value,
-        "T must be fp16 or bf16");
+ public:
+  static_assert(platform::is_same<T, half_t>::value ||
+                    platform::is_same<T, bfloat16_t>::value,
+                "T must be fp16 or bf16");
 
   using ElementType = T;
   using Shape = Shape_;
@@ -72,7 +73,7 @@ public:
   using ElementLocalScale = typename IteratorLocalScale::Element;
   using LayoutLocalScale = typename IteratorLocalScale::Layout;
   static_assert(platform::is_same<ElementLocalScale, uint4b_t>::value,
-        "local_scale's type must be uint4b_t.");
+                "local_scale's type must be uint4b_t.");
 
   using ElementCodeScaleZp = typename IteratorCodeScaleZp::Element;
   using LayoutCodeScaleZp = typename IteratorCodeScaleZp::Layout;
@@ -80,24 +81,32 @@ public:
   /// 2 uint4b_t values are stored in a single uint8_t
   constexpr static int kStagesPerLocalScaleLoad = 2 * kGroupSize / Shape::kK;
   constexpr static int kLocalScaleRows =
-      IteratorLocalScale::Shape::kRow * IteratorLocalScale::Shape::kColumn * sizeof_bits<ElementLocalScale>::value / 8 / Shape::kN;
+      IteratorLocalScale::Shape::kRow * IteratorLocalScale::Shape::kColumn *
+      sizeof_bits<ElementLocalScale>::value / 8 / Shape::kN;
 
   using SmemElement = uint8_t;
-  constexpr static int kSmemRows =
-      kLocalScaleRows * kStages + sizeof(ElementSuperScale) + sizeof(ElementCodeScaleZp) * 2;
+  constexpr static int kSmemRows = kLocalScaleRows * kStages +
+                                   sizeof(ElementSuperScale) +
+                                   sizeof(ElementCodeScaleZp) * 2;
   constexpr static int kSmemColumns = Shape::kN;
 
   using QuantParamsShape = MatrixShape<kSmemRows, kSmemColumns>;
 
   constexpr static int kSuperScaleSmemOffset = 0;
-  constexpr static int kCodeScaleSmemOffset = kSmemColumns * sizeof(ElementSuperScale);
-  constexpr static int kCodeZpSmemOffset = kCodeScaleSmemOffset + kSmemColumns * sizeof(ElementCodeScaleZp);
-  constexpr static int kLocalScaleSmemOffset = kCodeZpSmemOffset + kSmemColumns * sizeof(ElementCodeScaleZp);
+  constexpr static int kCodeScaleSmemOffset =
+      kSmemColumns * sizeof(ElementSuperScale);
+  constexpr static int kCodeZpSmemOffset =
+      kCodeScaleSmemOffset + kSmemColumns * sizeof(ElementCodeScaleZp);
+  constexpr static int kLocalScaleSmemOffset =
+      kCodeZpSmemOffset + kSmemColumns * sizeof(ElementCodeScaleZp);
 
   /// TensorRef type for loading element from a tensor
-  using SuperTensorRef = cutlass::TensorRef<ElementSuperScale, LayoutSuperScale>;
-  using LocalTensorRef = cutlass::TensorRef<ElementLocalScale, LayoutLocalScale>;
-  using CodeTensorRef = cutlass::TensorRef<ElementCodeScaleZp, LayoutCodeScaleZp>;
+  using SuperTensorRef =
+      cutlass::TensorRef<ElementSuperScale, LayoutSuperScale>;
+  using LocalTensorRef =
+      cutlass::TensorRef<ElementLocalScale, LayoutLocalScale>;
+  using CodeTensorRef =
+      cutlass::TensorRef<ElementCodeScaleZp, LayoutCodeScaleZp>;
 
   struct Arguments {
     IteratorSuperScale iterator_super_scale;
@@ -113,14 +122,14 @@ public:
               IteratorCodeScaleZp iterator_code_scale,
               IteratorCodeScaleZp iterator_code_zp,
               int local_scale_pointer_offset)
-      : iterator_super_scale(iterator_super_scale),
-        iterator_local_scale(iterator_local_scale),
-        iterator_code_scale(iterator_code_scale),
-        iterator_code_zp(iterator_code_zp),
-        local_scale_pointer_offset(local_scale_pointer_offset) {}
+        : iterator_super_scale(iterator_super_scale),
+          iterator_local_scale(iterator_local_scale),
+          iterator_code_scale(iterator_code_scale),
+          iterator_code_zp(iterator_code_zp),
+          local_scale_pointer_offset(local_scale_pointer_offset) {}
   };
 
-private:
+ private:
   //
   // Data members
   //
@@ -128,13 +137,17 @@ private:
   /// Begin address of shared memory
   uint8_t* smem_pointer_;
 
-  /// Iterator to write threadblock-scoped tile of super scale operand to shared memory
+  /// Iterator to write threadblock-scoped tile of super scale operand to shared
+  /// memory
   SmemIteratorSuperScale smem_iterator_super_scale_;
-  /// Iterator to write threadblock-scoped tile of local scale operand to shared memory
+  /// Iterator to write threadblock-scoped tile of local scale operand to shared
+  /// memory
   SmemIteratorLocalScale smem_iterator_local_scale_;
-  /// Iterator to write threadblock-scoped tile of code scale operand to shared memory
+  /// Iterator to write threadblock-scoped tile of code scale operand to shared
+  /// memory
   SmemIteratorCodeScaleZp smem_iterator_code_scale_;
-  /// Iterator to write threadblock-scoped tile of code zp operand to shared memory
+  /// Iterator to write threadblock-scoped tile of code zp operand to shared
+  /// memory
   SmemIteratorCodeScaleZp smem_iterator_code_zp_;
 
   /// Shared memory write stage index
@@ -145,25 +158,29 @@ private:
 
   CUTLASS_DEVICE
   ElementSuperScale* get_super_scale_smem_ptr() {
-    return reinterpret_cast<ElementSuperScale*>(smem_pointer_ + kSuperScaleSmemOffset);
+    return reinterpret_cast<ElementSuperScale*>(smem_pointer_ +
+                                                kSuperScaleSmemOffset);
   }
 
   CUTLASS_DEVICE
   ElementLocalScale* get_local_scale_smem_ptr() {
-    return reinterpret_cast<ElementLocalScale*>(smem_pointer_ + kLocalScaleSmemOffset);
+    return reinterpret_cast<ElementLocalScale*>(smem_pointer_ +
+                                                kLocalScaleSmemOffset);
   }
 
   CUTLASS_DEVICE
   ElementCodeScaleZp* get_code_scale_smem_ptr() {
-    return reinterpret_cast<ElementCodeScaleZp*>(smem_pointer_ + kCodeScaleSmemOffset);
+    return reinterpret_cast<ElementCodeScaleZp*>(smem_pointer_ +
+                                                 kCodeScaleSmemOffset);
   }
 
   CUTLASS_DEVICE
   ElementCodeScaleZp* get_code_zp_smem_ptr() {
-    return reinterpret_cast<ElementCodeScaleZp*>(smem_pointer_ + kCodeZpSmemOffset);
+    return reinterpret_cast<ElementCodeScaleZp*>(smem_pointer_ +
+                                                 kCodeZpSmemOffset);
   }
 
-public:
+ public:
   /// Construct from tensor references
   CUTLASS_DEVICE
   Wint2ParamsAccessor(
@@ -175,55 +192,74 @@ public:
       int warp_idx,
       ///< ID of each thread within a warp
       int lane_idx)
-    : smem_pointer_(smem_pointer),
-      smem_iterator_super_scale_(LayoutSuperScale(IteratorSuperScale::Shape::kColumn),
-          get_super_scale_smem_ptr(), {1, IteratorSuperScale::Shape::kColumn}, thread_idx),
-      smem_iterator_local_scale_(LayoutLocalScale(IteratorLocalScale::Shape::kColumn),
-          get_local_scale_smem_ptr(), {1, IteratorLocalScale::Shape::kColumn}, thread_idx),
-      smem_iterator_code_scale_(LayoutCodeScaleZp(IteratorCodeScaleZp::Shape::kColumn),
-          get_code_scale_smem_ptr(), {1, IteratorCodeScaleZp::Shape::kColumn}, thread_idx),
-      smem_iterator_code_zp_(LayoutCodeScaleZp(IteratorCodeScaleZp::Shape::kColumn),
-          get_code_zp_smem_ptr(), {1, IteratorCodeScaleZp::Shape::kColumn}, thread_idx),
-      smem_write_stage_idx_(0),
-      smem_read_stage_idx_(0) {}
+      : smem_pointer_(smem_pointer),
+        smem_iterator_super_scale_(
+            LayoutSuperScale(IteratorSuperScale::Shape::kColumn),
+            get_super_scale_smem_ptr(),
+            {1, IteratorSuperScale::Shape::kColumn},
+            thread_idx),
+        smem_iterator_local_scale_(
+            LayoutLocalScale(IteratorLocalScale::Shape::kColumn),
+            get_local_scale_smem_ptr(),
+            {1, IteratorLocalScale::Shape::kColumn},
+            thread_idx),
+        smem_iterator_code_scale_(
+            LayoutCodeScaleZp(IteratorCodeScaleZp::Shape::kColumn),
+            get_code_scale_smem_ptr(),
+            {1, IteratorCodeScaleZp::Shape::kColumn},
+            thread_idx),
+        smem_iterator_code_zp_(
+            LayoutCodeScaleZp(IteratorCodeScaleZp::Shape::kColumn),
+            get_code_zp_smem_ptr(),
+            {1, IteratorCodeScaleZp::Shape::kColumn},
+            thread_idx),
+        smem_write_stage_idx_(0),
+        smem_read_stage_idx_(0) {}
 
   CUTLASS_DEVICE
   SuperTensorRef super_scale_ref() {
-    return {get_super_scale_smem_ptr(), LayoutSuperScale(IteratorSuperScale::Shape::kColumn)};
+    return {get_super_scale_smem_ptr(),
+            LayoutSuperScale(IteratorSuperScale::Shape::kColumn)};
   }
 
   CUTLASS_DEVICE
   LocalTensorRef local_scale_ref() {
-    return {get_local_scale_smem_ptr(), LayoutLocalScale(IteratorLocalScale::Shape::kColumn)};
+    return {get_local_scale_smem_ptr(),
+            LayoutLocalScale(IteratorLocalScale::Shape::kColumn)};
   }
 
   CUTLASS_DEVICE
   CodeTensorRef code_scale_ref() {
-    return {get_code_scale_smem_ptr(), LayoutCodeScaleZp(IteratorCodeScaleZp::Shape::kColumn)};
+    return {get_code_scale_smem_ptr(),
+            LayoutCodeScaleZp(IteratorCodeScaleZp::Shape::kColumn)};
   }
 
   CUTLASS_DEVICE
   CodeTensorRef code_zp_ref() {
-    return {get_code_zp_smem_ptr(), LayoutCodeScaleZp(IteratorCodeScaleZp::Shape::kColumn)};
+    return {get_code_zp_smem_ptr(),
+            LayoutCodeScaleZp(IteratorCodeScaleZp::Shape::kColumn)};
   }
 
   template <bool IsFirstStage>
-  CUTLASS_DEVICE
-  void copy_tiles_and_advance_per_stage(Arguments &quant_args, int stage) {
+  CUTLASS_DEVICE void copy_tiles_and_advance_per_stage(Arguments& quant_args,
+                                                       int stage) {
     if constexpr (IsFirstStage) {
-      // Load channel-wise super_scale to shared memory, which only needs to be done once.
+      // Load channel-wise super_scale to shared memory, which only needs to be
+      // done once.
       typename IteratorSuperScale::Fragment tb_frag_super_scale;
       tb_frag_super_scale.clear();
       quant_args.iterator_super_scale.load(tb_frag_super_scale);
       this->smem_iterator_super_scale_.store(tb_frag_super_scale);
 
-      // Load channel-wise code_scale to shared memory, which only needs to be done once.
+      // Load channel-wise code_scale to shared memory, which only needs to be
+      // done once.
       typename IteratorCodeScaleZp::Fragment tb_frag_code_scale;
       tb_frag_code_scale.clear();
       quant_args.iterator_code_scale.load(tb_frag_code_scale);
       this->smem_iterator_code_scale_.store(tb_frag_code_scale);
 
-      // Load channel-wise code_zp to shared memory, which only needs to be done once.
+      // Load channel-wise code_zp to shared memory, which only needs to be done
+      // once.
       typename IteratorCodeScaleZp::Fragment tb_frag_code_zp;
       tb_frag_code_zp.clear();
       quant_args.iterator_code_zp.load(tb_frag_code_zp);
@@ -231,20 +267,24 @@ public:
     }
 
     if ((stage % kStagesPerLocalScaleLoad) == 0) {
-      // Load group-wise local_scale to shared memory, which only needs to be done at each stage.
-      // Since 2 uint4b_t values of local_scale are saved in a single uint8_t, local_scale needs to be loaded once every two stages.
+      // Load group-wise local_scale to shared memory, which only needs to be
+      // done at each stage. Since 2 uint4b_t values of local_scale are saved in
+      // a single uint8_t, local_scale needs to be loaded once every two stages.
       using AccessType = typename IteratorLocalScale::AccessType;
-      cutlass::arch::CacheOperation::Kind const kCacheOp = (sizeof_bits<AccessType>::value == 128)
-          ? cutlass::arch::CacheOperation::Global : cutlass::arch::CacheOperation::Always;
+      cutlass::arch::CacheOperation::Kind const kCacheOp =
+          (sizeof_bits<AccessType>::value == 128)
+              ? cutlass::arch::CacheOperation::Global
+              : cutlass::arch::CacheOperation::Always;
 
       quant_args.iterator_local_scale.set_iteration_index(0);
       this->smem_iterator_local_scale_.set_iteration_index(0);
 
       // Async Copy for local_scale
       CUTLASS_PRAGMA_UNROLL
-      for (int j = 0; j < IteratorLocalScale::ThreadMap::Iterations::kCount; ++j) {
-        AccessType *dst_ptr =
-            reinterpret_cast<AccessType *>(this->smem_iterator_local_scale_.get());
+      for (int j = 0; j < IteratorLocalScale::ThreadMap::Iterations::kCount;
+           ++j) {
+        AccessType* dst_ptr = reinterpret_cast<AccessType*>(
+            this->smem_iterator_local_scale_.get());
 
         CUTLASS_PRAGMA_UNROLL
         for (int v = 0; v < IteratorLocalScale::kAccessesPerVector; ++v) {
@@ -255,8 +295,8 @@ public:
               IteratorLocalScale::ThreadMap::kElementsPerAccess /
               IteratorLocalScale::kAccessesPerVector / 8;
 
-              cutlass::arch::cp_async<kSrcBytes, kCacheOp>(
-                  dst_ptr + v, gmem_ptr, quant_args.iterator_local_scale.valid());
+          cutlass::arch::cp_async<kSrcBytes, kCacheOp>(
+              dst_ptr + v, gmem_ptr, quant_args.iterator_local_scale.valid());
         }
         ++quant_args.iterator_local_scale;
       }
@@ -265,13 +305,15 @@ public:
   }
 
   CUTLASS_DEVICE
-  void advance_smem_write_stage(Arguments &quant_args) {
+  void advance_smem_write_stage(Arguments& quant_args) {
     if (smem_write_stage_idx_ % kStagesPerLocalScaleLoad == 0) {
       // Advance global iterators
-      quant_args.iterator_local_scale.add_pointer_offset(quant_args.local_scale_pointer_offset);
+      quant_args.iterator_local_scale.add_pointer_offset(
+          quant_args.local_scale_pointer_offset);
 
       // Advance shared iterators
-      int smem_pointer_offset = IteratorLocalScale::Shape::kRow * IteratorLocalScale::Shape::kColumn;
+      int smem_pointer_offset =
+          IteratorLocalScale::Shape::kRow * IteratorLocalScale::Shape::kColumn;
       smem_iterator_local_scale_.add_pointer_offset(smem_pointer_offset);
     }
 
@@ -280,7 +322,8 @@ public:
 
     if (smem_write_stage_idx_ == kStagesPerLocalScaleLoad * kStages) {
       // Wrap back around to the 'start' of the circular buffer in shared memory
-      int pointer_offset = - kStages * IteratorLocalScale::Shape::kRow * IteratorLocalScale::Shape::kColumn;
+      int pointer_offset = -kStages * IteratorLocalScale::Shape::kRow *
+                           IteratorLocalScale::Shape::kColumn;
       smem_iterator_local_scale_.add_pointer_offset(pointer_offset);
       smem_write_stage_idx_ = 0;
     }
@@ -298,14 +341,14 @@ public:
 
     if (smem_read_stage_idx_ == kStagesPerLocalScaleLoad * kStages) {
       smem_read_stage_idx_ = 0;
-      byte_offset = - (kStages - 1) * kLocalScaleRows * kSmemColumns;
+      byte_offset = -(kStages - 1) * kLocalScaleRows * kSmemColumns;
     }
 
     return byte_offset;
   }
 
   CUTLASS_DEVICE
-  int clear_mask(Arguments &quant_args, bool cond) {
+  int clear_mask(Arguments& quant_args, bool cond) {
     quant_args.iterator_local_scale.clear_mask(cond);
   }
 };
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_unzip.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_unzip.h
index 9d49d5eb53..a1bc5a0ecf 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_unzip.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/threadblock/wint2x_unzip.h
@@ -29,18 +29,27 @@ namespace gemm {
 namespace threadblock {
 
 template <typename T, int N>
-using UnzipArray = cutlass::AlignedArray<T, N, (N * cutlass::sizeof_bits<T>::value / 8)>;
+using UnzipArray =
+    cutlass::AlignedArray<T, N, (N * cutlass::sizeof_bits<T>::value / 8)>;
 
-template <typename T, WintQuantMethod QuantMethod, int TileRows,
-          int TileColumns, int NumThreads = 128>
+template <typename T,
+          WintQuantMethod QuantMethod,
+          int TileRows,
+          int TileColumns,
+          int NumThreads = 128>
 struct UnzipAndDequantFunctor {
-  __device__ void operator()(const T *in_ptr, const T *supper_scale_ptr,
-                             T *out_ptr, const int64_t in_stride) {}
+  __device__ void operator()(const T *in_ptr,
+                             const T *supper_scale_ptr,
+                             T *out_ptr,
+                             const int64_t in_stride) {}
 };
 
 template <typename T, int TileRows, int TileColumns, int NumThreads>
-struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt25, TileRows,
-                              TileColumns, NumThreads> {
+struct UnzipAndDequantFunctor<T,
+                              WintQuantMethod::kWeightOnlyInt25,
+                              TileRows,
+                              TileColumns,
+                              NumThreads> {
   using ZippedT = uint16_t;
   using ScaleComputeT = float;
 
@@ -52,7 +61,8 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt25, TileRows,
   static constexpr int32_t kLocalScaleMask = 0x1FFF;
   static constexpr int32_t kBZP = 4;
 
-  __device__ inline T Compute(int32_t zipped_value, int32_t shift_bit,
+  __device__ inline T Compute(int32_t zipped_value,
+                              int32_t shift_bit,
                               ScaleComputeT scale) {
     int32_t shifted_value = (zipped_value >> shift_bit) & kWeightMask;
     int32_t value = shifted_value - kBZP;
@@ -61,8 +71,10 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt25, TileRows,
     return static_cast<T>(scaled_value);
   }
 
-  __device__ void operator()(const uint16_t *in_ptr, const T *super_scale_ptr,
-                             T *out_ptr, const int64_t in_stride) {
+  __device__ void operator()(const uint16_t *in_ptr,
+                             const T *super_scale_ptr,
+                             T *out_ptr,
+                             const int64_t in_stride) {
     int32_t shift_bits[7] = {13, 11, 9, 6, 4, 2, 0};
 
     int tid = threadIdx.x;
@@ -111,8 +123,11 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt25, TileRows,
 };
 
 template <typename T, int TileRows, int TileColumns, int NumThreads>
-struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
-                              TileColumns, NumThreads> {
+struct UnzipAndDequantFunctor<T,
+                              WintQuantMethod::kWeightOnlyInt2,
+                              TileRows,
+                              TileColumns,
+                              NumThreads> {
   using ZippedT = uint8_t;
   using ScaleComputeT = float;
 
@@ -129,9 +144,11 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
   // super_scale          [N]         T
 
   // code_scale, code_zp and super_scale
-  static constexpr int32_t kColumnWiseSmemBytes = (2 * sizeof(float) + sizeof(T)) * TileColumns;
+  static constexpr int32_t kColumnWiseSmemBytes =
+      (2 * sizeof(float) + sizeof(T)) * TileColumns;
   // zipped weights and local_scale
-  static constexpr int32_t kZippedSmemBytes = (TileRows / 4 + (TileRows + 127) / 128) * TileColumns;
+  static constexpr int32_t kZippedSmemBytes =
+      (TileRows / 4 + (TileRows + 127) / 128) * TileColumns;
 
   struct Arguments {
     uint8_t *weight_ptr;
@@ -140,14 +157,20 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
     float *code_zp_ptr;
     T *super_scale_ptr;
 
-    __device__ Arguments() : weight_ptr(nullptr), local_scale_ptr(nullptr), code_scale_ptr(nullptr), code_zp_ptr(nullptr), super_scale_ptr(nullptr) {}
+    __device__ Arguments()
+        : weight_ptr(nullptr),
+          local_scale_ptr(nullptr),
+          code_scale_ptr(nullptr),
+          code_zp_ptr(nullptr),
+          super_scale_ptr(nullptr) {}
 
     __device__ explicit Arguments(uint8_t *smem_ptr) {
       SetZippedPtrs(smem_ptr);
       SetColumnWisePtrs(smem_ptr + kZippedSmemBytes);
     }
 
-    __device__ Arguments(uint8_t *zipped_smem_ptr, uint8_t *column_wise_smem_ptr) {
+    __device__ Arguments(uint8_t *zipped_smem_ptr,
+                         uint8_t *column_wise_smem_ptr) {
       SetZippedPtrs(zipped_smem_ptr);
       SetColumnWisePtrs(column_wise_smem_ptr);
     }
@@ -159,15 +182,21 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
 
     __device__ void SetColumnWisePtrs(uint8_t *column_wise_smem_ptr) {
       code_scale_ptr = reinterpret_cast<float *>(column_wise_smem_ptr);
-      code_zp_ptr = reinterpret_cast<float *>(column_wise_smem_ptr + sizeof(float) * TileColumns);
-      super_scale_ptr = reinterpret_cast<T *>(column_wise_smem_ptr + 2 * sizeof(float) * TileColumns);
+      code_zp_ptr = reinterpret_cast<float *>(column_wise_smem_ptr +
+                                              sizeof(float) * TileColumns);
+      super_scale_ptr = reinterpret_cast<T *>(column_wise_smem_ptr +
+                                              2 * sizeof(float) * TileColumns);
     }
   };
 
-  __device__ void Load(const uint8_t *g_weight_ptr, const uint8_t *g_local_scale_ptr,
-                       const float *g_code_scale_ptr, const float *g_code_zp_ptr,
+  __device__ void Load(const uint8_t *g_weight_ptr,
+                       const uint8_t *g_local_scale_ptr,
+                       const float *g_code_scale_ptr,
+                       const float *g_code_zp_ptr,
                        const T *g_super_scale_ptr,
-                       Arguments *args, const int64_t in_stride, bool need_preload) {
+                       Arguments *args,
+                       const int64_t in_stride,
+                       bool need_preload) {
     int tid = threadIdx.x;
 
 #pragma unroll
@@ -186,7 +215,8 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
 #pragma unroll
       for (int ls_row_id = 0; ls_row_id < TileRows / 128; ++ls_row_id) {
         int local_scale_offset = ls_row_id * in_stride + col;
-        args->local_scale_ptr[ls_row_id * TileColumns + col] = g_local_scale_ptr[local_scale_offset];
+        args->local_scale_ptr[ls_row_id * TileColumns + col] =
+            g_local_scale_ptr[local_scale_offset];
       }
 
 #pragma unroll
@@ -205,10 +235,12 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
                             const float *g_code_scale_ptr,
                             const float *g_code_zp_ptr,
                             const T *g_super_scale_ptr,
-                            Arguments *args, const int64_t in_stride, bool need_preload) {
+                            Arguments *args,
+                            const int64_t in_stride,
+                            bool need_preload) {
     int tid = threadIdx.x;
 
-    constexpr int kBytesPerThread = 16; // 16B per thread
+    constexpr int kBytesPerThread = 16;  // 16B per thread
 
     constexpr int weight_size = TileRows / 4 * TileColumns;
     constexpr int local_scale_size = (TileRows + 127) / 128 * TileColumns;
@@ -216,87 +248,130 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
     constexpr int code_zp_size = sizeof(float) * TileColumns;
     constexpr int super_scale_size = sizeof(T) * TileColumns;
 
-    constexpr int total_size = weight_size + local_scale_size + code_scale_size + code_zp_size + super_scale_size;
+    constexpr int total_size = weight_size + local_scale_size +
+                               code_scale_size + code_zp_size +
+                               super_scale_size;
     constexpr int total_tasks = total_size / kBytesPerThread;
 
-    constexpr int cur_num_threads = total_tasks / ((total_tasks + NumThreads - 1) / NumThreads);
+    constexpr int cur_num_threads =
+        total_tasks / ((total_tasks + NumThreads - 1) / NumThreads);
 
     constexpr int weight_threads = weight_size * cur_num_threads / total_size;
-    constexpr int local_scale_threads = local_scale_size * cur_num_threads / total_size;
-    constexpr int code_scale_threads = code_scale_size * cur_num_threads / total_size;
+    constexpr int local_scale_threads =
+        local_scale_size * cur_num_threads / total_size;
+    constexpr int code_scale_threads =
+        code_scale_size * cur_num_threads / total_size;
     constexpr int code_zp_threads = code_zp_size * cur_num_threads / total_size;
-    constexpr int super_scale_threads = super_scale_size * cur_num_threads / total_size;
+    constexpr int super_scale_threads =
+        super_scale_size * cur_num_threads / total_size;
 
     static_assert(TileColumns % weight_threads == 0,
-                  "TileColumns must be divisible by weight_threads to ensure correct thread mapping.");
+                  "TileColumns must be divisible by weight_threads to ensure "
+                  "correct thread mapping.");
 
     static_assert(TileColumns % local_scale_threads == 0,
-                  "TileColumns must be divisible by local_scale_threads to ensure correct thread mapping.");
+                  "TileColumns must be divisible by local_scale_threads to "
+                  "ensure correct thread mapping.");
 
     if (tid < weight_threads) {
       constexpr int weight_per_thread_size = weight_size / weight_threads;
-      constexpr int kIterations = (weight_per_thread_size + kBytesPerThread - 1) / kBytesPerThread;
+      constexpr int kIterations =
+          (weight_per_thread_size + kBytesPerThread - 1) / kBytesPerThread;
 
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < kIterations; ++i) {
-          int z_offset = (tid * weight_per_thread_size + i * kBytesPerThread);
-          int g_offset = z_offset / TileColumns * in_stride + z_offset % TileColumns;
-          cutlass::arch::cp_async<kBytesPerThread, cutlass::arch::CacheOperation::Global>(
-              args->weight_ptr + z_offset, g_weight_ptr + g_offset, true);
+        int z_offset = (tid * weight_per_thread_size + i * kBytesPerThread);
+        int g_offset =
+            z_offset / TileColumns * in_stride + z_offset % TileColumns;
+        cutlass::arch::cp_async<kBytesPerThread,
+                                cutlass::arch::CacheOperation::Global>(
+            args->weight_ptr + z_offset, g_weight_ptr + g_offset, true);
       }
     } else if (tid < weight_threads + local_scale_threads) {
       constexpr int start_thread_id = weight_threads;
-      constexpr int local_scale_per_thread_size = local_scale_size / local_scale_threads;
-      constexpr int kIterations = (local_scale_per_thread_size + kBytesPerThread - 1) / kBytesPerThread;
+      constexpr int local_scale_per_thread_size =
+          local_scale_size / local_scale_threads;
+      constexpr int kIterations =
+          (local_scale_per_thread_size + kBytesPerThread - 1) / kBytesPerThread;
 
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < kIterations; ++i) {
-          int z_offset = (tid - start_thread_id) * local_scale_per_thread_size + i * kBytesPerThread;
-          int g_offset = z_offset / TileColumns * in_stride + z_offset % TileColumns;
-          cutlass::arch::cp_async<kBytesPerThread, cutlass::arch::CacheOperation::Global>(
-              args->local_scale_ptr + z_offset, g_local_scale_ptr + g_offset, true);
+        int z_offset = (tid - start_thread_id) * local_scale_per_thread_size +
+                       i * kBytesPerThread;
+        int g_offset =
+            z_offset / TileColumns * in_stride + z_offset % TileColumns;
+        cutlass::arch::cp_async<kBytesPerThread,
+                                cutlass::arch::CacheOperation::Global>(
+            args->local_scale_ptr + z_offset,
+            g_local_scale_ptr + g_offset,
+            true);
       }
     } else if (need_preload) {
       if (tid < weight_threads + local_scale_threads + code_scale_threads) {
         constexpr int start_thread_id = weight_threads + local_scale_threads;
-        constexpr int code_scale_per_thread_size = code_scale_size / code_scale_threads;
-        constexpr int kIterations = (code_scale_per_thread_size + kBytesPerThread - 1) / kBytesPerThread;
+        constexpr int code_scale_per_thread_size =
+            code_scale_size / code_scale_threads;
+        constexpr int kIterations =
+            (code_scale_per_thread_size + kBytesPerThread - 1) /
+            kBytesPerThread;
 
         CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < kIterations; ++i) {
-          int offset = ((tid - start_thread_id) * code_scale_per_thread_size + i * kBytesPerThread) / sizeof(float);
-          cutlass::arch::cp_async<kBytesPerThread, cutlass::arch::CacheOperation::Global>(
+          int offset = ((tid - start_thread_id) * code_scale_per_thread_size +
+                        i * kBytesPerThread) /
+                       sizeof(float);
+          cutlass::arch::cp_async<kBytesPerThread,
+                                  cutlass::arch::CacheOperation::Global>(
               args->code_scale_ptr + offset, g_code_scale_ptr + offset, true);
         }
-      } else if (tid < weight_threads + local_scale_threads + code_scale_threads + code_zp_threads) {
-        constexpr int start_thread_id = weight_threads + local_scale_threads + code_scale_threads;
+      } else if (tid < weight_threads + local_scale_threads +
+                           code_scale_threads + code_zp_threads) {
+        constexpr int start_thread_id =
+            weight_threads + local_scale_threads + code_scale_threads;
         constexpr int code_zp_per_thread_size = code_zp_size / code_zp_threads;
-        constexpr int kIterations = (code_zp_per_thread_size + kBytesPerThread - 1) / kBytesPerThread;
+        constexpr int kIterations =
+            (code_zp_per_thread_size + kBytesPerThread - 1) / kBytesPerThread;
 
         CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < kIterations; ++i) {
-          int offset = ((tid - start_thread_id) * code_zp_per_thread_size + i * kBytesPerThread) / sizeof(float);
-          cutlass::arch::cp_async<kBytesPerThread, cutlass::arch::CacheOperation::Global>(
+          int offset = ((tid - start_thread_id) * code_zp_per_thread_size +
+                        i * kBytesPerThread) /
+                       sizeof(float);
+          cutlass::arch::cp_async<kBytesPerThread,
+                                  cutlass::arch::CacheOperation::Global>(
               args->code_zp_ptr + offset, g_code_zp_ptr + offset, true);
         }
-      } else if (tid < weight_threads + local_scale_threads + code_scale_threads + code_zp_threads + super_scale_threads) {
+      } else if (tid < weight_threads + local_scale_threads +
+                           code_scale_threads + code_zp_threads +
+                           super_scale_threads) {
         if (g_super_scale_ptr) {
-          constexpr int start_thread_id = weight_threads + local_scale_threads + code_scale_threads + code_zp_threads;
-          constexpr int super_scale_per_thread_size = super_scale_size / super_scale_threads;
-          constexpr int kIterations = (super_scale_per_thread_size + kBytesPerThread - 1) / kBytesPerThread;
+          constexpr int start_thread_id = weight_threads + local_scale_threads +
+                                          code_scale_threads + code_zp_threads;
+          constexpr int super_scale_per_thread_size =
+              super_scale_size / super_scale_threads;
+          constexpr int kIterations =
+              (super_scale_per_thread_size + kBytesPerThread - 1) /
+              kBytesPerThread;
 
           CUTLASS_PRAGMA_UNROLL
           for (int i = 0; i < kIterations; ++i) {
-            int offset = ((tid - start_thread_id) * super_scale_per_thread_size + i * kBytesPerThread) / sizeof(T);
-            cutlass::arch::cp_async<kBytesPerThread, cutlass::arch::CacheOperation::Global>(
-                args->super_scale_ptr + offset, g_super_scale_ptr + offset, true);
+            int offset =
+                ((tid - start_thread_id) * super_scale_per_thread_size +
+                 i * kBytesPerThread) /
+                sizeof(T);
+            cutlass::arch::cp_async<kBytesPerThread,
+                                    cutlass::arch::CacheOperation::Global>(
+                args->super_scale_ptr + offset,
+                g_super_scale_ptr + offset,
+                true);
           }
         }
       }
     }
   }
 
-  __device__ void Compute(const Arguments &args, T *out_ptr,
+  __device__ void Compute(const Arguments &args,
+                          T *out_ptr,
                           const int64_t block_start_row) {
     int32_t shift_bits[4] = {9, 6, 3, 0};
 
@@ -333,9 +408,9 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
 
 #pragma unroll
         for (int zipped_row = 0; zipped_row < 16; ++zipped_row) {
-          int32_t decode_value =
-              static_cast<int32_t>(floor(zipped_value[zipped_row] * code_scale + code_zp +
-                                         static_cast<ScaleComputeT>(0.5)));
+          int32_t decode_value = static_cast<int32_t>(
+              floor(zipped_value[zipped_row] * code_scale + code_zp +
+                    static_cast<ScaleComputeT>(0.5)));
 
           int row = group_id * 64 + zipped_row * 4;
 
@@ -355,14 +430,17 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
     __syncthreads();
   }
 
-  __device__ void ComputeVectorized(const Arguments &args, T *out_ptr,
+  __device__ void ComputeVectorized(const Arguments &args,
+                                    T *out_ptr,
                                     const int64_t block_start_row) {
-    constexpr int kNumWeightsPerThread = TileRows * TileColumns / (4 * NumThreads);
+    constexpr int kNumWeightsPerThread =
+        TileRows * TileColumns / (4 * NumThreads);
     constexpr int N = (kNumWeightsPerThread >= 32) ? 4 : 2;
     constexpr int RowStride = NumThreads * N / TileColumns;
     constexpr int kNumIters = kNumWeightsPerThread / N;
 
-    static_assert(N * NumThreads >= TileColumns, "N * NumThreads should be no less than TileColumns.");
+    static_assert(N * NumThreads >= TileColumns,
+                  "N * NumThreads should be no less than TileColumns.");
 
     constexpr ScaleComputeT decode_value_zp = static_cast<ScaleComputeT>(0.5);
 
@@ -373,19 +451,22 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
     static_assert(TileRows <= 128, "TileRows is expected to no more than 128.");
 
     UnzipArray<uint8_t, N> local_scales =
-        *reinterpret_cast<const UnzipArray<uint8_t, N> *>(args.local_scale_ptr + begin_col_id);
+        *reinterpret_cast<const UnzipArray<uint8_t, N> *>(args.local_scale_ptr +
+                                                          begin_col_id);
 
     UnzipArray<uint8_t, N> zipped_values[2];
     int zipped_offset = begin_row_id * TileColumns + begin_col_id;
-    zipped_values[0] =
-        *reinterpret_cast<const UnzipArray<uint8_t, N> *>(args.weight_ptr + zipped_offset);
+    zipped_values[0] = *reinterpret_cast<const UnzipArray<uint8_t, N> *>(
+        args.weight_ptr + zipped_offset);
 
-    UnzipArray<T, N> super_scales =
-        *reinterpret_cast<const UnzipArray<T, N> *>(args.super_scale_ptr + begin_col_id);
+    UnzipArray<T, N> super_scales = *reinterpret_cast<const UnzipArray<T, N> *>(
+        args.super_scale_ptr + begin_col_id);
     UnzipArray<float, N> code_scales =
-        *reinterpret_cast<const UnzipArray<float, N> *>(args.code_scale_ptr + begin_col_id);
+        *reinterpret_cast<const UnzipArray<float, N> *>(args.code_scale_ptr +
+                                                        begin_col_id);
     UnzipArray<float, N> code_zps =
-        *reinterpret_cast<const UnzipArray<float, N> *>(args.code_zp_ptr + begin_col_id);
+        *reinterpret_cast<const UnzipArray<float, N> *>(args.code_zp_ptr +
+                                                        begin_col_id);
 
     // special for TileRows = 64
     int local_scale_shift = (((block_start_row / 64) + 1) & 1) * 4;
@@ -394,9 +475,10 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
 #pragma unroll
     for (int i = 0; i < N; ++i) {
       int32_t shifted_local_scale =
-          (static_cast<int32_t>(local_scales[i]) >> local_scale_shift) & kLocalScaleMask;
-      scales[i] =
-            static_cast<ScaleComputeT>(shifted_local_scale) * static_cast<ScaleComputeT>(super_scales[i]);
+          (static_cast<int32_t>(local_scales[i]) >> local_scale_shift) &
+          kLocalScaleMask;
+      scales[i] = static_cast<ScaleComputeT>(shifted_local_scale) *
+                  static_cast<ScaleComputeT>(super_scales[i]);
     }
 
 #pragma unroll
@@ -405,26 +487,33 @@ struct UnzipAndDequantFunctor<T, WintQuantMethod::kWeightOnlyInt2, TileRows,
       int row = zipped_row * 4;
 
       if (iter_id < kNumIters - 1) {
-        int zipped_offset = (zipped_row + RowStride) * TileColumns + begin_col_id;
+        int zipped_offset =
+            (zipped_row + RowStride) * TileColumns + begin_col_id;
         zipped_values[(iter_id + 1) & 1] =
-            *reinterpret_cast<const UnzipArray<uint8_t, N> *>(args.weight_ptr + zipped_offset);
+            *reinterpret_cast<const UnzipArray<uint8_t, N> *>(args.weight_ptr +
+                                                              zipped_offset);
       }
 
       UnzipArray<T, N> outs[4];
 
 #pragma unroll
       for (int i = 0; i < N; ++i) {
-        int32_t decode_value =
-            static_cast<int32_t>(floor(static_cast<ScaleComputeT>(zipped_values[iter_id & 1][i]) * code_scales[i]
-                                       + code_zps[i] + decode_value_zp));
+        int32_t decode_value = static_cast<int32_t>(
+            floor(static_cast<ScaleComputeT>(zipped_values[iter_id & 1][i]) *
+                      code_scales[i] +
+                  code_zps[i] + decode_value_zp));
 
-        ScaleComputeT value_3 = static_cast<ScaleComputeT>((decode_value & kWeightMask) - kBZP);
+        ScaleComputeT value_3 =
+            static_cast<ScaleComputeT>((decode_value & kWeightMask) - kBZP);
         decode_value >>= 3;
-        ScaleComputeT value_2 = static_cast<ScaleComputeT>((decode_value & kWeightMask) - kBZP);
+        ScaleComputeT value_2 =
+            static_cast<ScaleComputeT>((decode_value & kWeightMask) - kBZP);
         decode_value >>= 3;
-        ScaleComputeT value_1 = static_cast<ScaleComputeT>((decode_value & kWeightMask) - kBZP);
+        ScaleComputeT value_1 =
+            static_cast<ScaleComputeT>((decode_value & kWeightMask) - kBZP);
         decode_value >>= 3;
-        ScaleComputeT value_0 = static_cast<ScaleComputeT>((decode_value & kWeightMask) - kBZP);
+        ScaleComputeT value_0 =
+            static_cast<ScaleComputeT>((decode_value & kWeightMask) - kBZP);
         outs[0][i] = static_cast<T>(scales[i] * value_0);
         outs[1][i] = static_cast<T>(scales[i] * value_1);
         outs[2][i] = static_cast<T>(scales[i] * value_2);
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/default_mma_tensor_op.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/default_mma_tensor_op.h
index af4298df5e..8245aff71c 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/default_mma_tensor_op.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/default_mma_tensor_op.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,18 +18,20 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
-    \brief Default warp-level GEMM operators selected by data type, size, and layouts of operands.
+    \brief Default warp-level GEMM operators selected by data type, size, and
+   layouts of operands.
 */
 
 #pragma once
@@ -70,35 +72,60 @@ template <
     /// Store the accumulators in row major or column major.  Row major is used
     /// when output layout is interleaved.
     bool AccumulatorsInRowMajor>
-struct DefaultMmaTensorOp<WarpShape_, InstructionShape_, ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
-    arch::OpMultiplyAddDequantizeInterleavedBToA, PartitionsK, AccumulatorsInRowMajor>
-{
+struct DefaultMmaTensorOp<WarpShape_,
+                          InstructionShape_,
+                          ElementA,
+                          LayoutA,
+                          ElementB,
+                          LayoutB,
+                          ElementC,
+                          LayoutC,
+                          arch::OpMultiplyAddDequantizeInterleavedBToA,
+                          PartitionsK,
+                          AccumulatorsInRowMajor> {
+ private:
+  // Shape for computing the FP16s
+  using ComputeInstructionShape = InstructionShape_;
 
-private:
-    // Shape for computing the FP16s
-    using ComputeInstructionShape = InstructionShape_;
+  // Chosen so we get K=16 for int8, K=32 for int4, K=64 for int2.
+  static constexpr int LoadInstructionK = 128 / sizeof_bits<ElementB>::value;
 
-    // Chosen so we get K=16 for int8, K=32 for int4, K=64 for int2.
-    static constexpr int LoadInstructionK = 128 / sizeof_bits<ElementB>::value;
+  // Shape for loading the narrow data type from shared memory
+  using LoadInstructionShape =
+      GemmShape<InstructionShape_::kM, InstructionShape_::kN, LoadInstructionK>;
 
-    // Shape for loading the narrow data type from shared memory
-    using LoadInstructionShape = GemmShape<InstructionShape_::kM, InstructionShape_::kN, LoadInstructionK>;
+ public:
+  using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<
+      cutlass::arch::Mma<InstructionShape_,
+                         32,
+                         ElementA,
+                         cutlass::layout::RowMajor,
+                         ElementA,
+                         cutlass::layout::ColumnMajor,
+                         ElementC,
+                         cutlass::layout::RowMajor,
+                         arch::OpMultiplyAdd>,
+      cutlass::MatrixShape<1, 1>>;
 
-public:
-    using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<
-        cutlass::arch::Mma<InstructionShape_, 32, ElementA, cutlass::layout::RowMajor, ElementA,
-            cutlass::layout::ColumnMajor, ElementC, cutlass::layout::RowMajor, arch::OpMultiplyAdd>,
-        cutlass::MatrixShape<1, 1>>;
-
-    // Define the warp-level tensor op
-    using Type = cutlass::gemm::warp::MmaTensorOpComputeBWithF16<WarpShape_, ElementA, LayoutA, ElementB, LayoutB,
-        ElementC, LayoutC, Policy, LoadInstructionShape, PartitionsK, AccumulatorsInRowMajor>;
+  // Define the warp-level tensor op
+  using Type =
+      cutlass::gemm::warp::MmaTensorOpComputeBWithF16<WarpShape_,
+                                                      ElementA,
+                                                      LayoutA,
+                                                      ElementB,
+                                                      LayoutB,
+                                                      ElementC,
+                                                      LayoutC,
+                                                      Policy,
+                                                      LoadInstructionShape,
+                                                      PartitionsK,
+                                                      AccumulatorsInRowMajor>;
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace warp
-} // namespace gemm
-} // namespace cutlass
+}  // namespace warp
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_compute_B_with_f16.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_compute_B_with_f16.h
index 64136a9758..edc37d72a1 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_compute_B_with_f16.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_compute_B_with_f16.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,19 +18,20 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
-    \brief Templates implementing warp-level matrix multiply-accumulate operations targeting
-      Tensor Cores.
+    \brief Templates implementing warp-level matrix multiply-accumulate
+   operations targeting Tensor Cores.
 */
 
 #pragma once
@@ -63,7 +64,8 @@ namespace gemm {
 namespace warp {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
-/// Structure to compute the matrix product targeting Tensor Cores, for the case when A is floating point and B is quantized integer.
+/// Structure to compute the matrix product targeting Tensor Cores, for the case
+/// when A is floating point and B is quantized integer.
 template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
@@ -90,213 +92,241 @@ template <
     bool AccumulatorsInRowMajor = false,
     /// Used for partial specialization
     typename Enable = bool>
-class MmaTensorOpComputeBWithF16
-{
-public:
-    /// Shape of warp-level matrix operation (concept: GemmShape)
-    using Shape = Shape_;
+class MmaTensorOpComputeBWithF16 {
+ public:
+  /// Shape of warp-level matrix operation (concept: GemmShape)
+  using Shape = Shape_;
 
-    /// Data type of multiplicand A
-    using ElementA = ElementA_;
+  /// Data type of multiplicand A
+  using ElementA = ElementA_;
 
-    /// Layout of multiplicand A
-    using LayoutA = LayoutA_;
+  /// Layout of multiplicand A
+  using LayoutA = LayoutA_;
 
-    /// Data type of multiplicand B
-    using ElementB = ElementB_;
+  /// Data type of multiplicand B
+  using ElementB = ElementB_;
 
-    /// Layout of multiplicand B
-    using LayoutB = LayoutB_;
+  /// Layout of multiplicand B
+  using LayoutB = LayoutB_;
 
-    /// Data type of accumulator matrix C
-    using ElementC = ElementC_;
+  /// Data type of accumulator matrix C
+  using ElementC = ElementC_;
 
-    /// Layout of accumulator matrix C
-    using LayoutC = LayoutC_;
+  /// Layout of accumulator matrix C
+  using LayoutC = LayoutC_;
 
-    /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
-    using Policy = Policy_;
+  /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
+  using Policy = Policy_;
 
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    using ArchMmaOperator = typename Policy::Operator;
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  using ArchMmaOperator = typename Policy::Operator;
 
-    /// Indicates math operator
-    using MathOperator = typename ArchMmaOperator::Operator;
+  /// Indicates math operator
+  using MathOperator = typename ArchMmaOperator::Operator;
 
-    /// Architecture tag from underlying instruction
-    using ArchTag = typename ArchMmaOperator::ArchTag;
-    static_assert((platform::is_same<typename ArchMmaOperator::ElementA, half_t>::value
-                      && platform::is_same<typename ArchMmaOperator::ElementB, half_t>::value)
-            || (platform::is_same<typename ArchMmaOperator::ElementA, bfloat16_t>::value
-                && platform::is_same<typename ArchMmaOperator::ElementB, bfloat16_t>::value
-                && ArchTag::kMinComputeCapability >= 80)
-            || (platform::is_same<typename ArchMmaOperator::ElementA, float_e4m3_t>::value
-                && platform::is_same<typename ArchMmaOperator::ElementB, float_e4m3_t>::value
-                && ArchTag::kMinComputeCapability >= 89),
-        "MmaTensorOpCvtBToA only supports underlying HMMA/QMMA");
+  /// Architecture tag from underlying instruction
+  using ArchTag = typename ArchMmaOperator::ArchTag;
+  static_assert(
+      (platform::is_same<typename ArchMmaOperator::ElementA, half_t>::value &&
+       platform::is_same<typename ArchMmaOperator::ElementB, half_t>::value) ||
+          (platform::is_same<typename ArchMmaOperator::ElementA,
+                             bfloat16_t>::value &&
+           platform::is_same<typename ArchMmaOperator::ElementB,
+                             bfloat16_t>::value &&
+           ArchTag::kMinComputeCapability >= 80) ||
+          (platform::is_same<typename ArchMmaOperator::ElementA,
+                             float_e4m3_t>::value &&
+           platform::is_same<typename ArchMmaOperator::ElementB,
+                             float_e4m3_t>::value &&
+           ArchTag::kMinComputeCapability >= 89),
+      "MmaTensorOpCvtBToA only supports underlying HMMA/QMMA");
 
-    static_assert(platform::is_same<ElementA, half_t>::value
-            || (platform::is_same<ElementA, bfloat16_t>::value && ArchTag::kMinComputeCapability >= 80)
-            || (platform::is_same<ElementA, float_e4m3_t>::value && ArchTag::kMinComputeCapability >= 89),
-        "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on Ampere+, or FP8 on Ada");
+  static_assert(platform::is_same<ElementA, half_t>::value ||
+                    (platform::is_same<ElementA, bfloat16_t>::value &&
+                     ArchTag::kMinComputeCapability >= 80) ||
+                    (platform::is_same<ElementA, float_e4m3_t>::value &&
+                     ArchTag::kMinComputeCapability >= 89),
+                "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on Ampere+, "
+                "or FP8 on Ada");
 
-    /// Indicates class of matrix operator
-    using OperatorClass = arch::OpClassTensorOp;
+  /// Indicates class of matrix operator
+  using OperatorClass = arch::OpClassTensorOp;
 
-    /// Shape of underlying instruction
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  /// Shape of underlying instruction
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    /// Instruction shape to override shared memory iterators with
-    using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
+  /// Instruction shape to override shared memory iterators with
+  using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
+
+  static_assert(SharedMemoryInstructionShape::kM == InstructionShape::kM,
+                "M dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kN == InstructionShape::kN,
+                "N dimension of compute instruction must match load");
+
+  static constexpr int kExpansionFactor =
+      SharedMemoryInstructionShape::kK / InstructionShape::kK;
+
+  static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
+
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
+
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
+
+  /// Number of threads participating in warp-level matrix product
+  static int const kThreadCount = 32;
+
+  /// Number of partitions along K dimension
+  static int const kPartitionsK = PartitionsK_;
+
+ public:
+  /// Iterates over the A operand in memory
+  using IteratorA = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kM, Shape::kK>,
+      Operand::kA,
+      ElementA,
+      LayoutA,
+      MatrixShape<InstructionShape::kM, InstructionShape::kK>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
+
+  /// Storage for A tile
+  using FragmentA = typename IteratorA::Fragment;
+
+  /// Storage for transformed A tile
+  using TransformedFragmentA =
+      Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
+
+  /// Iterates over the B operand in memory
+  using IteratorB = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kK, Shape::kN>,
+      Operand::kB,
+      ElementB,
+      LayoutB,
+      MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
+
+  /// Storage for B tile
+  using FragmentB = typename IteratorB::Fragment;
+
+  /// Storage for transformed B tile
+  using TransformedFragmentB =
+      Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
+
+  /// Iterates over the C operand in memory
+  using IteratorC =
+      MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>,
+                                         ElementC,
+                                         LayoutC,
+                                         typename ArchMmaOperator::Shape,
+                                         typename Policy::OpDelta>;
+
+  /// Storage for C tile
+  using FragmentC = typename IteratorC::Fragment;
+
+  /// Number of mma operations performed
+  using MmaIterations =
+      MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) /
+                      ArchMmaOperator::Shape::kM,
+                  (Shape::kN + ArchMmaOperator::Shape::kN - 1) /
+                      ArchMmaOperator::Shape::kN>;
+
+ public:
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  ArchMmaOperator mma;
+
+ public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  CUTLASS_DEVICE
+  MmaTensorOpComputeBWithF16() {}
+
+  /// Performs a warp-level matrix multiply-accumulate operation
+  CUTLASS_DEVICE
+  void operator()(FragmentC& D,
+                  TransformedFragmentA const& A,
+                  TransformedFragmentB const& B,
+                  FragmentC const& C,
+                  int const warp_tileB_k_offset) const {
+    using MmaOperandA = typename ArchMmaOperator::FragmentA;
+    using MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using MmaOperandC = typename ArchMmaOperator::FragmentC;
 
     static_assert(
-        SharedMemoryInstructionShape::kM == InstructionShape::kM, "M dimension of compute instruction must match load");
-    static_assert(
-        SharedMemoryInstructionShape::kN == InstructionShape::kN, "N dimension of compute instruction must match load");
+        TransformedFragmentB::kElements ==
+            MmaOperandB::kElements * kExpansionFactor * MmaIterations::kColumn,
+        "Each thread should have a pack of mma registers for each column "
+        "iteration AND for the expanded K dim of "
+        "B");
 
-    static constexpr int kExpansionFactor = SharedMemoryInstructionShape::kK / InstructionShape::kK;
+    D = C;
 
-    static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
-
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = ComplexTransform::kNone;
-
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = ComplexTransform::kNone;
-
-    /// Number of threads participating in warp-level matrix product
-    static int const kThreadCount = 32;
-
-    /// Number of partitions along K dimension
-    static int const kPartitionsK = PartitionsK_;
-
-public:
-    /// Iterates over the A operand in memory
-    using IteratorA
-        = MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kM, Shape::kK>, Operand::kA, ElementA, LayoutA,
-            MatrixShape<InstructionShape::kM, InstructionShape::kK>, Policy::OpDelta::kRow, kThreadCount, kPartitionsK>;
-
-    /// Storage for A tile
-    using FragmentA = typename IteratorA::Fragment;
-
-    /// Storage for transformed A tile
-    using TransformedFragmentA = Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
-
-    /// Iterates over the B operand in memory
-    using IteratorB = MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kK, Shape::kN>, Operand::kB, ElementB,
-        LayoutB, MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>, Policy::OpDelta::kRow,
-        kThreadCount, kPartitionsK>;
-
-    /// Storage for B tile
-    using FragmentB = typename IteratorB::Fragment;
-
-    /// Storage for transformed B tile
-    using TransformedFragmentB = Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
-
-    /// Iterates over the C operand in memory
-    using IteratorC = MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>, ElementC, LayoutC,
-        typename ArchMmaOperator::Shape, typename Policy::OpDelta>;
-
-    /// Storage for C tile
-    using FragmentC = typename IteratorC::Fragment;
-
-    /// Number of mma operations performed
-    using MmaIterations = MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM,
-        (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN>;
-
-public:
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    ArchMmaOperator mma;
-
-public:
-    //
-    // Methods
-    //
-
-    /// Ctor
-    CUTLASS_DEVICE
-    MmaTensorOpComputeBWithF16() {}
-
-    /// Performs a warp-level matrix multiply-accumulate operation
-    CUTLASS_DEVICE
-    void operator()(FragmentC& D, TransformedFragmentA const& A, TransformedFragmentB const& B, FragmentC const& C,
-        int const warp_tileB_k_offset) const
-    {
-
-        using MmaOperandA = typename ArchMmaOperator::FragmentA;
-        using MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using MmaOperandC = typename ArchMmaOperator::FragmentC;
-
-        static_assert(
-            TransformedFragmentB::kElements == MmaOperandB::kElements * kExpansionFactor * MmaIterations::kColumn,
-            "Each thread should have a pack of mma registers for each column iteration AND for the expanded K dim of "
-            "B");
-
-        D = C;
-
-        MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
-        MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
-        MmaOperandC* ptr_D = reinterpret_cast<MmaOperandC*>(&D);
+    MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
+    MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
+    MmaOperandC* ptr_D = reinterpret_cast<MmaOperandC*>(&D);
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
-        // Serpentine visitation order maximizing reuse of Rb
-        CUTLASS_PRAGMA_UNROLL
-        for (int n = 0; n < MmaIterations::kColumn; ++n)
-        {
+    // Serpentine visitation order maximizing reuse of Rb
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < MmaIterations::kColumn; ++n) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int m = 0; m < MmaIterations::kRow; ++m) {
+        int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
 
-            CUTLASS_PRAGMA_UNROLL
-            for (int m = 0; m < MmaIterations::kRow; ++m)
-            {
-
-                int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
-
-                int n_offsetB = warp_tileB_k_offset + kExpansionFactor * n;
-                if (AccumulatorsInRowMajor)
-                { // matrix B is reordered
-                    mma(ptr_D[n + m_serpentine * MmaIterations::kColumn], ptr_A[m_serpentine], ptr_B[n_offsetB],
-                        ptr_D[n + m_serpentine * MmaIterations::kColumn]);
-                }
-                else
-                {
-                    mma(ptr_D[m_serpentine + n * MmaIterations::kRow], ptr_A[m_serpentine], ptr_B[n_offsetB],
-                        ptr_D[m_serpentine + n * MmaIterations::kRow]);
-                }
-            }
+        int n_offsetB = warp_tileB_k_offset + kExpansionFactor * n;
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n + m_serpentine * MmaIterations::kColumn],
+              ptr_A[m_serpentine],
+              ptr_B[n_offsetB],
+              ptr_D[n + m_serpentine * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m_serpentine + n * MmaIterations::kRow],
+              ptr_A[m_serpentine],
+              ptr_B[n_offsetB],
+              ptr_D[m_serpentine + n * MmaIterations::kRow]);
         }
-#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
-        // Serpentine visitation order maximizing reuse of Ra
-        CUTLASS_PRAGMA_UNROLL
-        for (int m = 0; m < MmaIterations::kRow; ++m)
-        {
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int n = 0; n < MmaIterations::kColumn; ++n)
-            {
-
-                int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
-
-                int n_serpentine_offsetB = warp_tileB_k_offset + kExpansionFactor * n_serpentine;
-                if (AccumulatorsInRowMajor)
-                { // matrix B is reordered
-                    mma(ptr_D[n_serpentine + m * MmaIterations::kColumn], ptr_A[m], ptr_B[n_serpentine_offsetB],
-                        ptr_D[n_serpentine + m * MmaIterations::kColumn]);
-                }
-                else
-                {
-                    mma(ptr_D[m + n_serpentine * MmaIterations::kRow], ptr_A[m], ptr_B[n_serpentine_offsetB],
-                        ptr_D[m + n_serpentine * MmaIterations::kRow]);
-                }
-            }
-        }
-#else
-        assert(0);
-#endif
+      }
     }
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+    // Serpentine visitation order maximizing reuse of Ra
+    CUTLASS_PRAGMA_UNROLL
+    for (int m = 0; m < MmaIterations::kRow; ++m) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < MmaIterations::kColumn; ++n) {
+        int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
+
+        int n_serpentine_offsetB =
+            warp_tileB_k_offset + kExpansionFactor * n_serpentine;
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n_serpentine + m * MmaIterations::kColumn],
+              ptr_A[m],
+              ptr_B[n_serpentine_offsetB],
+              ptr_D[n_serpentine + m * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m + n_serpentine * MmaIterations::kRow],
+              ptr_A[m],
+              ptr_B[n_serpentine_offsetB],
+              ptr_D[m + n_serpentine * MmaIterations::kRow]);
+        }
+      }
+    }
+#else
+    assert(0);
+#endif
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
-/// Structure to compute the matrix product targeting Tensor Cores, for the case when A is floating point and B is quantized integer.
-/// Specialization for B of uint2b_t.
+/// Structure to compute the matrix product targeting Tensor Cores, for the case
+/// when A is floating point and B is quantized integer. Specialization for B of
+/// uint2b_t.
 template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
@@ -319,214 +349,232 @@ template <
     /// Store the accumulators in row major or column major.  Row major is used
     /// when output layout is interleaved.
     bool AccumulatorsInRowMajor>
-class MmaTensorOpComputeBWithF16<
-    Shape_,
-    ElementA_,
-    LayoutA_,
-    uint2b_t,
-    LayoutB_,
-    ElementC_,
-    LayoutC_,
-    Policy_,
-    SharedMemoryInstructionShape_,
-    PartitionsK_,
-    AccumulatorsInRowMajor>
-{
-public:
-    /// Shape of warp-level matrix operation (concept: GemmShape)
-    using Shape = Shape_;
+class MmaTensorOpComputeBWithF16<Shape_,
+                                 ElementA_,
+                                 LayoutA_,
+                                 uint2b_t,
+                                 LayoutB_,
+                                 ElementC_,
+                                 LayoutC_,
+                                 Policy_,
+                                 SharedMemoryInstructionShape_,
+                                 PartitionsK_,
+                                 AccumulatorsInRowMajor> {
+ public:
+  /// Shape of warp-level matrix operation (concept: GemmShape)
+  using Shape = Shape_;
 
-    /// Data type of multiplicand A
-    using ElementA = ElementA_;
+  /// Data type of multiplicand A
+  using ElementA = ElementA_;
 
-    /// Layout of multiplicand A
-    using LayoutA = LayoutA_;
+  /// Layout of multiplicand A
+  using LayoutA = LayoutA_;
 
-    /// Data type of multiplicand B
-    using ElementB = uint2b_t;
+  /// Data type of multiplicand B
+  using ElementB = uint2b_t;
 
-    /// Layout of multiplicand B
-    using LayoutB = LayoutB_;
+  /// Layout of multiplicand B
+  using LayoutB = LayoutB_;
 
-    /// Data type of accumulator matrix C
-    using ElementC = ElementC_;
+  /// Data type of accumulator matrix C
+  using ElementC = ElementC_;
 
-    /// Layout of accumulator matrix C
-    using LayoutC = LayoutC_;
+  /// Layout of accumulator matrix C
+  using LayoutC = LayoutC_;
 
-    /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
-    using Policy = Policy_;
+  /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
+  using Policy = Policy_;
 
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    using ArchMmaOperator = typename Policy::Operator;
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  using ArchMmaOperator = typename Policy::Operator;
 
-    /// Indicates math operator
-    using MathOperator = typename ArchMmaOperator::Operator;
+  /// Indicates math operator
+  using MathOperator = typename ArchMmaOperator::Operator;
 
-    /// Architecture tag from underlying instruction
-    using ArchTag = typename ArchMmaOperator::ArchTag;
-    static_assert((platform::is_same<typename ArchMmaOperator::ElementA, half_t>::value
-                      && platform::is_same<typename ArchMmaOperator::ElementB, half_t>::value)
-            || (platform::is_same<typename ArchMmaOperator::ElementA, bfloat16_t>::value
-                && platform::is_same<typename ArchMmaOperator::ElementB, bfloat16_t>::value
-                && ArchTag::kMinComputeCapability >= 80),
-        "MmaTensorOpCvtBToA only supports underlying HMMA/QMMA");
+  /// Architecture tag from underlying instruction
+  using ArchTag = typename ArchMmaOperator::ArchTag;
+  static_assert(
+      (platform::is_same<typename ArchMmaOperator::ElementA, half_t>::value &&
+       platform::is_same<typename ArchMmaOperator::ElementB, half_t>::value) ||
+          (platform::is_same<typename ArchMmaOperator::ElementA,
+                             bfloat16_t>::value &&
+           platform::is_same<typename ArchMmaOperator::ElementB,
+                             bfloat16_t>::value &&
+           ArchTag::kMinComputeCapability >= 80),
+      "MmaTensorOpCvtBToA only supports underlying HMMA/QMMA");
 
-    static_assert(platform::is_same<ElementA, half_t>::value
-            || (platform::is_same<ElementA, bfloat16_t>::value && ArchTag::kMinComputeCapability >= 80),
-        "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on Ampere+");
+  static_assert(platform::is_same<ElementA, half_t>::value ||
+                    (platform::is_same<ElementA, bfloat16_t>::value &&
+                     ArchTag::kMinComputeCapability >= 80),
+                "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on Ampere+");
 
-    /// Indicates class of matrix operator
-    using OperatorClass = arch::OpClassTensorOp;
+  /// Indicates class of matrix operator
+  using OperatorClass = arch::OpClassTensorOp;
 
-    /// Shape of underlying instruction
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  /// Shape of underlying instruction
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    /// Instruction shape to override shared memory iterators with
-    using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
+  /// Instruction shape to override shared memory iterators with
+  using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
 
-    static_assert(
-        SharedMemoryInstructionShape::kM == InstructionShape::kM, "M dimension of compute instruction must match load");
-    static_assert(
-        SharedMemoryInstructionShape::kN == InstructionShape::kN, "N dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kM == InstructionShape::kM,
+                "M dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kN == InstructionShape::kN,
+                "N dimension of compute instruction must match load");
 
-    static constexpr int kExpansionFactor = SharedMemoryInstructionShape::kK / InstructionShape::kK;
+  static constexpr int kExpansionFactor =
+      SharedMemoryInstructionShape::kK / InstructionShape::kK;
 
-    static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
+  static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = ComplexTransform::kNone;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = ComplexTransform::kNone;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
 
-    /// Number of threads participating in warp-level matrix product
-    static int const kThreadCount = 32;
+  /// Number of threads participating in warp-level matrix product
+  static int const kThreadCount = 32;
 
-    /// Number of partitions along K dimension
-    static int const kPartitionsK = PartitionsK_;
+  /// Number of partitions along K dimension
+  static int const kPartitionsK = PartitionsK_;
 
-public:
-    /// Iterates over the A operand in memory
-    using IteratorA
-        = MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kM, Shape::kK>, Operand::kA, ElementA, LayoutA,
-            MatrixShape<InstructionShape::kM, InstructionShape::kK>, Policy::OpDelta::kRow, kThreadCount, kPartitionsK>;
+ public:
+  /// Iterates over the A operand in memory
+  using IteratorA = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kM, Shape::kK>,
+      Operand::kA,
+      ElementA,
+      LayoutA,
+      MatrixShape<InstructionShape::kM, InstructionShape::kK>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
 
-    /// Storage for A tile
-    using FragmentA = typename IteratorA::Fragment;
+  /// Storage for A tile
+  using FragmentA = typename IteratorA::Fragment;
 
-    /// Storage for transformed A tile
-    using TransformedFragmentA = Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
+  /// Storage for transformed A tile
+  using TransformedFragmentA =
+      Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
 
-    /// Iterates over the B operand in memory
-    using IteratorB = MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kK, Shape::kN>, Operand::kB, ElementB,
-        LayoutB, MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>, Policy::OpDelta::kRow,
-        kThreadCount, kPartitionsK>;
+  /// Iterates over the B operand in memory
+  using IteratorB = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kK, Shape::kN>,
+      Operand::kB,
+      ElementB,
+      LayoutB,
+      MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
 
-    /// Storage for B tile
-    using FragmentB = typename IteratorB::Fragment;
+  /// Storage for B tile
+  using FragmentB = typename IteratorB::Fragment;
 
-    /// Storage for transformed B tile
-    using TransformedFragmentB =
-        Array<typename ArchMmaOperator::ElementB, FragmentB::kElements / kExpansionFactor>;
+  /// Storage for transformed B tile
+  using TransformedFragmentB = Array<typename ArchMmaOperator::ElementB,
+                                     FragmentB::kElements / kExpansionFactor>;
 
-    /// Iterates over the C operand in memory
-    using IteratorC = MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>, ElementC, LayoutC,
-        typename ArchMmaOperator::Shape, typename Policy::OpDelta>;
+  /// Iterates over the C operand in memory
+  using IteratorC =
+      MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>,
+                                         ElementC,
+                                         LayoutC,
+                                         typename ArchMmaOperator::Shape,
+                                         typename Policy::OpDelta>;
 
-    /// Storage for C tile
-    using FragmentC = typename IteratorC::Fragment;
+  /// Storage for C tile
+  using FragmentC = typename IteratorC::Fragment;
 
-    /// Number of mma operations performed
-    using MmaIterations = MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM,
-        (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN>;
+  /// Number of mma operations performed
+  using MmaIterations =
+      MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) /
+                      ArchMmaOperator::Shape::kM,
+                  (Shape::kN + ArchMmaOperator::Shape::kN - 1) /
+                      ArchMmaOperator::Shape::kN>;
 
-public:
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    ArchMmaOperator mma;
+ public:
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  ArchMmaOperator mma;
 
-public:
-    //
-    // Methods
-    //
+ public:
+  //
+  // Methods
+  //
 
-    /// Ctor
-    CUTLASS_DEVICE
-    MmaTensorOpComputeBWithF16() {}
+  /// Ctor
+  CUTLASS_DEVICE
+  MmaTensorOpComputeBWithF16() {}
 
-    /// Performs a warp-level matrix multiply-accumulate operation
-    CUTLASS_DEVICE
-    void operator()(FragmentC& D, TransformedFragmentA const& A, TransformedFragmentB const& B, FragmentC const& C) const
-    {
+  /// Performs a warp-level matrix multiply-accumulate operation
+  CUTLASS_DEVICE
+  void operator()(FragmentC& D,
+                  TransformedFragmentA const& A,
+                  TransformedFragmentB const& B,
+                  FragmentC const& C) const {
+    using MmaOperandA = typename ArchMmaOperator::FragmentA;
+    using MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using MmaOperandC = typename ArchMmaOperator::FragmentC;
 
-        using MmaOperandA = typename ArchMmaOperator::FragmentA;
-        using MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using MmaOperandC = typename ArchMmaOperator::FragmentC;
+    D = C;
 
-        D = C;
-
-        MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
-        MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
-        MmaOperandC* ptr_D = reinterpret_cast<MmaOperandC*>(&D);
+    MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
+    MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
+    MmaOperandC* ptr_D = reinterpret_cast<MmaOperandC*>(&D);
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
-        // Serpentine visitation order maximizing reuse of Rb
-        CUTLASS_PRAGMA_UNROLL
-        for (int n = 0; n < MmaIterations::kColumn; ++n)
-        {
+    // Serpentine visitation order maximizing reuse of Rb
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < MmaIterations::kColumn; ++n) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int m = 0; m < MmaIterations::kRow; ++m) {
+        int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
 
-            CUTLASS_PRAGMA_UNROLL
-            for (int m = 0; m < MmaIterations::kRow; ++m)
-            {
-
-                int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
-
-                if (AccumulatorsInRowMajor)
-                { // matrix B is reordered
-                    mma(ptr_D[n + m_serpentine * MmaIterations::kColumn], ptr_A[m_serpentine], ptr_B[n],
-                        ptr_D[n + m_serpentine * MmaIterations::kColumn]);
-                }
-                else
-                {
-                    mma(ptr_D[m_serpentine + n * MmaIterations::kRow], ptr_A[m_serpentine], ptr_B[n],
-                        ptr_D[m_serpentine + n * MmaIterations::kRow]);
-                }
-            }
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n + m_serpentine * MmaIterations::kColumn],
+              ptr_A[m_serpentine],
+              ptr_B[n],
+              ptr_D[n + m_serpentine * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m_serpentine + n * MmaIterations::kRow],
+              ptr_A[m_serpentine],
+              ptr_B[n],
+              ptr_D[m_serpentine + n * MmaIterations::kRow]);
         }
-#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
-        // Serpentine visitation order maximizing reuse of Ra
-        CUTLASS_PRAGMA_UNROLL
-        for (int m = 0; m < MmaIterations::kRow; ++m)
-        {
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int n = 0; n < MmaIterations::kColumn; ++n)
-            {
-
-                int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
-
-                if (AccumulatorsInRowMajor)
-                { // matrix B is reordered
-                    mma(ptr_D[n_serpentine + m * MmaIterations::kColumn], ptr_A[m], ptr_B[n_serpentine],
-                        ptr_D[n_serpentine + m * MmaIterations::kColumn]);
-                }
-                else
-                {
-                    mma(ptr_D[m + n_serpentine * MmaIterations::kRow], ptr_A[m], ptr_B[n_serpentine],
-                        ptr_D[m + n_serpentine * MmaIterations::kRow]);
-                }
-            }
-        }
-#else
-        assert(0);
-#endif
+      }
     }
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+    // Serpentine visitation order maximizing reuse of Ra
+    CUTLASS_PRAGMA_UNROLL
+    for (int m = 0; m < MmaIterations::kRow; ++m) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < MmaIterations::kColumn; ++n) {
+        int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
+
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n_serpentine + m * MmaIterations::kColumn],
+              ptr_A[m],
+              ptr_B[n_serpentine],
+              ptr_D[n_serpentine + m * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m + n_serpentine * MmaIterations::kRow],
+              ptr_A[m],
+              ptr_B[n_serpentine],
+              ptr_D[m + n_serpentine * MmaIterations::kRow]);
+        }
+      }
+    }
+#else
+    assert(0);
+#endif
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace warp
-} // namespace gemm
-} // namespace cutlass
+}  // namespace warp
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h
index 24e844abca..5bb06a5414 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,18 +18,20 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
-  \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores.
+  \brief Defines iterators used by warp-level matrix multiply operations
+  targeting Tensor Cores.
 */
 
 #pragma once
@@ -57,12 +59,9 @@
 
 ////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace gemm
-{
-namespace warp
-{
+namespace cutlass {
+namespace gemm {
+namespace warp {
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -94,193 +93,216 @@ template <
     typename Shape_,
     ///
     WeightOnlyQuantOp QuantOp_>
-class MmaTensorOpDequantizer<MmaOperator_, Shape_, Operand::kB, bfloat16_t, layout::RowMajor, 32, QuantOp_,
-    typename platform::enable_if<MmaOperator_::ArchTag::kMinComputeCapability >= 80
-        && platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB, layout::ColumnMajor>::value>::type>
-{
+class MmaTensorOpDequantizer<
+    MmaOperator_,
+    Shape_,
+    Operand::kB,
+    bfloat16_t,
+    layout::RowMajor,
+    32,
+    QuantOp_,
+    typename platform::enable_if<
+        MmaOperator_::ArchTag::kMinComputeCapability >= 80 &&
+        platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB,
+                          layout::ColumnMajor>::value>::type> {
+ public:
+  /// Mma Operator
+  using MmaOperator = MmaOperator_;
 
-public:
-    /// Mma Operator
-    using MmaOperator = MmaOperator_;
+  // The architecture specific mma ooperator being used
+  using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
 
-    // The architecture specific mma ooperator being used
-    using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
+  // Mma Instruction Shape
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    // Mma Instruction Shape
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  // This is the ratio of the load instruction vs the compute instruction.
+  static constexpr int kExpansionFactor =
+      MmaOperator::IteratorB::InstructionShape::kRow / InstructionShape::kK;
 
-    // This is the ratio of the load instruction vs the compute instruction.
-    static constexpr int kExpansionFactor = MmaOperator::IteratorB::InstructionShape::kRow / InstructionShape::kK;
+  /// Type of the scales
+  using ElementScale = bfloat16_t;
 
-    /// Type of the scales
-    using ElementScale = bfloat16_t;
+  /// Fragment to hold B data before Mma
+  using FragmentDequantizedOperand =
+      Array<ElementScale, MmaOperator::FragmentB::kElements>;
 
-    /// Fragment to hold B data before Mma
-    using FragmentDequantizedOperand = Array<ElementScale, MmaOperator::FragmentB::kElements>;
+  // Fragment to hold scale data to apply to B before mma
+  // We need 1 fp16 per matrix iteration in the N dimension
+  static constexpr int kColsPerMmaPerThread = 1;
+  using FragmentScale =
+      Array<ElementScale,
+            kColsPerMmaPerThread * MmaOperator::MmaIterations::kColumn>;
+  using FragmentZero =
+      Array<ElementScale,
+            kColsPerMmaPerThread * MmaOperator::MmaIterations::kColumn>;
 
-    // Fragment to hold scale data to apply to B before mma
-    // We need 1 fp16 per matrix iteration in the N dimension
-    static constexpr int kColsPerMmaPerThread = 1;
-    using FragmentScale = Array<ElementScale, kColsPerMmaPerThread * MmaOperator::MmaIterations::kColumn>;
-    using FragmentZero = Array<ElementScale, kColsPerMmaPerThread * MmaOperator::MmaIterations::kColumn>;
+  /// Warp mma shape
+  using Shape = Shape_;
 
-    /// Warp mma shape
-    using Shape = Shape_;
+  /// Layout of the scales in shared memory
+  using Layout = layout::RowMajor;
 
-    /// Layout of the scales in shared memory
-    using Layout = layout::RowMajor;
+  /// TensorRef type for loading element from a tensor
+  using TensorRef = TensorRef<ElementScale, Layout>;
 
-    /// TensorRef type for loading element from a tensor
-    using TensorRef = TensorRef<ElementScale, Layout>;
+  static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
 
-    static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
-
-    CUTLASS_DEVICE
-    MmaTensorOpDequantizer(TensorRef smem_scales, TensorRef smem_zeros, int const warp_idx_n, int const lane_idx)
-    {
-        int const warp_offset = warp_idx_n * Shape::kN;
-        int const quad = lane_idx / 4;
-        int const thread_offset = warp_offset + quad;
-        pointer_scale_ = smem_scales.data() + thread_offset;
-        if constexpr (hasZero(QuantOp))
-        {
-            pointer_zero_ = smem_zeros.data() + thread_offset;
-        }
+  CUTLASS_DEVICE
+  MmaTensorOpDequantizer(TensorRef smem_scales,
+                         TensorRef smem_zeros,
+                         int const warp_idx_n,
+                         int const lane_idx) {
+    int const warp_offset = warp_idx_n * Shape::kN;
+    int const quad = lane_idx / 4;
+    int const thread_offset = warp_offset + quad;
+    pointer_scale_ = smem_scales.data() + thread_offset;
+    if constexpr (hasZero(QuantOp)) {
+      pointer_zero_ = smem_zeros.data() + thread_offset;
     }
+  }
 
-    CUTLASS_DEVICE
-    MmaTensorOpDequantizer(TensorRef smem_scales, int const warp_idx_n, int const lane_idx)
-        : MmaTensorOpDequantizer(smem_scales, TensorRef(), warp_idx_n, lane_idx)
-    {
+  CUTLASS_DEVICE
+  MmaTensorOpDequantizer(TensorRef smem_scales,
+                         int const warp_idx_n,
+                         int const lane_idx)
+      : MmaTensorOpDequantizer(smem_scales, TensorRef(), warp_idx_n, lane_idx) {
+  }
+
+  CUTLASS_DEVICE
+  void load(FragmentScale& scale_frag) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+         ++mma_n_iter) {
+      scale_frag[mma_n_iter] =
+          pointer_scale_[mma_n_iter * InstructionShape::kN];
     }
+  }
 
-    CUTLASS_DEVICE
-    void load(FragmentScale& scale_frag)
-    {
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-        {
-            scale_frag[mma_n_iter] = pointer_scale_[mma_n_iter * InstructionShape::kN];
-        }
-    }
-
-    CUTLASS_DEVICE
-    void dequantize(FragmentDequantizedOperand& operand_frag, FragmentScale const& scale_frag)
-    {
+  CUTLASS_DEVICE
+  void dequantize(FragmentDequantizedOperand& operand_frag,
+                  FragmentScale const& scale_frag) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) && defined(ENABLE_BF16))
-        using _MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using ExpandedMmaOperandB = Array<typename _MmaOperandB::Element, kExpansionFactor * _MmaOperandB::kElements>;
-        static_assert(ExpandedMmaOperandB::kElements * MmaOperator::MmaIterations::kColumn
-                == FragmentDequantizedOperand::kElements,
-            "");
+    using _MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using ExpandedMmaOperandB =
+        Array<typename _MmaOperandB::Element,
+              kExpansionFactor * _MmaOperandB::kElements>;
+    static_assert(
+        ExpandedMmaOperandB::kElements * MmaOperator::MmaIterations::kColumn ==
+            FragmentDequantizedOperand::kElements,
+        "");
 
-        __nv_bfloat16 const* scale_ptr = reinterpret_cast<__nv_bfloat16 const*>(&scale_frag);
-        ExpandedMmaOperandB* operand_frag_ptr = reinterpret_cast<ExpandedMmaOperandB*>(&operand_frag);
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-        {
-            static_assert(ExpandedMmaOperandB::kElements % 2 == 0, "");
+    __nv_bfloat16 const* scale_ptr =
+        reinterpret_cast<__nv_bfloat16 const*>(&scale_frag);
+    ExpandedMmaOperandB* operand_frag_ptr =
+        reinterpret_cast<ExpandedMmaOperandB*>(&operand_frag);
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+         ++mma_n_iter) {
+      static_assert(ExpandedMmaOperandB::kElements % 2 == 0, "");
 
-            __nv_bfloat162 scalex2 = __bfloat162bfloat162(scale_ptr[mma_n_iter]);
-            __nv_bfloat162* operand_bf16x2_ptr = reinterpret_cast<__nv_bfloat162*>(&operand_frag_ptr[mma_n_iter]);
+      __nv_bfloat162 scalex2 = __bfloat162bfloat162(scale_ptr[mma_n_iter]);
+      __nv_bfloat162* operand_bf16x2_ptr =
+          reinterpret_cast<__nv_bfloat162*>(&operand_frag_ptr[mma_n_iter]);
 
-            CUTLASS_PRAGMA_UNROLL
-            for (int ii = 0; ii < ExpandedMmaOperandB::kElements / 2; ++ii)
-            {
-                operand_bf16x2_ptr[ii] = __hmul2(operand_bf16x2_ptr[ii], scalex2);
-            }
-        }
+      CUTLASS_PRAGMA_UNROLL
+      for (int ii = 0; ii < ExpandedMmaOperandB::kElements / 2; ++ii) {
+        operand_bf16x2_ptr[ii] = __hmul2(operand_bf16x2_ptr[ii], scalex2);
+      }
+    }
 #else
-        // Slow path not implemented here on purpose. If we need to do HMMA on older arch, scale conversion should
-        // happen before scales are stored to shared memory and we should use the fp16 dequantizer. This will avoid
-        // numerous conversion instructions in GEMM main loop.
-        arch::device_breakpoint();
+    // Slow path not implemented here on purpose. If we need to do HMMA on older
+    // arch, scale conversion should happen before scales are stored to shared
+    // memory and we should use the fp16 dequantizer. This will avoid numerous
+    // conversion instructions in GEMM main loop.
+    arch::device_breakpoint();
 #endif
-    }
+  }
 
-    CUTLASS_DEVICE
-    void load(FragmentScale& scale_frag, FragmentScale& zero_frag)
-    {
-        if constexpr (hasZero(QuantOp))
-        {
-            CUTLASS_PRAGMA_UNROLL
-            for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-            {
-                scale_frag[mma_n_iter] = pointer_scale_[mma_n_iter * InstructionShape::kN];
-                zero_frag[mma_n_iter] = pointer_zero_[mma_n_iter * InstructionShape::kN];
-            }
-        }
-        else
-        {
-            CUTLASS_PRAGMA_UNROLL
-            for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-            {
-                scale_frag[mma_n_iter] = pointer_scale_[mma_n_iter * InstructionShape::kN];
-            }
-        }
+  CUTLASS_DEVICE
+  void load(FragmentScale& scale_frag, FragmentScale& zero_frag) {
+    if constexpr (hasZero(QuantOp)) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+           ++mma_n_iter) {
+        scale_frag[mma_n_iter] =
+            pointer_scale_[mma_n_iter * InstructionShape::kN];
+        zero_frag[mma_n_iter] =
+            pointer_zero_[mma_n_iter * InstructionShape::kN];
+      }
+    } else {
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+           ++mma_n_iter) {
+        scale_frag[mma_n_iter] =
+            pointer_scale_[mma_n_iter * InstructionShape::kN];
+      }
     }
+  }
 
-    CUTLASS_DEVICE
-    void dequantize(
-        FragmentDequantizedOperand& operand_frag, FragmentScale const& scale_frag, FragmentScale const& zero_frag)
-    {
+  CUTLASS_DEVICE
+  void dequantize(FragmentDequantizedOperand& operand_frag,
+                  FragmentScale const& scale_frag,
+                  FragmentScale const& zero_frag) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) && defined(ENABLE_BF16))
-        using _MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using ExpandedMmaOperandB = Array<typename _MmaOperandB::Element, kExpansionFactor * _MmaOperandB::kElements>;
-        static_assert(ExpandedMmaOperandB::kElements * MmaOperator::MmaIterations::kColumn
-                == FragmentDequantizedOperand::kElements,
-            "");
+    using _MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using ExpandedMmaOperandB =
+        Array<typename _MmaOperandB::Element,
+              kExpansionFactor * _MmaOperandB::kElements>;
+    static_assert(
+        ExpandedMmaOperandB::kElements * MmaOperator::MmaIterations::kColumn ==
+            FragmentDequantizedOperand::kElements,
+        "");
 
-        __nv_bfloat16 const* scale_ptr = reinterpret_cast<__nv_bfloat16 const*>(&scale_frag);
-        __nv_bfloat16 const* zero_ptr = reinterpret_cast<__nv_bfloat16 const*>(&zero_frag);
+    __nv_bfloat16 const* scale_ptr =
+        reinterpret_cast<__nv_bfloat16 const*>(&scale_frag);
+    __nv_bfloat16 const* zero_ptr =
+        reinterpret_cast<__nv_bfloat16 const*>(&zero_frag);
 
-        ExpandedMmaOperandB* operand_frag_ptr = reinterpret_cast<ExpandedMmaOperandB*>(&operand_frag);
+    ExpandedMmaOperandB* operand_frag_ptr =
+        reinterpret_cast<ExpandedMmaOperandB*>(&operand_frag);
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+         ++mma_n_iter) {
+      static_assert(ExpandedMmaOperandB::kElements % 2 == 0, "");
+
+      __nv_bfloat162 scalex2 = __bfloat162bfloat162(scale_ptr[mma_n_iter]);
+      __nv_bfloat162 zerox2 = __bfloat162bfloat162(zero_ptr[mma_n_iter]);
+      __nv_bfloat162* operand_bf16x2_ptr =
+          reinterpret_cast<__nv_bfloat162*>(&operand_frag_ptr[mma_n_iter]);
+
+      if constexpr (hasZero(QuantOp)) {
         CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-        {
-            static_assert(ExpandedMmaOperandB::kElements % 2 == 0, "");
-
-            __nv_bfloat162 scalex2 = __bfloat162bfloat162(scale_ptr[mma_n_iter]);
-            __nv_bfloat162 zerox2 = __bfloat162bfloat162(zero_ptr[mma_n_iter]);
-            __nv_bfloat162* operand_bf16x2_ptr = reinterpret_cast<__nv_bfloat162*>(&operand_frag_ptr[mma_n_iter]);
-
-            if constexpr (hasZero(QuantOp))
-            {
-                CUTLASS_PRAGMA_UNROLL
-                for (int ii = 0; ii < ExpandedMmaOperandB::kElements / 2; ++ii)
-                {
-                    operand_bf16x2_ptr[ii] = __hfma2(operand_bf16x2_ptr[ii], scalex2, zerox2);
-                }
-            }
-            else
-            {
-                CUTLASS_PRAGMA_UNROLL
-                for (int ii = 0; ii < ExpandedMmaOperandB::kElements / 2; ++ii)
-                {
-                    operand_bf16x2_ptr[ii] = __hmul2(operand_bf16x2_ptr[ii], scalex2);
-                }
-            }
+        for (int ii = 0; ii < ExpandedMmaOperandB::kElements / 2; ++ii) {
+          operand_bf16x2_ptr[ii] =
+              __hfma2(operand_bf16x2_ptr[ii], scalex2, zerox2);
         }
+      } else {
+        CUTLASS_PRAGMA_UNROLL
+        for (int ii = 0; ii < ExpandedMmaOperandB::kElements / 2; ++ii) {
+          operand_bf16x2_ptr[ii] = __hmul2(operand_bf16x2_ptr[ii], scalex2);
+        }
+      }
+    }
 #else
-        // Slow path not implemented here on purpose. If we need to do HMMA on older arch, scale conversion should
-        // happen before scales are stored to shared memory and we should use the fp16 dequantizer. This will avoid
-        // numerous conversion instructions in GEMM main loop.
-        arch::device_breakpoint();
+    // Slow path not implemented here on purpose. If we need to do HMMA on older
+    // arch, scale conversion should happen before scales are stored to shared
+    // memory and we should use the fp16 dequantizer. This will avoid numerous
+    // conversion instructions in GEMM main loop.
+    arch::device_breakpoint();
 #endif
-    }
+  }
 
-    // Adds a pointer offset in units of elements.
-    CUTLASS_DEVICE
-    void add_pointer_offset(int64_t const& offset)
-    {
-        static_assert(sizeof(ElementScale) > 1, "");
-        pointer_scale_ += offset;
-        pointer_zero_ += offset;
-    }
+  // Adds a pointer offset in units of elements.
+  CUTLASS_DEVICE
+  void add_pointer_offset(int64_t const& offset) {
+    static_assert(sizeof(ElementScale) > 1, "");
+    pointer_scale_ += offset;
+    pointer_zero_ += offset;
+  }
 
-private:
-    ElementScale const* pointer_scale_;
-    ElementScale const* pointer_zero_;
+ private:
+  ElementScale const* pointer_scale_;
+  ElementScale const* pointer_zero_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -293,170 +315,190 @@ template <
     typename Shape_,
     ///
     WeightOnlyQuantOp QuantOp_>
-class MmaTensorOpDequantizer<MmaOperator_, Shape_, Operand::kB, half_t, layout::RowMajor, 32, QuantOp_,
-    typename platform::enable_if<MmaOperator_::ArchTag::kMinComputeCapability >= 75
-        && platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB, layout::ColumnMajor>::value>::type>
-{
+class MmaTensorOpDequantizer<
+    MmaOperator_,
+    Shape_,
+    Operand::kB,
+    half_t,
+    layout::RowMajor,
+    32,
+    QuantOp_,
+    typename platform::enable_if<
+        MmaOperator_::ArchTag::kMinComputeCapability >= 75 &&
+        platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB,
+                          layout::ColumnMajor>::value>::type> {
+ public:
+  /// Mma Operator
+  using MmaOperator = MmaOperator_;
 
-public:
-    /// Mma Operator
-    using MmaOperator = MmaOperator_;
+  // The architecture specific mma ooperator being used
+  using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
 
-    // The architecture specific mma ooperator being used
-    using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
+  // Mma Instruction Shape
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    // Mma Instruction Shape
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  // This is the ratio of the load instruction vs the compute instruction.
+  static constexpr int kExpansionFactor =
+      MmaOperator::IteratorB::InstructionShape::kRow / InstructionShape::kK;
 
-    // This is the ratio of the load instruction vs the compute instruction.
-    static constexpr int kExpansionFactor = MmaOperator::IteratorB::InstructionShape::kRow / InstructionShape::kK;
+  /// Type of the scales
+  using ElementScale = half_t;
 
-    /// Type of the scales
-    using ElementScale = half_t;
+  /// Fragment to hold B data before Mma
+  using FragmentDequantizedOperand =
+      Array<ElementScale, MmaOperator::FragmentB::kElements>;
 
-    /// Fragment to hold B data before Mma
-    using FragmentDequantizedOperand = Array<ElementScale, MmaOperator::FragmentB::kElements>;
+  // Fragment to hold scale data to apply to B before mma
+  // We need 1 fp16 per matrix iteration in the N dimension
+  static constexpr int kColsPerMmaPerThread = 1;
+  using FragmentScale =
+      Array<ElementScale,
+            kColsPerMmaPerThread * MmaOperator::MmaIterations::kColumn>;
+  using FragmentZero =
+      Array<ElementScale,
+            kColsPerMmaPerThread * MmaOperator::MmaIterations::kColumn>;
 
-    // Fragment to hold scale data to apply to B before mma
-    // We need 1 fp16 per matrix iteration in the N dimension
-    static constexpr int kColsPerMmaPerThread = 1;
-    using FragmentScale = Array<ElementScale, kColsPerMmaPerThread * MmaOperator::MmaIterations::kColumn>;
-    using FragmentZero = Array<ElementScale, kColsPerMmaPerThread * MmaOperator::MmaIterations::kColumn>;
+  /// Warp mma shape
+  using Shape = Shape_;
 
-    /// Warp mma shape
-    using Shape = Shape_;
+  /// Layout of the scales in shared memory
+  using Layout = layout::RowMajor;
 
-    /// Layout of the scales in shared memory
-    using Layout = layout::RowMajor;
+  /// TensorRef type for loading element from a tensor
+  using TensorRef = TensorRef<ElementScale, Layout>;
 
-    /// TensorRef type for loading element from a tensor
-    using TensorRef = TensorRef<ElementScale, Layout>;
+  static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
 
-    static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
-
-    CUTLASS_DEVICE
-    MmaTensorOpDequantizer(TensorRef smem_scales, TensorRef smem_zeros, int const warp_idx_n, int const lane_idx)
-    {
-        int const warp_offset = warp_idx_n * Shape::kN;
-        int const quad = lane_idx / 4;
-        int const thread_offset = warp_offset + quad;
-        pointer_scale_ = smem_scales.data() + thread_offset;
-        if constexpr (hasZero(QuantOp))
-        {
-            pointer_zero_ = smem_zeros.data() + thread_offset;
-        }
+  CUTLASS_DEVICE
+  MmaTensorOpDequantizer(TensorRef smem_scales,
+                         TensorRef smem_zeros,
+                         int const warp_idx_n,
+                         int const lane_idx) {
+    int const warp_offset = warp_idx_n * Shape::kN;
+    int const quad = lane_idx / 4;
+    int const thread_offset = warp_offset + quad;
+    pointer_scale_ = smem_scales.data() + thread_offset;
+    if constexpr (hasZero(QuantOp)) {
+      pointer_zero_ = smem_zeros.data() + thread_offset;
     }
+  }
 
-    CUTLASS_DEVICE
-    MmaTensorOpDequantizer(TensorRef smem_scales, int const warp_idx_n, int const lane_idx)
-        : MmaTensorOpDequantizer(smem_scales, TensorRef(), warp_idx_n, lane_idx)
-    {
+  CUTLASS_DEVICE
+  MmaTensorOpDequantizer(TensorRef smem_scales,
+                         int const warp_idx_n,
+                         int const lane_idx)
+      : MmaTensorOpDequantizer(smem_scales, TensorRef(), warp_idx_n, lane_idx) {
+  }
+
+  CUTLASS_DEVICE
+  void load(FragmentScale& scale_frag) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+         ++mma_n_iter) {
+      scale_frag[mma_n_iter] =
+          pointer_scale_[mma_n_iter * InstructionShape::kN];
     }
+  }
 
-    CUTLASS_DEVICE
-    void load(FragmentScale& scale_frag)
-    {
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-        {
-            scale_frag[mma_n_iter] = pointer_scale_[mma_n_iter * InstructionShape::kN];
-        }
+  CUTLASS_DEVICE
+  void dequantize(FragmentDequantizedOperand& operand_frag,
+                  FragmentScale const& scale_frag) {
+    using _MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using ExpandedMmaOperandB =
+        Array<typename FragmentDequantizedOperand::Element,
+              kExpansionFactor * _MmaOperandB::kElements>;
+    static_assert(
+        ExpandedMmaOperandB::kElements * MmaOperator::MmaIterations::kColumn ==
+            FragmentDequantizedOperand::kElements,
+        "");
+
+    multiplies<ExpandedMmaOperandB> mul_op;
+
+    ExpandedMmaOperandB* operand_frag_ptr =
+        reinterpret_cast<ExpandedMmaOperandB*>(&operand_frag);
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+         ++mma_n_iter) {
+      operand_frag_ptr[mma_n_iter] =
+          mul_op(operand_frag_ptr[mma_n_iter], scale_frag[mma_n_iter]);
     }
+  }
 
-    CUTLASS_DEVICE
-    void dequantize(FragmentDequantizedOperand& operand_frag, FragmentScale const& scale_frag)
-    {
-        using _MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using ExpandedMmaOperandB
-            = Array<typename FragmentDequantizedOperand::Element, kExpansionFactor * _MmaOperandB::kElements>;
-        static_assert(ExpandedMmaOperandB::kElements * MmaOperator::MmaIterations::kColumn
-                == FragmentDequantizedOperand::kElements,
-            "");
-
-        multiplies<ExpandedMmaOperandB> mul_op;
-
-        ExpandedMmaOperandB* operand_frag_ptr = reinterpret_cast<ExpandedMmaOperandB*>(&operand_frag);
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-        {
-            operand_frag_ptr[mma_n_iter] = mul_op(operand_frag_ptr[mma_n_iter], scale_frag[mma_n_iter]);
-        }
+  CUTLASS_DEVICE
+  void load(FragmentScale& scale_frag, FragmentScale& zero_frag) {
+    if constexpr (hasZero(QuantOp)) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+           ++mma_n_iter) {
+        scale_frag[mma_n_iter] =
+            pointer_scale_[mma_n_iter * InstructionShape::kN];
+        zero_frag[mma_n_iter] =
+            pointer_zero_[mma_n_iter * InstructionShape::kN];
+      }
+    } else {
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+           ++mma_n_iter) {
+        scale_frag[mma_n_iter] =
+            pointer_scale_[mma_n_iter * InstructionShape::kN];
+      }
     }
+  }
 
-    CUTLASS_DEVICE
-    void load(FragmentScale& scale_frag, FragmentScale& zero_frag)
-    {
-        if constexpr (hasZero(QuantOp))
-        {
-            CUTLASS_PRAGMA_UNROLL
-            for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-            {
-                scale_frag[mma_n_iter] = pointer_scale_[mma_n_iter * InstructionShape::kN];
-                zero_frag[mma_n_iter] = pointer_zero_[mma_n_iter * InstructionShape::kN];
-            }
-        }
-        else
-        {
-            CUTLASS_PRAGMA_UNROLL
-            for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-            {
-                scale_frag[mma_n_iter] = pointer_scale_[mma_n_iter * InstructionShape::kN];
-            }
-        }
+  CUTLASS_DEVICE
+  void dequantize(FragmentDequantizedOperand& operand_frag,
+                  FragmentScale const& scale_frag,
+                  FragmentScale const& zero_frag) {
+    using _MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using ExpandedMmaOperandB =
+        Array<typename FragmentDequantizedOperand::Element,
+              kExpansionFactor * _MmaOperandB::kElements>;
+    static_assert(
+        ExpandedMmaOperandB::kElements * MmaOperator::MmaIterations::kColumn ==
+            FragmentDequantizedOperand::kElements,
+        "");
+
+    multiplies<ExpandedMmaOperandB> mul_op;
+    ExpandedMmaOperandB* operand_frag_ptr =
+        reinterpret_cast<ExpandedMmaOperandB*>(&operand_frag);
+
+    if constexpr (hasZero(QuantOp)) {
+      plus<ExpandedMmaOperandB> plus_op;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+           ++mma_n_iter) {
+        operand_frag_ptr[mma_n_iter] = plus_op(
+            mul_op(operand_frag_ptr[mma_n_iter], scale_frag[mma_n_iter]),
+            zero_frag[mma_n_iter]);
+      }
+    } else {
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
+           ++mma_n_iter) {
+        operand_frag_ptr[mma_n_iter] =
+            mul_op(operand_frag_ptr[mma_n_iter], scale_frag[mma_n_iter]);
+      }
     }
+  }
 
-    CUTLASS_DEVICE
-    void dequantize(
-        FragmentDequantizedOperand& operand_frag, FragmentScale const& scale_frag, FragmentScale const& zero_frag)
-    {
-        using _MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using ExpandedMmaOperandB
-            = Array<typename FragmentDequantizedOperand::Element, kExpansionFactor * _MmaOperandB::kElements>;
-        static_assert(ExpandedMmaOperandB::kElements * MmaOperator::MmaIterations::kColumn
-                == FragmentDequantizedOperand::kElements,
-            "");
+  // Adds a pointer offset in units of elements.
+  CUTLASS_DEVICE
+  void add_pointer_offset(int64_t const& offset) {
+    static_assert(sizeof(ElementScale) > 1, "");
+    pointer_scale_ += offset;
+    pointer_zero_ += offset;
+  }
 
-        multiplies<ExpandedMmaOperandB> mul_op;
-        ExpandedMmaOperandB* operand_frag_ptr = reinterpret_cast<ExpandedMmaOperandB*>(&operand_frag);
-
-        if constexpr (hasZero(QuantOp))
-        {
-            plus<ExpandedMmaOperandB> plus_op;
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-            {
-                operand_frag_ptr[mma_n_iter]
-                    = plus_op(mul_op(operand_frag_ptr[mma_n_iter], scale_frag[mma_n_iter]), zero_frag[mma_n_iter]);
-            }
-        }
-        else
-        {
-            CUTLASS_PRAGMA_UNROLL
-            for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn; ++mma_n_iter)
-            {
-                operand_frag_ptr[mma_n_iter] = mul_op(operand_frag_ptr[mma_n_iter], scale_frag[mma_n_iter]);
-            }
-        }
-    }
-
-    // Adds a pointer offset in units of elements.
-    CUTLASS_DEVICE
-    void add_pointer_offset(int64_t const& offset)
-    {
-        static_assert(sizeof(ElementScale) > 1, "");
-        pointer_scale_ += offset;
-        pointer_zero_ += offset;
-    }
-
-private:
-    ElementScale const* pointer_scale_;
-    ElementScale const* pointer_zero_;
+ private:
+  ElementScale const* pointer_scale_;
+  ElementScale const* pointer_zero_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-// Specialization for Volta A x RowMajor B tensorOp, for 32x32x4 interleaved gemm
+// Specialization for Volta A x RowMajor B tensorOp, for 32x32x4 interleaved
+// gemm
 template <
     /// Underlying matrix multiply operator (concept: MmaTensorOp)
     typename MmaOperator_,
@@ -464,86 +506,98 @@ template <
     typename Shape_,
     ///
     WeightOnlyQuantOp QuantOp_>
-class MmaTensorOpDequantizer<MmaOperator_, Shape_, Operand::kB, half_t, layout::RowMajor, 32, QuantOp_,
-    typename platform::enable_if<platform::is_same<typename MmaOperator_::ArchTag, arch::Sm70>::value
-        && platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB, layout::RowMajor>::value>::type>
-{
+class MmaTensorOpDequantizer<
+    MmaOperator_,
+    Shape_,
+    Operand::kB,
+    half_t,
+    layout::RowMajor,
+    32,
+    QuantOp_,
+    typename platform::enable_if<
+        platform::is_same<typename MmaOperator_::ArchTag, arch::Sm70>::value &&
+        platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB,
+                          layout::RowMajor>::value>::type> {
+ public:
+  static_assert(platform::is_same<typename MmaOperator_::InterleavedTileShape,
+                                  GemmShape<32, 32, 4>>::value,
+                "");
 
-public:
-    static_assert(platform::is_same<typename MmaOperator_::InterleavedTileShape, GemmShape<32, 32, 4>>::value, "");
+  /// Mma Operator
+  using MmaOperator = MmaOperator_;
 
-    /// Mma Operator
-    using MmaOperator = MmaOperator_;
+  // The architecture specific mma ooperator being used
+  using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
 
-    // The architecture specific mma ooperator being used
-    using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
+  // Mma Instruction Shape
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    // Mma Instruction Shape
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  /// Type of the scales
+  using ElementScale = half_t;
 
-    /// Type of the scales
-    using ElementScale = half_t;
+  /// Fragment to hold B data before Mma
+  using FragmentDequantizedOperand =
+      Array<ElementScale, MmaOperator::FragmentB::kElements>;
 
-    /// Fragment to hold B data before Mma
-    using FragmentDequantizedOperand = Array<ElementScale, MmaOperator::FragmentB::kElements>;
+  /// Warp mma shape
+  using Shape = Shape_;
 
-    /// Warp mma shape
-    using Shape = Shape_;
+  // Fragment to hold scale data to apply to B before mma
+  // Each 32x32x4 matmul uses 8 elements from B.
+  static constexpr int ColsPerMmaTile = 32;
+  static constexpr int TileNIterations = Shape::kN / ColsPerMmaTile;
+  using FragmentScale = Array<ElementScale, TileNIterations * 8>;
+  using AccessType = Array<ElementScale, 8>;
 
-    // Fragment to hold scale data to apply to B before mma
-    // Each 32x32x4 matmul uses 8 elements from B.
-    static constexpr int ColsPerMmaTile = 32;
-    static constexpr int TileNIterations = Shape::kN / ColsPerMmaTile;
-    using FragmentScale = Array<ElementScale, TileNIterations * 8>;
-    using AccessType = Array<ElementScale, 8>;
+  /// Layout of the scales in shared memory
+  using Layout = layout::RowMajor;
 
-    /// Layout of the scales in shared memory
-    using Layout = layout::RowMajor;
+  /// TensorRef type for loading element from a tensor
+  using TensorRef = TensorRef<ElementScale, Layout>;
 
-    /// TensorRef type for loading element from a tensor
-    using TensorRef = TensorRef<ElementScale, Layout>;
+  static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
+  static_assert(QuantOp == WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY, "");
 
-    static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
-    static_assert(QuantOp == WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY, "");
+  CUTLASS_DEVICE
+  MmaTensorOpDequantizer(TensorRef smem_scales,
+                         int const warp_idx_n,
+                         int const lane_idx) {
+    int const warp_offset = warp_idx_n * Shape::kN;
+    int const base_col = lane_idx & 0xF8;
+    int const thread_offset = warp_offset + base_col;
+    pointer_ = smem_scales.data() + thread_offset;
+  }
 
-    CUTLASS_DEVICE
-    MmaTensorOpDequantizer(TensorRef smem_scales, int const warp_idx_n, int const lane_idx)
-    {
-        int const warp_offset = warp_idx_n * Shape::kN;
-        int const base_col = lane_idx & 0xF8;
-        int const thread_offset = warp_offset + base_col;
-        pointer_ = smem_scales.data() + thread_offset;
+  CUTLASS_DEVICE
+  void load(FragmentScale& scale_frag) {
+    AccessType* scale_frag_ptr = reinterpret_cast<AccessType*>(&scale_frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int tile_iter = 0; tile_iter < TileNIterations; ++tile_iter) {
+      // We jump by 32 here since volta does <32x32x4> super mmas inside a warp.
+      scale_frag_ptr[tile_iter] = *reinterpret_cast<AccessType const*>(
+          pointer_ + ColsPerMmaTile * tile_iter);
     }
+  }
 
-    CUTLASS_DEVICE
-    void load(FragmentScale& scale_frag)
-    {
-        AccessType* scale_frag_ptr = reinterpret_cast<AccessType*>(&scale_frag);
+  CUTLASS_DEVICE
+  void dequantize(FragmentDequantizedOperand& operand_frag,
+                  FragmentScale const& scale_frag) {
+    static_assert(
+        FragmentScale::kElements == FragmentDequantizedOperand::kElements, "");
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int tile_iter = 0; tile_iter < TileNIterations; ++tile_iter)
-        {
-            // We jump by 32 here since volta does <32x32x4> super mmas inside a warp.
-            scale_frag_ptr[tile_iter] = *reinterpret_cast<AccessType const*>(pointer_ + ColsPerMmaTile * tile_iter);
-        }
-    }
+    multiplies<FragmentDequantizedOperand> mul_op;
+    operand_frag = mul_op(operand_frag, scale_frag);
+  }
 
-    CUTLASS_DEVICE
-    void dequantize(FragmentDequantizedOperand& operand_frag, FragmentScale const& scale_frag)
-    {
-        static_assert(FragmentScale::kElements == FragmentDequantizedOperand::kElements, "");
-
-        multiplies<FragmentDequantizedOperand> mul_op;
-        operand_frag = mul_op(operand_frag, scale_frag);
-    }
-
-private:
-    ElementScale const* pointer_;
+ private:
+  ElementScale const* pointer_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-// Specialization for Volta A x ColumnMajor B tensorOp, for 32x32x4 interleaved gemm
+// Specialization for Volta A x ColumnMajor B tensorOp, for 32x32x4 interleaved
+// gemm
 template <
     /// Underlying matrix multiply operator (concept: MmaTensorOp)
     typename MmaOperator_,
@@ -551,98 +605,110 @@ template <
     typename Shape_,
     ///
     WeightOnlyQuantOp QuantOp_>
-class MmaTensorOpDequantizer<MmaOperator_, Shape_, Operand::kB, half_t, layout::RowMajor, 32, QuantOp_,
-    typename platform::enable_if<platform::is_same<typename MmaOperator_::ArchTag, arch::Sm70>::value
-        && platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB, layout::ColumnMajor>::value>::type>
-{
+class MmaTensorOpDequantizer<
+    MmaOperator_,
+    Shape_,
+    Operand::kB,
+    half_t,
+    layout::RowMajor,
+    32,
+    QuantOp_,
+    typename platform::enable_if<
+        platform::is_same<typename MmaOperator_::ArchTag, arch::Sm70>::value &&
+        platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB,
+                          layout::ColumnMajor>::value>::type> {
+ public:
+  static_assert(platform::is_same<typename MmaOperator_::InterleavedTileShape,
+                                  GemmShape<32, 32, 4>>::value,
+                "");
 
-public:
-    static_assert(platform::is_same<typename MmaOperator_::InterleavedTileShape, GemmShape<32, 32, 4>>::value, "");
+  /// Mma Operator
+  using MmaOperator = MmaOperator_;
 
-    /// Mma Operator
-    using MmaOperator = MmaOperator_;
+  // The architecture specific mma ooperator being used
+  using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
 
-    // The architecture specific mma ooperator being used
-    using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
+  // Mma Instruction Shape
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    // Mma Instruction Shape
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  /// Type of the scales
+  using ElementScale = half_t;
 
-    /// Type of the scales
-    using ElementScale = half_t;
+  /// Fragment to hold B data before Mma
+  using FragmentDequantizedOperand =
+      Array<ElementScale, MmaOperator::FragmentB::kElements>;
 
-    /// Fragment to hold B data before Mma
-    using FragmentDequantizedOperand = Array<ElementScale, MmaOperator::FragmentB::kElements>;
+  /// Warp mma shape
+  using Shape = Shape_;
 
-    /// Warp mma shape
-    using Shape = Shape_;
+  // Fragment to hold scale data to apply to B before mma
+  // Each 32x32x4 matmul uses 8 elements from B.
+  static constexpr int ColsPerMmaTile = 32;
+  static constexpr int TileNIterations = Shape::kN / ColsPerMmaTile;
+  using FragmentScale = Array<ElementScale, TileNIterations * 2>;
 
-    // Fragment to hold scale data to apply to B before mma
-    // Each 32x32x4 matmul uses 8 elements from B.
-    static constexpr int ColsPerMmaTile = 32;
-    static constexpr int TileNIterations = Shape::kN / ColsPerMmaTile;
-    using FragmentScale = Array<ElementScale, TileNIterations * 2>;
+  /// Layout of the scales in shared memory
+  using Layout = layout::RowMajor;
 
-    /// Layout of the scales in shared memory
-    using Layout = layout::RowMajor;
+  /// TensorRef type for loading element from a tensor
+  using TensorRef = TensorRef<ElementScale, Layout>;
 
-    /// TensorRef type for loading element from a tensor
-    using TensorRef = TensorRef<ElementScale, Layout>;
+  static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
+  static_assert(QuantOp == WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY, "");
 
-    static constexpr WeightOnlyQuantOp QuantOp = QuantOp_;
-    static_assert(QuantOp == WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY, "");
+  CUTLASS_DEVICE
+  MmaTensorOpDequantizer(TensorRef smem_scales,
+                         int const warp_idx_n,
+                         int const lane_idx) {
+    int const warp_offset = warp_idx_n * Shape::kN;
+    int const base_col = lane_idx & 0xF8 + lane_idx % 4;
+    int const thread_offset = warp_offset + base_col;
+    pointer_ = smem_scales.data() + thread_offset;
+  }
 
-    CUTLASS_DEVICE
-    MmaTensorOpDequantizer(TensorRef smem_scales, int const warp_idx_n, int const lane_idx)
-    {
-        int const warp_offset = warp_idx_n * Shape::kN;
-        int const base_col = lane_idx & 0xF8 + lane_idx % 4;
-        int const thread_offset = warp_offset + base_col;
-        pointer_ = smem_scales.data() + thread_offset;
+  CUTLASS_DEVICE
+  void load(FragmentScale& scale_frag) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int tile_iter = 0; tile_iter < TileNIterations; ++tile_iter) {
+      // We jump by 32 here since volta does <32x32x4> super mmas inside a warp.
+      // For col major B, each thread will jump 4 cols to get its next value
+      // inside of the super mma.
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_iter = 0; mma_iter < 2; ++mma_iter) {
+        scale_frag[tile_iter * 2 + mma_iter] =
+            pointer_[ColsPerMmaTile * tile_iter + 4 * mma_iter];
+      }
     }
+  }
 
-    CUTLASS_DEVICE
-    void load(FragmentScale& scale_frag)
-    {
-        CUTLASS_PRAGMA_UNROLL
-        for (int tile_iter = 0; tile_iter < TileNIterations; ++tile_iter)
-        {
-            // We jump by 32 here since volta does <32x32x4> super mmas inside a warp.
-            // For col major B, each thread will jump 4 cols to get its next value inside
-            // of the super mma.
-            CUTLASS_PRAGMA_UNROLL
-            for (int mma_iter = 0; mma_iter < 2; ++mma_iter)
-            {
-                scale_frag[tile_iter * 2 + mma_iter] = pointer_[ColsPerMmaTile * tile_iter + 4 * mma_iter];
-            }
-        }
+  CUTLASS_DEVICE
+  void dequantize(FragmentDequantizedOperand& operand_frag,
+                  FragmentScale const& scale_frag) {
+    using MmaOperandB = typename ArchMmaOperator::FragmentB;
+    static constexpr int total_n_mmas = 2 * TileNIterations;
+    static_assert(MmaOperandB::kElements * total_n_mmas ==
+                      FragmentDequantizedOperand::kElements,
+                  "");
+
+    multiplies<MmaOperandB> mul_op;
+
+    MmaOperandB* operand_frag_ptr =
+        reinterpret_cast<MmaOperandB*>(&operand_frag);
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < total_n_mmas; ++mma_n_iter) {
+      operand_frag_ptr[mma_n_iter] =
+          mul_op(operand_frag_ptr[mma_n_iter], scale_frag[mma_n_iter]);
     }
+  }
 
-    CUTLASS_DEVICE
-    void dequantize(FragmentDequantizedOperand& operand_frag, FragmentScale const& scale_frag)
-    {
-        using MmaOperandB = typename ArchMmaOperator::FragmentB;
-        static constexpr int total_n_mmas = 2 * TileNIterations;
-        static_assert(MmaOperandB::kElements * total_n_mmas == FragmentDequantizedOperand::kElements, "");
-
-        multiplies<MmaOperandB> mul_op;
-
-        MmaOperandB* operand_frag_ptr = reinterpret_cast<MmaOperandB*>(&operand_frag);
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < total_n_mmas; ++mma_n_iter)
-        {
-            operand_frag_ptr[mma_n_iter] = mul_op(operand_frag_ptr[mma_n_iter], scale_frag[mma_n_iter]);
-        }
-    }
-
-private:
-    ElementScale const* pointer_;
+ private:
+  ElementScale const* pointer_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-} // namespace warp
-} // namespace gemm
-} // namespace cutlass
+}  // namespace warp
+}  // namespace gemm
+}  // namespace cutlass
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_wint2x_dequantizer.h b/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_wint2x_dequantizer.h
index 4678b58e48..ac17315008 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_wint2x_dequantizer.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm/warp/mma_tensorop_wint2x_dequantizer.h
@@ -64,139 +64,152 @@ struct DataTypeTraits;
 
 template <>
 struct DataTypeTraits<bfloat16_t> {
-    using Type = __nv_bfloat16;
-    using DualType = __nv_bfloat162;
+  using Type = __nv_bfloat16;
+  using DualType = __nv_bfloat162;
 };
 
 template <>
 struct DataTypeTraits<half_t> {
-    using Type = __half;
-    using DualType = __half2;
+  using Type = __half;
+  using DualType = __half2;
 };
 
 template <typename T, int N, typename Enable = void>
 struct LocalScaleConverter {
-    using FragmentSource = Array<uint8_t, N>;
-    using FragmentResult = Array<T, N>;
+  using FragmentSource = Array<uint8_t, N>;
+  using FragmentResult = Array<T, N>;
 
-    CUTLASS_DEVICE
-    static void Apply(FragmentSource const& local_scale_frag,
-                      FragmentResult const& super_scale_frag,
-                      FragmentResult& scale_frag,
-                      int shift_bit) {
-        constexpr uint32_t kLocalScaleMask = 0xf;
+  CUTLASS_DEVICE
+  static void Apply(FragmentSource const& local_scale_frag,
+                    FragmentResult const& super_scale_frag,
+                    FragmentResult& scale_frag,
+                    int shift_bit) {
+    constexpr uint32_t kLocalScaleMask = 0xf;
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N; ++i) {
-            int32_t shifted_value = (static_cast<int32_t>(local_scale_frag[i]) >> shift_bit) & kLocalScaleMask;
-            scale_frag[i] = static_cast<T>(shifted_value) * super_scale_frag[i];
-        }
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      int32_t shifted_value =
+          (static_cast<int32_t>(local_scale_frag[i]) >> shift_bit) &
+          kLocalScaleMask;
+      scale_frag[i] = static_cast<T>(shifted_value) * super_scale_frag[i];
     }
+  }
 };
 
 template <int N>
-struct LocalScaleConverter<half_t, N, typename platform::enable_if<N % 4 == 0>::type> {
-    using FragmentSource = Array<uint8_t, N>;
-    using FragmentResult = Array<half_t, N>;
+struct LocalScaleConverter<half_t,
+                           N,
+                           typename platform::enable_if<N % 4 == 0>::type> {
+  using FragmentSource = Array<uint8_t, N>;
+  using FragmentResult = Array<half_t, N>;
 
-    CUTLASS_DEVICE
-    static void Apply(FragmentSource const& local_scale_frag,
-                      FragmentResult const& super_scale_frag,
-                      FragmentResult& scale_frag,
-                      int shift_bit) {
-        constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
-        constexpr uint32_t MASK = 0x000f000f;
-        // 2^10 = 1024
-        constexpr uint32_t I4s_TO_FP16s_MAGIC_NUM = 0x64006400;
+  CUTLASS_DEVICE
+  static void Apply(FragmentSource const& local_scale_frag,
+                    FragmentResult const& super_scale_frag,
+                    FragmentResult& scale_frag,
+                    int shift_bit) {
+    constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
+    constexpr uint32_t MASK = 0x000f000f;
+    // 2^10 = 1024
+    constexpr uint32_t I4s_TO_FP16s_MAGIC_NUM = 0x64006400;
 
-        // -2^10 = -1024
-        constexpr uint32_t FP16_BIAS = 0xE400E400;
-        // 1.0
-        constexpr uint32_t FP16_ONE = 0x3C003C00;
+    // -2^10 = -1024
+    constexpr uint32_t FP16_BIAS = 0xE400E400;
+    // 1.0
+    constexpr uint32_t FP16_ONE = 0x3C003C00;
 
-        __half2* scale_ptr = reinterpret_cast<__half2 *>(&scale_frag);
-        __half2 const* super_scale_ptr = reinterpret_cast<__half2 const*>(&super_scale_frag);
+    __half2* scale_ptr = reinterpret_cast<__half2*>(&scale_frag);
+    __half2 const* super_scale_ptr =
+        reinterpret_cast<__half2 const*>(&super_scale_frag);
 
-        uint32_t const* local_scale_ptr = reinterpret_cast<uint32_t const*>(&local_scale_frag);
+    uint32_t const* local_scale_ptr =
+        reinterpret_cast<uint32_t const*>(&local_scale_frag);
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / 4; ++i) {
-            int i4s = local_scale_ptr[i] >> shift_bit;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 4; ++i) {
+      int i4s = local_scale_ptr[i] >> shift_bit;
 
-            // unpack: 0, 1
-            int32_t low = __byte_perm(i4s, i4s, 0xF1F0);
-            int32_t unpack0 = lop3<immLut>(low, MASK, I4s_TO_FP16s_MAGIC_NUM);
-            // unpack: 2, 3
-            int32_t high = __byte_perm(i4s, i4s, 0xF3F2);
-            int32_t unpack1 = lop3<immLut>(high, MASK, I4s_TO_FP16s_MAGIC_NUM);
+      // unpack: 0, 1
+      int32_t low = __byte_perm(i4s, i4s, 0xF1F0);
+      int32_t unpack0 = lop3<immLut>(low, MASK, I4s_TO_FP16s_MAGIC_NUM);
+      // unpack: 2, 3
+      int32_t high = __byte_perm(i4s, i4s, 0xF3F2);
+      int32_t unpack1 = lop3<immLut>(high, MASK, I4s_TO_FP16s_MAGIC_NUM);
 
-            __half2 scale0 = __hfma2(*reinterpret_cast<__half2*>(&unpack0),
-                                     *reinterpret_cast<const __half2*>(&FP16_ONE),
-                                     *reinterpret_cast<const __half2*>(&FP16_BIAS));
-            __half2 scale1 = __hfma2(*reinterpret_cast<__half2*>(&unpack1),
-                                     *reinterpret_cast<const __half2*>(&FP16_ONE),
-                                     *reinterpret_cast<const __half2*>(&FP16_BIAS));
+      __half2 scale0 = __hfma2(*reinterpret_cast<__half2*>(&unpack0),
+                               *reinterpret_cast<const __half2*>(&FP16_ONE),
+                               *reinterpret_cast<const __half2*>(&FP16_BIAS));
+      __half2 scale1 = __hfma2(*reinterpret_cast<__half2*>(&unpack1),
+                               *reinterpret_cast<const __half2*>(&FP16_ONE),
+                               *reinterpret_cast<const __half2*>(&FP16_BIAS));
 
-            scale_ptr[2 * i] = __hmul2(scale0, super_scale_ptr[2 * i]);
-            scale_ptr[2 * i + 1] = __hmul2(scale1, super_scale_ptr[2 * i + 1]);
-        }
+      scale_ptr[2 * i] = __hmul2(scale0, super_scale_ptr[2 * i]);
+      scale_ptr[2 * i + 1] = __hmul2(scale1, super_scale_ptr[2 * i + 1]);
     }
+  }
 };
 
 template <int N>
-struct LocalScaleConverter<bfloat16_t, N, typename platform::enable_if<N % 4 == 0>::type> {
-    using FragmentSource = Array<uint8_t, N>;
-    using FragmentResult = Array<bfloat16_t, N>;
+struct LocalScaleConverter<bfloat16_t,
+                           N,
+                           typename platform::enable_if<N % 4 == 0>::type> {
+  using FragmentSource = Array<uint8_t, N>;
+  using FragmentResult = Array<bfloat16_t, N>;
 
-    CUTLASS_DEVICE
-    static void Apply(FragmentSource const& local_scale_frag,
-                      FragmentResult const& super_scale_frag,
-                      FragmentResult& scale_frag,
-                      int shift_bit) {
+  CUTLASS_DEVICE
+  static void Apply(FragmentSource const& local_scale_frag,
+                    FragmentResult const& super_scale_frag,
+                    FragmentResult& scale_frag,
+                    int shift_bit) {
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) && defined(ENABLE_BF16))
-        constexpr uint32_t immLut = (0xF0 & 0xCC) | 0xAA;
-        constexpr uint32_t MASK = 0x000F000F;
-        constexpr uint32_t I4s_TO_BF16s_MAGIC_NUM = 0x43004300;
+    constexpr uint32_t immLut = (0xF0 & 0xCC) | 0xAA;
+    constexpr uint32_t MASK = 0x000F000F;
+    constexpr uint32_t I4s_TO_BF16s_MAGIC_NUM = 0x43004300;
 
-        constexpr uint32_t BF16_BIAS = 0xC300C300;
-        constexpr uint32_t BF16_ONE = 0x3F803F80;
+    constexpr uint32_t BF16_BIAS = 0xC300C300;
+    constexpr uint32_t BF16_ONE = 0x3F803F80;
 
-        __nv_bfloat162* scale_ptr = reinterpret_cast<__nv_bfloat162 *>(&scale_frag);
-        __nv_bfloat162 const* super_scale_ptr = reinterpret_cast<__nv_bfloat162 const*>(&super_scale_frag);
+    __nv_bfloat162* scale_ptr = reinterpret_cast<__nv_bfloat162*>(&scale_frag);
+    __nv_bfloat162 const* super_scale_ptr =
+        reinterpret_cast<__nv_bfloat162 const*>(&super_scale_frag);
 
-        uint32_t const* local_scale_ptr = reinterpret_cast<uint32_t const*>(&local_scale_frag);
+    uint32_t const* local_scale_ptr =
+        reinterpret_cast<uint32_t const*>(&local_scale_frag);
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / 4; ++i) {
-            int i4s = local_scale_ptr[i] >> shift_bit;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 4; ++i) {
+      int i4s = local_scale_ptr[i] >> shift_bit;
 
-            // unpack: 0, 1
-            int32_t low = __byte_perm(i4s, i4s, 0xF1F0);
-            int32_t unpack0 = lop3<immLut>(low, MASK, I4s_TO_BF16s_MAGIC_NUM);
-            // unpack: 2, 3
-            int32_t high = __byte_perm(i4s, i4s, 0xF3F2);
-            int32_t unpack1 = lop3<immLut>(high, MASK, I4s_TO_BF16s_MAGIC_NUM);
+      // unpack: 0, 1
+      int32_t low = __byte_perm(i4s, i4s, 0xF1F0);
+      int32_t unpack0 = lop3<immLut>(low, MASK, I4s_TO_BF16s_MAGIC_NUM);
+      // unpack: 2, 3
+      int32_t high = __byte_perm(i4s, i4s, 0xF3F2);
+      int32_t unpack1 = lop3<immLut>(high, MASK, I4s_TO_BF16s_MAGIC_NUM);
 
-            nv_bfloat162 scale0 = __hfma2(*reinterpret_cast<nv_bfloat162*>(&unpack0),
-                                          *reinterpret_cast<const nv_bfloat162*>(&BF16_ONE),
-                                          *reinterpret_cast<const nv_bfloat162*>(&BF16_BIAS));
-            nv_bfloat162 scale1 = __hfma2(*reinterpret_cast<nv_bfloat162*>(&unpack1),
-                                          *reinterpret_cast<const nv_bfloat162*>(&BF16_ONE),
-                                          *reinterpret_cast<const nv_bfloat162*>(&BF16_BIAS));
+      nv_bfloat162 scale0 =
+          __hfma2(*reinterpret_cast<nv_bfloat162*>(&unpack0),
+                  *reinterpret_cast<const nv_bfloat162*>(&BF16_ONE),
+                  *reinterpret_cast<const nv_bfloat162*>(&BF16_BIAS));
+      nv_bfloat162 scale1 =
+          __hfma2(*reinterpret_cast<nv_bfloat162*>(&unpack1),
+                  *reinterpret_cast<const nv_bfloat162*>(&BF16_ONE),
+                  *reinterpret_cast<const nv_bfloat162*>(&BF16_BIAS));
 
-            scale_ptr[2 * i] = __hmul2(scale0, super_scale_ptr[2 * i]);
-            scale_ptr[2 * i + 1] = __hmul2(scale1, super_scale_ptr[2 * i + 1]);
-        }
-#else
-        // Slow path not implemented here on purpose. If we need to do HMMA on older arch, scale conversion should
-        // happen before scales are stored to shared memory and we should use the fp16 dequantizer. This will avoid
-        // numerous conversion instructions in GEMM main loop.
-        arch::device_breakpoint();
-#endif
+      scale_ptr[2 * i] = __hmul2(scale0, super_scale_ptr[2 * i]);
+      scale_ptr[2 * i + 1] = __hmul2(scale1, super_scale_ptr[2 * i + 1]);
     }
+#else
+    // Slow path not implemented here on purpose. If we need to do HMMA on older
+    // arch, scale conversion should happen before scales are stored to shared
+    // memory and we should use the fp16 dequantizer. This will avoid numerous
+    // conversion instructions in GEMM main loop.
+    arch::device_breakpoint();
+#endif
+  }
 };
 
-} // namespace detail
+}  // namespace detail
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -216,7 +229,7 @@ template <
     ///
     typename Enable = void>
 class MmaTensorOpWin2xDequantizer {
-    //static_assert(false, "Not Supported!");
+  // static_assert(false, "Not Supported!");
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -230,207 +243,237 @@ template <
     typename ElementOperand_,
     /// Group size for quantization
     int GroupSize_>
-class MmaTensorOpWin2xDequantizer<
-    MmaOperator_,
-    Shape_,
-    Operand::kB,
-    ElementOperand_,
-    layout::RowMajor,
-    GroupSize_>
-    //typename platform::enable_if<MmaOperator_::ArchTag::kMinComputeCapability >= 80
-    //    && platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB, layout::ColumnMajor>::value>::type>
+class MmaTensorOpWin2xDequantizer<MmaOperator_,
+                                  Shape_,
+                                  Operand::kB,
+                                  ElementOperand_,
+                                  layout::RowMajor,
+                                  GroupSize_>
+// typename platform::enable_if<MmaOperator_::ArchTag::kMinComputeCapability >=
+// 80
+//     && platform::is_same<typename MmaOperator_::ArchMmaOperator::LayoutB,
+//     layout::ColumnMajor>::value>::type>
 {
-public:
-    static_assert(platform::is_same<ElementOperand_, half_t>::value || platform::is_same<ElementOperand_, bfloat16_t>::value,
-        "T must be fp16 or bf16");
+ public:
+  static_assert(platform::is_same<ElementOperand_, half_t>::value ||
+                    platform::is_same<ElementOperand_, bfloat16_t>::value,
+                "T must be fp16 or bf16");
 
-    /// Mma Operator
-    using MmaOperator = MmaOperator_;
+  /// Mma Operator
+  using MmaOperator = MmaOperator_;
 
-    // The architecture specific mma ooperator being used
-    using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
+  // The architecture specific mma ooperator being used
+  using ArchMmaOperator = typename MmaOperator::ArchMmaOperator;
 
-    // Mma Instruction Shape
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  // Mma Instruction Shape
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    /// Warp mma shape
-    using Shape = Shape_;
+  /// Warp mma shape
+  using Shape = Shape_;
 
-    /// Type of mma operand
-    using ElementOperand = ElementOperand_;
+  /// Type of mma operand
+  using ElementOperand = ElementOperand_;
 
-    /// Layout of the scales in shared memory
-    using Layout = layout::RowMajor;
+  /// Layout of the scales in shared memory
+  using Layout = layout::RowMajor;
 
-    /// Group size for quantization
-    static constexpr int kGroupSize = GroupSize_;
+  /// Group size for quantization
+  static constexpr int kGroupSize = GroupSize_;
 
-    /// Type of input
-    using ElementB = typename MmaOperator::FragmentB::Element;
-    static_assert(platform::is_same<ElementB, uint2b_t>::value, "ElementB must be uint2b_t");
+  /// Type of input
+  using ElementB = typename MmaOperator::FragmentB::Element;
+  static_assert(platform::is_same<ElementB, uint2b_t>::value,
+                "ElementB must be uint2b_t");
 
-    /// Type of the scales
-    using ElementLocalScale = uint4b_t;
-    using ElementSuperScale = ElementOperand;
-    using ElementCodeScaleZp = float;
+  /// Type of the scales
+  using ElementLocalScale = uint4b_t;
+  using ElementSuperScale = ElementOperand;
+  using ElementCodeScaleZp = float;
 
-    // Fragment to hold scale data to apply to B before mma
-    // We need 1 fp16 per matrix iteration in the N dimension
-    static constexpr int kWarpIterationsAlongN = MmaOperator::MmaIterations::kColumn;
+  // Fragment to hold scale data to apply to B before mma
+  // We need 1 fp16 per matrix iteration in the N dimension
+  static constexpr int kWarpIterationsAlongN =
+      MmaOperator::MmaIterations::kColumn;
 
-    // use uint8_t to save 2 4-bits local scales
-    using FragmentLocalScale = Array<uint8_t, kWarpIterationsAlongN>;
-    using FragmentSuperScale = Array<ElementSuperScale, kWarpIterationsAlongN>;
-    using FragmentCodeScaleZp = Array<ElementCodeScaleZp, kWarpIterationsAlongN>;
+  // use uint8_t to save 2 4-bits local scales
+  using FragmentLocalScale = Array<uint8_t, kWarpIterationsAlongN>;
+  using FragmentSuperScale = Array<ElementSuperScale, kWarpIterationsAlongN>;
+  using FragmentCodeScaleZp = Array<ElementCodeScaleZp, kWarpIterationsAlongN>;
 
-    /// Fragment to hold B data before Mma
-    using FragmentInput = Array<ElementB, MmaOperator::FragmentB::kElements>;
+  /// Fragment to hold B data before Mma
+  using FragmentInput = Array<ElementB, MmaOperator::FragmentB::kElements>;
 
-    // This is the ratio of the load instruction vs the compute instruction.
-    static constexpr int kExpansionFactor = MmaOperator::IteratorB::InstructionShape::kRow / InstructionShape::kK;
+  // This is the ratio of the load instruction vs the compute instruction.
+  static constexpr int kExpansionFactor =
+      MmaOperator::IteratorB::InstructionShape::kRow / InstructionShape::kK;
 
-    static constexpr int kNumPacks = sizeof_bits<uint8_t>::value / sizeof_bits<ElementB>::value;
-    static constexpr int kUnpackFactor = MmaOperator::FragmentB::kElements / (kWarpIterationsAlongN * kNumPacks);
-    static constexpr int kUnpackInterval = kExpansionFactor / kUnpackFactor;
+  static constexpr int kNumPacks =
+      sizeof_bits<uint8_t>::value / sizeof_bits<ElementB>::value;
+  static constexpr int kUnpackFactor =
+      MmaOperator::FragmentB::kElements / (kWarpIterationsAlongN * kNumPacks);
+  static constexpr int kUnpackInterval = kExpansionFactor / kUnpackFactor;
 
-    /// Unpack 4 uint2b_t values compreseed in a uint8_t to floating points.
-    using Uint2Converter = FastInterleavedAndBiasedNumericArrayConverter<
-        ElementOperand, ElementB, MmaOperator::FragmentB::kElements / kUnpackFactor>;
-    using FragmentInputUnpack = typename Uint2Converter::result_type;
+  /// Unpack 4 uint2b_t values compreseed in a uint8_t to floating points.
+  using Uint2Converter = FastInterleavedAndBiasedNumericArrayConverter<
+      ElementOperand,
+      ElementB,
+      MmaOperator::FragmentB::kElements / kUnpackFactor>;
+  using FragmentInputUnpack = typename Uint2Converter::result_type;
 
-    /// Fragment to hold internal scales before Mma
-    using FragmentScale = Array<ElementOperand, FragmentLocalScale::kElements>;
+  /// Fragment to hold internal scales before Mma
+  using FragmentScale = Array<ElementOperand, FragmentLocalScale::kElements>;
 
-    /// Fragment of dequantized B
-    using FragmentOutput = Array<ElementOperand, MmaOperator::FragmentB::kElements / kExpansionFactor>;
+  /// Fragment of dequantized B
+  using FragmentOutput =
+      Array<ElementOperand,
+            MmaOperator::FragmentB::kElements / kExpansionFactor>;
 
-    /// TensorRef type for loading element from a tensor
-    using SuperTensorRef = cutlass::TensorRef<ElementSuperScale, Layout>;
-    using LocalTensorRef = cutlass::TensorRef<ElementLocalScale, Layout>;
-    using CodeTensorRef = cutlass::TensorRef<ElementCodeScaleZp, Layout>;
+  /// TensorRef type for loading element from a tensor
+  using SuperTensorRef = cutlass::TensorRef<ElementSuperScale, Layout>;
+  using LocalTensorRef = cutlass::TensorRef<ElementLocalScale, Layout>;
+  using CodeTensorRef = cutlass::TensorRef<ElementCodeScaleZp, Layout>;
 
-private:
-    //
-    // Data members
-    //
+ private:
+  //
+  // Data members
+  //
 
-    uint8_t* pointer_local_scale_;
-    ElementCodeScaleZp* pointer_code_scale_;
-    ElementCodeScaleZp* pointer_code_zp_;
-    ElementSuperScale* pointer_super_scale_;
+  uint8_t* pointer_local_scale_;
+  ElementCodeScaleZp* pointer_code_scale_;
+  ElementCodeScaleZp* pointer_code_zp_;
+  ElementSuperScale* pointer_super_scale_;
 
-    //FragmentInputUnpack unpacked_frag_;
-    FragmentScale scale_frag_;
+  // FragmentInputUnpack unpacked_frag_;
+  FragmentScale scale_frag_;
 
-public:
-    CUTLASS_DEVICE
-    MmaTensorOpWin2xDequantizer(SuperTensorRef smem_super_scale,
-                                LocalTensorRef smem_local_scale,
-                                CodeTensorRef smem_code_scale,
-                                CodeTensorRef smem_code_zp,
-                                int warp_idx_n,
-                                int lane_idx) {
-        int warp_offset = warp_idx_n * Shape::kN;
-        int quad = lane_idx / 4;
-        int thread_offset = warp_offset + quad;
-        pointer_super_scale_ = smem_super_scale.data() + thread_offset;
-        pointer_code_scale_ = smem_code_scale.data() + thread_offset;
-        pointer_code_zp_ = smem_code_zp.data() + thread_offset;
-        pointer_local_scale_ = reinterpret_cast<uint8_t *>(smem_local_scale.data()) + thread_offset;
+ public:
+  CUTLASS_DEVICE
+  MmaTensorOpWin2xDequantizer(SuperTensorRef smem_super_scale,
+                              LocalTensorRef smem_local_scale,
+                              CodeTensorRef smem_code_scale,
+                              CodeTensorRef smem_code_zp,
+                              int warp_idx_n,
+                              int lane_idx) {
+    int warp_offset = warp_idx_n * Shape::kN;
+    int quad = lane_idx / 4;
+    int thread_offset = warp_offset + quad;
+    pointer_super_scale_ = smem_super_scale.data() + thread_offset;
+    pointer_code_scale_ = smem_code_scale.data() + thread_offset;
+    pointer_code_zp_ = smem_code_zp.data() + thread_offset;
+    pointer_local_scale_ =
+        reinterpret_cast<uint8_t*>(smem_local_scale.data()) + thread_offset;
+  }
+
+  /// Channel-wise params, need to load just once
+  CUTLASS_DEVICE
+  void load(FragmentCodeScaleZp& code_scale_frag,
+            FragmentCodeScaleZp& code_zp_frag,
+            FragmentSuperScale& super_scale_frag) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < kWarpIterationsAlongN; ++mma_n_iter) {
+      super_scale_frag[mma_n_iter] =
+          pointer_super_scale_[mma_n_iter *
+                               InstructionShape::kN];  // bank conflict
+      code_scale_frag[mma_n_iter] =
+          pointer_code_scale_[mma_n_iter * InstructionShape::kN];
+      code_zp_frag[mma_n_iter] =
+          pointer_code_zp_[mma_n_iter * InstructionShape::kN];
+    }
+  }
+
+  /// Group-wise params, need to load multiple times
+  CUTLASS_DEVICE
+  void load(FragmentLocalScale& local_scale_frag) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < kWarpIterationsAlongN; ++mma_n_iter) {
+      local_scale_frag[mma_n_iter] =
+          pointer_local_scale_[mma_n_iter *
+                               InstructionShape::kN];  // bank conflict
+    }
+  }
+
+  CUTLASS_DEVICE
+  void dequantize(const FragmentLocalScale& local_scale_frag,
+                  const FragmentCodeScaleZp& code_scale_frag,
+                  const FragmentCodeScaleZp& code_zp_frag,
+                  const FragmentSuperScale& super_scale_frag,
+                  const FragmentInput& input_frag,
+                  FragmentOutput& output_frag,
+                  int tb_offset_k,
+                  int warp_k_compute_offset) {
+    if constexpr (kUnpackInterval != 1) {
+      // unsupport now
+      arch::device_breakpoint();
     }
 
-    /// Channel-wise params, need to load just once
-    CUTLASS_DEVICE
-    void load(FragmentCodeScaleZp& code_scale_frag,
-              FragmentCodeScaleZp& code_zp_frag,
-              FragmentSuperScale& super_scale_frag) {
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < kWarpIterationsAlongN; ++mma_n_iter) {
-            super_scale_frag[mma_n_iter] = pointer_super_scale_[mma_n_iter * InstructionShape::kN]; // bank conflict
-            code_scale_frag[mma_n_iter] = pointer_code_scale_[mma_n_iter * InstructionShape::kN];
-            code_zp_frag[mma_n_iter] = pointer_code_zp_[mma_n_iter * InstructionShape::kN];
-        }
+    typename Uint2Converter::source_type source_frag;
+
+    int in_offset = warp_k_compute_offset * kUnpackInterval;
+
+    uint8_t const* ptr_input = reinterpret_cast<uint8_t const*>(&input_frag);
+    uint8_t* ptr_source = reinterpret_cast<uint8_t*>(&source_frag);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < kWarpIterationsAlongN; ++mma_n_iter) {
+      ptr_source[mma_n_iter] =
+          ptr_input[mma_n_iter * kUnpackFactor + in_offset];
+    }
+    FragmentInputUnpack unpacked_frag =
+        Uint2Converter::convert(source_frag, code_scale_frag, code_zp_frag);
+
+    // dequantize local_scale
+    if (warp_k_compute_offset == 0) {
+      using LocalScaleConverter =
+          detail::LocalScaleConverter<ElementOperand,
+                                      FragmentLocalScale::kElements>;
+
+      // special for TileRows = 64
+      int local_scale_shift = (((tb_offset_k / kGroupSize) + 1) & 1) * 4;
+      LocalScaleConverter::Apply(
+          local_scale_frag, super_scale_frag, scale_frag_, local_scale_shift);
     }
 
-    /// Group-wise params, need to load multiple times
-    CUTLASS_DEVICE
-    void load(FragmentLocalScale& local_scale_frag) {
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < kWarpIterationsAlongN; ++mma_n_iter) {
-            local_scale_frag[mma_n_iter] = pointer_local_scale_[mma_n_iter * InstructionShape::kN]; // bank conflict
-        }
+    // unscale
+    // After applying LOP3 optimizations for performance, the B operand requires
+    // data rearrangement. reorder: [0, 4, 1, 5, 2, 6, 3, 7, 8, 12, 9, 13, 10,
+    // 14, 11, 15]
+    const int kWarpIterationsAlongK =
+        FragmentOutput::kElements / kWarpIterationsAlongN;
+
+    using Type = typename detail::DataTypeTraits<ElementOperand>::Type;
+    using DualType = typename detail::DataTypeTraits<ElementOperand>::DualType;
+
+    Type* output_ptr = reinterpret_cast<Type*>(&output_frag);
+    DualType const* unpacked_ptr =
+        reinterpret_cast<DualType const*>(&unpacked_frag);
+    DualType const* scale_ptr = reinterpret_cast<DualType const*>(&scale_frag_);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int mma_n_iter = 0; mma_n_iter < kWarpIterationsAlongN;
+         mma_n_iter += 2) {
+      int mapped_idx_base = (mma_n_iter / 2) * kWarpIterationsAlongK;
+
+      DualType scalex2 = scale_ptr[mma_n_iter / 2];
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int mma_k_iter = 0; mma_k_iter < kWarpIterationsAlongK;
+           ++mma_k_iter) {
+        DualType unpacked_valuex2 = unpacked_ptr[mapped_idx_base + mma_k_iter];
+        DualType scaled_value = __hmul2(unpacked_valuex2, scalex2);
+        output_ptr[mma_n_iter * kWarpIterationsAlongK + mma_k_iter] =
+            scaled_value.x;
+        output_ptr[(mma_n_iter + 1) * kWarpIterationsAlongK + mma_k_iter] =
+            scaled_value.y;
+      }
     }
+  }
 
-    CUTLASS_DEVICE
-    void dequantize(const FragmentLocalScale& local_scale_frag,
-                    const FragmentCodeScaleZp& code_scale_frag,
-                    const FragmentCodeScaleZp& code_zp_frag,
-                    const FragmentSuperScale& super_scale_frag,
-                    const FragmentInput& input_frag,
-                    FragmentOutput& output_frag,
-                    int tb_offset_k,
-                    int warp_k_compute_offset) {
-        if constexpr (kUnpackInterval != 1) {
-            // unsupport now
-            arch::device_breakpoint();
-        }
-
-        typename Uint2Converter::source_type source_frag;
-
-        int in_offset = warp_k_compute_offset * kUnpackInterval;
-
-        uint8_t const* ptr_input = reinterpret_cast<uint8_t const*>(&input_frag);
-        uint8_t* ptr_source = reinterpret_cast<uint8_t *>(&source_frag);
-
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < kWarpIterationsAlongN; ++mma_n_iter) {
-            ptr_source[mma_n_iter] = ptr_input[mma_n_iter * kUnpackFactor + in_offset];
-        }
-        FragmentInputUnpack unpacked_frag = Uint2Converter::convert(source_frag, code_scale_frag, code_zp_frag);
-
-        // dequantize local_scale
-        if (warp_k_compute_offset == 0) {
-            using LocalScaleConverter = detail::LocalScaleConverter<ElementOperand, FragmentLocalScale::kElements>;
-
-            // special for TileRows = 64
-            int local_scale_shift = (((tb_offset_k / kGroupSize) + 1) & 1) * 4;
-            LocalScaleConverter::Apply(local_scale_frag, super_scale_frag, scale_frag_, local_scale_shift);
-        }
-
-        // unscale
-        // After applying LOP3 optimizations for performance, the B operand requires data rearrangement.
-        // reorder: [0, 4, 1, 5, 2, 6, 3, 7, 8, 12, 9, 13, 10, 14, 11, 15]
-        const int kWarpIterationsAlongK = FragmentOutput::kElements / kWarpIterationsAlongN;
-
-        using Type = typename detail::DataTypeTraits<ElementOperand>::Type;
-        using DualType = typename detail::DataTypeTraits<ElementOperand>::DualType;
-
-        Type* output_ptr = reinterpret_cast<Type *>(&output_frag);
-        DualType const* unpacked_ptr = reinterpret_cast<DualType const*>(&unpacked_frag);
-        DualType const* scale_ptr = reinterpret_cast<DualType const*>(&scale_frag_);
-
-        CUTLASS_PRAGMA_UNROLL
-        for (int mma_n_iter = 0; mma_n_iter < kWarpIterationsAlongN; mma_n_iter += 2) {
-            int mapped_idx_base = (mma_n_iter / 2) * kWarpIterationsAlongK;
-
-            DualType scalex2 = scale_ptr[mma_n_iter / 2];
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int mma_k_iter = 0; mma_k_iter < kWarpIterationsAlongK; ++mma_k_iter) {
-                DualType unpacked_valuex2 = unpacked_ptr[mapped_idx_base + mma_k_iter];
-                DualType scaled_value = __hmul2(unpacked_valuex2, scalex2);
-                output_ptr[mma_n_iter * kWarpIterationsAlongK + mma_k_iter] = scaled_value.x;
-                output_ptr[(mma_n_iter + 1) * kWarpIterationsAlongK + mma_k_iter] = scaled_value.y;
-            }
-        }
-    }
-
-    /// Add an offset to pointer in units of elements.
-    /// Only group-wise params needs.
-    CUTLASS_DEVICE
-    void add_pointer_offset(int64_t const& offset) {
-        pointer_local_scale_ += offset;
-    }
+  /// Add an offset to pointer in units of elements.
+  /// Only group-wise params needs.
+  CUTLASS_DEVICE
+  void add_pointer_offset(int64_t const& offset) {
+    pointer_local_scale_ += offset;
+  }
 };
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/gemm_configs.h b/custom_ops/gpu_ops/cutlass_extensions/gemm_configs.h
index 81e58f20ef..02becf2388 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/gemm_configs.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/gemm_configs.h
@@ -21,301 +21,322 @@
 #include <sstream>
 #include <string>
 
-namespace cutlass_extensions
-{
-// Note: The shapes are in the format MxNxK. The K shape of the runtime config MUST match the K shape
+namespace cutlass_extensions {
+// Note: The shapes are in the format MxNxK. The K shape of the runtime config
+// MUST match the K shape
 //       in the kernel layout details when doing weight only quantization.
-enum class CutlassTileConfig
-{
-    // Signals that we should run heuristics do choose a config
-    Undefined,
+enum class CutlassTileConfig {
+  // Signals that we should run heuristics do choose a config
+  Undefined,
 
-    // Signals that we should run heuristics do choose a config
-    ChooseWithHeuristic,
+  // Signals that we should run heuristics do choose a config
+  ChooseWithHeuristic,
 
-    // SiMT config
-    CtaShape128x128x8_WarpShape64x64x8,
+  // SiMT config
+  CtaShape128x128x8_WarpShape64x64x8,
 
-    // TensorCore configs CTA_N = 128, CTA_K = 64
-    // Warp configs for M=16
-    CtaShape16x128x64_WarpShape16x32x64,
-    // Warp configs for M=32
-    CtaShape32x128x64_WarpShape32x32x64,
+  // TensorCore configs CTA_N = 128, CTA_K = 64
+  // Warp configs for M=16
+  CtaShape16x128x64_WarpShape16x32x64,
+  // Warp configs for M=32
+  CtaShape32x128x64_WarpShape32x32x64,
 
-    // Warp configs for M=64
-    CtaShape64x128x64_WarpShape32x64x64,
-    CtaShape64x64x128_WarpShape32x64x64,
-    CtaShape64x128x64_WarpShape64x32x64,
+  // Warp configs for M=64
+  CtaShape64x128x64_WarpShape32x64x64,
+  CtaShape64x64x128_WarpShape32x64x64,
+  CtaShape64x128x64_WarpShape64x32x64,
 
-    // Warp configs for M=128
-    CtaShape128x64x64_WarpShape64x32x64,
-    CtaShape128x128x64_WarpShape64x32x64,
-    CtaShape128x128x64_WarpShape64x64x64,
-    CtaShape128x128x64_WarpShape128x32x64,
-    CtaShape128x256x64_WarpShape64x64x64,
+  // Warp configs for M=128
+  CtaShape128x64x64_WarpShape64x32x64,
+  CtaShape128x128x64_WarpShape64x32x64,
+  CtaShape128x128x64_WarpShape64x64x64,
+  CtaShape128x128x64_WarpShape128x32x64,
+  CtaShape128x256x64_WarpShape64x64x64,
 
-    // Warp configs for M=256
-    CtaShape256x128x64_WarpShape64x64x64,
+  // Warp configs for M=256
+  CtaShape256x128x64_WarpShape64x64x64,
 
-    // TensorCore config CTA_N = 64, CTA_K = 128
-    CtaShape128x64x128_WarpShape64x32x128,
+  // TensorCore config CTA_N = 64, CTA_K = 128
+  CtaShape128x64x128_WarpShape64x32x128,
 
-    // TensorCore config CTA_N = 256, CTA_K = 64
-    CtaShape16x256x64_WarpShape16x64x64,
+  // TensorCore config CTA_N = 256, CTA_K = 64
+  CtaShape16x256x64_WarpShape16x64x64,
 
-    // TensorCore config CTA_N = 256, CTA_K = 128
-    CtaShape16x256x128_WarpShape16x64x128
+  // TensorCore config CTA_N = 256, CTA_K = 128
+  CtaShape16x256x128_WarpShape16x64x128
 
 };
 
-enum class SplitKStyle
-{
-    NO_SPLIT_K,
-    SPLIT_K_SERIAL,
-    STREAM_K, // Sm80+
-    // SPLIT_K_PARALLEL // Not supported yet
+enum class SplitKStyle {
+  NO_SPLIT_K,
+  SPLIT_K_SERIAL,
+  STREAM_K,  // Sm80+
+             // SPLIT_K_PARALLEL // Not supported yet
 };
 
 // New enum for SM100 (Blackwell) Tile Configs
 // Placeholder values - actual optimal values need research
-enum class CutlassTileConfigSM100
-{
-    // Signals that we should run heuristics do choose a config
-    Undefined,
+enum class CutlassTileConfigSM100 {
+  // Signals that we should run heuristics do choose a config
+  Undefined,
 
-    // Signals that we should run heuristics do choose a config
-    ChooseWithHeuristic,
+  // Signals that we should run heuristics do choose a config
+  ChooseWithHeuristic,
 
-    // Actual SM100 tile configs based on user input (K-tile is 128B)
-    CtaShape64x64x128B,
-    CtaShape64x128x128B,
-    CtaShape64x256x128B,
-    CtaShape128x64x128B,
-    CtaShape128x128x128B,
-    CtaShape128x256x128B,
-    CtaShape256x64x128B,
-    CtaShape256x128x128B,
-    CtaShape256x256x128B
-    // Note: The user-provided list for get_candidate_tiles_sm100 also includes
-    // CtaShape128x64x128B and CtaShape256x64x128B for specific FP4 grouped gemm cases.
-    // These are already covered by the list above if general suffices.
-    // If they need distinct enum values, they should be added.
-    // For now, keeping the enum concise with unique shapes mentioned for general use.
+  // Actual SM100 tile configs based on user input (K-tile is 128B)
+  CtaShape64x64x128B,
+  CtaShape64x128x128B,
+  CtaShape64x256x128B,
+  CtaShape128x64x128B,
+  CtaShape128x128x128B,
+  CtaShape128x256x128B,
+  CtaShape256x64x128B,
+  CtaShape256x128x128B,
+  CtaShape256x256x128B
+  // Note: The user-provided list for get_candidate_tiles_sm100 also includes
+  // CtaShape128x64x128B and CtaShape256x64x128B for specific FP4 grouped gemm
+  // cases. These are already covered by the list above if general suffices. If
+  // they need distinct enum values, they should be added. For now, keeping the
+  // enum concise with unique shapes mentioned for general use.
 };
 
+enum class CutlassTileConfigSM90 {
+  // Signals that we should run heuristics do choose a config
+  Undefined,
 
-enum class CutlassTileConfigSM90
-{
-    // Signals that we should run heuristics do choose a config
-    Undefined,
+  // Signals that we should run heuristics do choose a config
+  ChooseWithHeuristic,
 
-    // Signals that we should run heuristics do choose a config
-    ChooseWithHeuristic,
+  // CTA configs for M=64
+  CtaShape64x16x128B,
+  CtaShape64x32x128B,
+  CtaShape64x64x128B,
+  CtaShape64x128x128B,
+  CtaShape64x256x128B,
 
-    // CTA configs for M=64
-    CtaShape64x16x128B,
-    CtaShape64x32x128B,
-    CtaShape64x64x128B,
-    CtaShape64x128x128B,
-    CtaShape64x256x128B,
+  // CTA configs for M=128
+  CtaShape128x16x128B,
+  CtaShape128x32x128B,
+  CtaShape128x64x128B,
+  CtaShape128x128x128B,
+  CtaShape128x256x128B,
 
-    // CTA configs for M=128
-    CtaShape128x16x128B,
-    CtaShape128x32x128B,
-    CtaShape128x64x128B,
-    CtaShape128x128x128B,
-    CtaShape128x256x128B,
-
-    // CTA configs for M=128
-    CtaShape256x128x128B,
+  // CTA configs for M=128
+  CtaShape256x128x128B,
 };
 
-enum class MainloopScheduleType
-{
-    AUTO // Automatically selects between pingpong and cooperative schedules on Hopper. On older architectures, this
-         // defaults to the "legacy" main loop schedule.
+enum class MainloopScheduleType {
+  AUTO  // Automatically selects between pingpong and cooperative schedules on
+        // Hopper. On older architectures, this defaults to the "legacy" main
+        // loop schedule.
 };
 
-enum class EpilogueScheduleType
-{
-    AUTO // Automatically chooses an epilogue schedule compatible with the selected main loop schedule for Hopper. For
-         // architectures older than hopper, the epilogue is always performed by the same thread block as the main loop.
+enum class EpilogueScheduleType {
+  AUTO  // Automatically chooses an epilogue schedule compatible with the
+        // selected main loop schedule for Hopper. For architectures older than
+        // hopper, the epilogue is always performed by the same thread block as
+        // the main loop.
 };
 
-enum class ClusterShape
-{
-    ClusterShape_1x1x1,
-    ClusterShape_2x1x1,
-    ClusterShape_1x2x1,
-    ClusterShape_2x2x1,
-    ClusterShape_1x8x1,
-    ClusterShape_8x1x1
+enum class ClusterShape {
+  ClusterShape_1x1x1,
+  ClusterShape_2x1x1,
+  ClusterShape_1x2x1,
+  ClusterShape_2x2x1,
+  ClusterShape_1x8x1,
+  ClusterShape_8x1x1
 };
 
-struct CutlassGemmConfig
-{
-    enum CandidateConfigTypeParam : int
-    {
-        NONE = 0,
-        WEIGHT_ONLY = 1u << 0,
-        SIMT_ONLY = 1u << 1,
-        INT8_ONLY = 1u << 2,
-        HOPPER = 1u << 3, // SM90
-        GROUPED_GEMM = 1u << 4,
-        FP8_ONLY = 1u << 5,
-        BLACKWELL = 1u << 6, // SM100
-        FP4_ONLY = 1u << 7, // For Blackwell FP4/MXFP4 paths
-    };
+struct CutlassGemmConfig {
+  enum CandidateConfigTypeParam : int {
+    NONE = 0,
+    WEIGHT_ONLY = 1u << 0,
+    SIMT_ONLY = 1u << 1,
+    INT8_ONLY = 1u << 2,
+    HOPPER = 1u << 3,  // SM90
+    GROUPED_GEMM = 1u << 4,
+    FP8_ONLY = 1u << 5,
+    BLACKWELL = 1u << 6,  // SM100
+    FP4_ONLY = 1u << 7,   // For Blackwell FP4/MXFP4 paths
+  };
 
-    CutlassTileConfig tile_config = CutlassTileConfig::ChooseWithHeuristic;
-    SplitKStyle split_k_style = SplitKStyle::NO_SPLIT_K;
-    int split_k_factor = -1;
-    int stages = -1;
+  CutlassTileConfig tile_config = CutlassTileConfig::ChooseWithHeuristic;
+  SplitKStyle split_k_style = SplitKStyle::NO_SPLIT_K;
+  int split_k_factor = -1;
+  int stages = -1;
 
-    // config options for sm90
-    CutlassTileConfigSM90 tile_config_sm90 = CutlassTileConfigSM90::ChooseWithHeuristic;
-    MainloopScheduleType mainloop_schedule = MainloopScheduleType::AUTO;
-    EpilogueScheduleType epilogue_schedule = EpilogueScheduleType::AUTO;
-    ClusterShape cluster_shape = ClusterShape::ClusterShape_1x1x1;
-    bool is_sm90 = false;
+  // config options for sm90
+  CutlassTileConfigSM90 tile_config_sm90 =
+      CutlassTileConfigSM90::ChooseWithHeuristic;
+  MainloopScheduleType mainloop_schedule = MainloopScheduleType::AUTO;
+  EpilogueScheduleType epilogue_schedule = EpilogueScheduleType::AUTO;
+  ClusterShape cluster_shape = ClusterShape::ClusterShape_1x1x1;
+  bool is_sm90 = false;
 
-    // config options for sm100 (Blackwell)
-    // Assuming SM100 might use similar schedule/cluster types as SM90 for now.
-    // These might need to become SM100-specific if Blackwell introduces new concepts.
-    CutlassTileConfigSM100 tile_config_sm100 = CutlassTileConfigSM100::ChooseWithHeuristic;
-    // MainloopScheduleType mainloop_schedule_sm100 = MainloopScheduleType::AUTO; // Example if SM100 has different types
-    // EpilogueScheduleType epilogue_schedule_sm100 = EpilogueScheduleType::AUTO; // Example
-    // ClusterShape cluster_shape_sm100 = ClusterShape::ClusterShape_1x1x1;       // Example
-    bool is_sm100 = false;
+  // config options for sm100 (Blackwell)
+  // Assuming SM100 might use similar schedule/cluster types as SM90 for now.
+  // These might need to become SM100-specific if Blackwell introduces new
+  // concepts.
+  CutlassTileConfigSM100 tile_config_sm100 =
+      CutlassTileConfigSM100::ChooseWithHeuristic;
+  // MainloopScheduleType mainloop_schedule_sm100 = MainloopScheduleType::AUTO;
+  // // Example if SM100 has different types EpilogueScheduleType
+  // epilogue_schedule_sm100 = EpilogueScheduleType::AUTO; // Example
+  // ClusterShape cluster_shape_sm100 = ClusterShape::ClusterShape_1x1x1; //
+  // Example
+  bool is_sm100 = false;
 
+  CutlassGemmConfig() : is_sm90(false), is_sm100(false) {}
 
-    CutlassGemmConfig() : is_sm90(false), is_sm100(false) {}
+  CutlassGemmConfig(CutlassTileConfig tile_config,
+                    SplitKStyle split_k_style,
+                    int split_k_factor,
+                    int stages)
+      : tile_config(tile_config),
+        split_k_style(split_k_style),
+        split_k_factor(split_k_factor),
+        stages(stages),
+        is_sm90(false),
+        is_sm100(false) {}
 
-    CutlassGemmConfig(CutlassTileConfig tile_config, SplitKStyle split_k_style, int split_k_factor, int stages)
-        : tile_config(tile_config)
-        , split_k_style(split_k_style)
-        , split_k_factor(split_k_factor)
-        , stages(stages)
-        , is_sm90(false)
-        , is_sm100(false)
-    {
+  // Constructor for SM90
+  CutlassGemmConfig(CutlassTileConfigSM90 tile_config_sm90_in,
+                    MainloopScheduleType mainloop_schedule_in,
+                    EpilogueScheduleType epilogue_schedule_in,
+                    ClusterShape cluster_shape_in)
+      : tile_config_sm90(tile_config_sm90_in),
+        mainloop_schedule(mainloop_schedule_in),
+        epilogue_schedule(epilogue_schedule_in),
+        cluster_shape(cluster_shape_in),
+        is_sm90(true),
+        is_sm100(false) {}
+
+  // Constructor for SM100 (Blackwell)
+  // Using existing MainloopScheduleType, EpilogueScheduleType, ClusterShape for
+  // now. These might need to be new SM100-specific types if Blackwell's TMA
+  // differs significantly.
+  CutlassGemmConfig(CutlassTileConfigSM100 tile_config_sm100_in,
+                    MainloopScheduleType mainloop_schedule_in,
+                    EpilogueScheduleType epilogue_schedule_in,
+                    ClusterShape cluster_shape_in)
+      : tile_config_sm100(tile_config_sm100_in),
+        mainloop_schedule(
+            mainloop_schedule_in)  // Potentially use mainloop_schedule_sm100 if
+                                   // types diverge
+        ,
+        epilogue_schedule(
+            epilogue_schedule_in)  // Potentially use epilogue_schedule_sm100
+        ,
+        cluster_shape(cluster_shape_in)  // Potentially use cluster_shape_sm100
+        ,
+        is_sm90(false)  // Explicitly false
+        ,
+        is_sm100(true) {}
+
+  std::string toString() const {
+    std::stringstream tactic;
+    tactic << "Cutlass GEMM Tactic";
+    if (is_sm100 &&
+        tile_config_sm100 !=
+            cutlass_extensions::CutlassTileConfigSM100::ChooseWithHeuristic) {
+      assert(is_sm100 && !is_sm90 && "Invalid cutlass GEMM config: SM100");
+      tactic
+          << "\n\tstyle=TMA_SM100"  // Indicate SM100 specific TMA if applicable
+          << "\n\ttile shape ID: " << (int)tile_config_sm100
+          << "\n\tcluster shape ID: " << (int)cluster_shape
+          << "\n\tmainloop sched: " << (int)mainloop_schedule
+          << "\n\tepi sched: " << (int)epilogue_schedule;
+    } else if (is_sm90 && tile_config_sm90 !=
+                              cutlass_extensions::CutlassTileConfigSM90::
+                                  ChooseWithHeuristic) {
+      assert(is_sm90 && !is_sm100 && "Invalid cutlass GEMM config: SM90");
+      tactic << "\n\tstyle=TMA_SM90"
+             << "\n\ttile shape ID: " << (int)tile_config_sm90
+             << "\n\tcluster shape ID: " << (int)cluster_shape
+             << "\n\tmainloop sched: " << (int)mainloop_schedule
+             << "\n\tepi sched: " << (int)epilogue_schedule;
+    } else if (tile_config !=
+               cutlass_extensions::CutlassTileConfig::ChooseWithHeuristic) {
+      assert(!is_sm90 && !is_sm100 &&
+             "Invalid cutlass GEMM config: Compatible");
+      tactic << "\n\tstyle=compatible"
+             << "\n\ttile shape ID: " << (int)tile_config
+             << "\n\tstages: " << (int)stages
+             << "\n\tsplit_k_style: " << (int)split_k_style
+             << "\n\tsplit k: " << (int)split_k_factor;
+    } else {
+      tactic << "\n\tundefined";
     }
+    tactic << "\n";
+    return tactic.str();
+  }
 
-    // Constructor for SM90
-    CutlassGemmConfig(CutlassTileConfigSM90 tile_config_sm90_in, MainloopScheduleType mainloop_schedule_in,
-        EpilogueScheduleType epilogue_schedule_in, ClusterShape cluster_shape_in)
-        : tile_config_sm90(tile_config_sm90_in)
-        , mainloop_schedule(mainloop_schedule_in)
-        , epilogue_schedule(epilogue_schedule_in)
-        , cluster_shape(cluster_shape_in)
-        , is_sm90(true)
-        , is_sm100(false)
-    {
-    }
+  void fromString(const std::string& str) {
+    std::istringstream stream(str);
+    std::string line;
 
-    // Constructor for SM100 (Blackwell)
-    // Using existing MainloopScheduleType, EpilogueScheduleType, ClusterShape for now.
-    // These might need to be new SM100-specific types if Blackwell's TMA differs significantly.
-    CutlassGemmConfig(CutlassTileConfigSM100 tile_config_sm100_in, MainloopScheduleType mainloop_schedule_in,
-                      EpilogueScheduleType epilogue_schedule_in, ClusterShape cluster_shape_in)
-        : tile_config_sm100(tile_config_sm100_in)
-        , mainloop_schedule(mainloop_schedule_in) // Potentially use mainloop_schedule_sm100 if types diverge
-        , epilogue_schedule(epilogue_schedule_in) // Potentially use epilogue_schedule_sm100
-        , cluster_shape(cluster_shape_in)         // Potentially use cluster_shape_sm100
-        , is_sm90(false) // Explicitly false
-        , is_sm100(true)
-    {
-    }
+    is_sm90 = false;  // Reset flags
+    is_sm100 = false;
 
-
-    std::string toString() const
-    {
-        std::stringstream tactic;
-        tactic << "Cutlass GEMM Tactic";
-        if (is_sm100 && tile_config_sm100 != cutlass_extensions::CutlassTileConfigSM100::ChooseWithHeuristic)
-        {
-            assert(is_sm100 && !is_sm90 && "Invalid cutlass GEMM config: SM100");
-            tactic << "\n\tstyle=TMA_SM100" // Indicate SM100 specific TMA if applicable
-                   << "\n\ttile shape ID: " << (int) tile_config_sm100
-                   << "\n\tcluster shape ID: " << (int) cluster_shape
-                   << "\n\tmainloop sched: " << (int) mainloop_schedule
-                   << "\n\tepi sched: " << (int) epilogue_schedule;
-        }
-        else if (is_sm90 && tile_config_sm90 != cutlass_extensions::CutlassTileConfigSM90::ChooseWithHeuristic)
-        {
-            assert(is_sm90 && !is_sm100 && "Invalid cutlass GEMM config: SM90");
-            tactic << "\n\tstyle=TMA_SM90"
-                   << "\n\ttile shape ID: " << (int) tile_config_sm90
-                   << "\n\tcluster shape ID: " << (int) cluster_shape
-                   << "\n\tmainloop sched: " << (int) mainloop_schedule
-                   << "\n\tepi sched: " << (int) epilogue_schedule;
-        }
-        else if (tile_config != cutlass_extensions::CutlassTileConfig::ChooseWithHeuristic)
-        {
-            assert(!is_sm90 && !is_sm100 && "Invalid cutlass GEMM config: Compatible");
-            tactic << "\n\tstyle=compatible"
-                   << "\n\ttile shape ID: " << (int) tile_config
-                   << "\n\tstages: " << (int) stages
-                   << "\n\tsplit_k_style: " << (int) split_k_style
-                   << "\n\tsplit k: " << (int) split_k_factor;
-        }
-        else
-        {
-            tactic << "\n\tundefined";
-        }
-        tactic << "\n";
-        return tactic.str();
-    }
-
-    void fromString(const std::string& str) {
-        std::istringstream stream(str);
-        std::string line;
-
-        is_sm90 = false; // Reset flags
+    while (std::getline(stream, line)) {
+      if (line.find("style=TMA_SM100") != std::string::npos) {
+        is_sm100 = true;
+        is_sm90 = false;
+        std::getline(stream, line);
+        tile_config_sm100 =
+            static_cast<cutlass_extensions::CutlassTileConfigSM100>(
+                std::stoi(line.substr(line.find(':') + 1)));
+        std::getline(stream, line);
+        cluster_shape = static_cast<cutlass_extensions::ClusterShape>(
+            std::stoi(line.substr(line.find(':') + 1)));
+        std::getline(stream, line);
+        mainloop_schedule =
+            static_cast<cutlass_extensions::MainloopScheduleType>(
+                std::stoi(line.substr(line.find(':') + 1)));
+        std::getline(stream, line);
+        epilogue_schedule =
+            static_cast<cutlass_extensions::EpilogueScheduleType>(
+                std::stoi(line.substr(line.find(':') + 1)));
+      } else if (line.find("style=TMA_SM90") !=
+                 std::string::npos) {  // Check for SM90 specific first
+        is_sm90 = true;
         is_sm100 = false;
-
-        while (std::getline(stream, line)) {
-            if (line.find("style=TMA_SM100") != std::string::npos) {
-                is_sm100 = true;
-                is_sm90 = false;
-                std::getline(stream, line);
-                tile_config_sm100 = static_cast<cutlass_extensions::CutlassTileConfigSM100>(std::stoi(line.substr(line.find(':') + 1)));
-                std::getline(stream, line);
-                cluster_shape = static_cast<cutlass_extensions::ClusterShape>(std::stoi(line.substr(line.find(':') + 1)));
-                std::getline(stream, line);
-                mainloop_schedule = static_cast<cutlass_extensions::MainloopScheduleType>(std::stoi(line.substr(line.find(':') + 1)));
-                std::getline(stream, line);
-                epilogue_schedule = static_cast<cutlass_extensions::EpilogueScheduleType>(std::stoi(line.substr(line.find(':') + 1)));
-            } else if (line.find("style=TMA_SM90") != std::string::npos) { // Check for SM90 specific first
-                is_sm90 = true;
-                is_sm100 = false;
-                std::getline(stream, line);
-                tile_config_sm90 = static_cast<cutlass_extensions::CutlassTileConfigSM90>(std::stoi(line.substr(line.find(':') + 1)));
-                std::getline(stream, line);
-                cluster_shape = static_cast<cutlass_extensions::ClusterShape>(std::stoi(line.substr(line.find(':') + 1)));
-                std::getline(stream, line);
-                mainloop_schedule = static_cast<cutlass_extensions::MainloopScheduleType>(std::stoi(line.substr(line.find(':') + 1)));
-                std::getline(stream, line);
-                epilogue_schedule = static_cast<cutlass_extensions::EpilogueScheduleType>(std::stoi(line.substr(line.find(':') + 1)));
-            } else if (line.find("style=compatible") != std::string::npos) {
-                is_sm90 = false;
-                is_sm100 = false;
-                std::getline(stream, line);
-                tile_config = static_cast<cutlass_extensions::CutlassTileConfig>(std::stoi(line.substr(line.find(':') + 1)));
-                std::getline(stream, line);
-                stages = std::stoi(line.substr(line.find(':') + 1));
-                std::getline(stream, line);
-                split_k_style = static_cast<cutlass_extensions::SplitKStyle>(std::stoi(line.substr(line.find(':') + 1)));
-                std::getline(stream, line);
-                split_k_factor = std::stoi(line.substr(line.find(':') + 1));
-            }
-        }
+        std::getline(stream, line);
+        tile_config_sm90 =
+            static_cast<cutlass_extensions::CutlassTileConfigSM90>(
+                std::stoi(line.substr(line.find(':') + 1)));
+        std::getline(stream, line);
+        cluster_shape = static_cast<cutlass_extensions::ClusterShape>(
+            std::stoi(line.substr(line.find(':') + 1)));
+        std::getline(stream, line);
+        mainloop_schedule =
+            static_cast<cutlass_extensions::MainloopScheduleType>(
+                std::stoi(line.substr(line.find(':') + 1)));
+        std::getline(stream, line);
+        epilogue_schedule =
+            static_cast<cutlass_extensions::EpilogueScheduleType>(
+                std::stoi(line.substr(line.find(':') + 1)));
+      } else if (line.find("style=compatible") != std::string::npos) {
+        is_sm90 = false;
+        is_sm100 = false;
+        std::getline(stream, line);
+        tile_config = static_cast<cutlass_extensions::CutlassTileConfig>(
+            std::stoi(line.substr(line.find(':') + 1)));
+        std::getline(stream, line);
+        stages = std::stoi(line.substr(line.find(':') + 1));
+        std::getline(stream, line);
+        split_k_style = static_cast<cutlass_extensions::SplitKStyle>(
+            std::stoi(line.substr(line.find(':') + 1)));
+        std::getline(stream, line);
+        split_k_factor = std::stoi(line.substr(line.find(':') + 1));
+      }
     }
+  }
 };
 
-inline std::ostream& operator<<(std::ostream& out, CutlassGemmConfig const& config)
-{
-    // clang-format off
+inline std::ostream& operator<<(std::ostream& out,
+                                CutlassGemmConfig const& config) {
+  // clang-format off
     if (config.is_sm100)
     {
         out << "tile_config_sm100_enum: " << int(config.tile_config_sm100)
@@ -337,8 +358,8 @@ inline std::ostream& operator<<(std::ostream& out, CutlassGemmConfig const& conf
             << ", split_k_factor: " << config.split_k_factor
             << ", stages: " << config.stages;
     }
-    // clang-format on
-    return out;
+  // clang-format on
+  return out;
 }
 
-} // namespace cutlass_extensions
+}  // namespace cutlass_extensions
diff --git a/custom_ops/gpu_ops/cutlass_extensions/interleaved_numeric_conversion.h b/custom_ops/gpu_ops/cutlass_extensions/interleaved_numeric_conversion.h
index e7e17657be..9a9b35a324 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/interleaved_numeric_conversion.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/interleaved_numeric_conversion.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,19 +18,21 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*!
     \file
-    \brief Boost-like numeric conversion operator for int8 and CUTLASS int4b_t interleaved in a register
+    \brief Boost-like numeric conversion operator for int8 and CUTLASS int4b_t
+   interleaved in a register
 */
 
 #pragma once
@@ -52,726 +54,801 @@ __device__ inline int lop3(int a, int b, int c) {
   return res;
 }
 
-// This converter is meant to be used with data interleaved in a 32-bit register where the even elements are in the low
-// bits and the odd elemeents are in the high bits of the register. In addition, it assumes elements were originally
-// signed and had a bias of 2**(b-1) added (where b is the number of bits in the type) to make all numbers unsigned.
-// This converter will uninterleave the data and subtract the bias while converting to the result type.
+// This converter is meant to be used with data interleaved in a 32-bit register
+// where the even elements are in the low bits and the odd elemeents are in the
+// high bits of the register. In addition, it assumes elements were originally
+// signed and had a bias of 2**(b-1) added (where b is the number of bits in the
+// type) to make all numbers unsigned. This converter will uninterleave the data
+// and subtract the bias while converting to the result type.
 template <typename T, typename S, int N>
 struct FastInterleavedAndBiasedNumericArrayConverter;
 
 template <>
-struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint8_t, 4>
-{
-    using result_type = Array<half_t, 4>;
-    using source_type = Array<uint8_t, 4>;
+struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint8_t, 4> {
+  using result_type = Array<half_t, 4>;
+  using source_type = Array<uint8_t, 4>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 
-        uint32_t* h = reinterpret_cast<uint32_t*>(&result);
-        uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
 
-        static constexpr uint32_t mask_for_elt_01 = 0x5250;
-        static constexpr uint32_t mask_for_elt_23 = 0x5351;
-        static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[0]) : "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_01));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[1]) : "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_23));
+    static constexpr uint32_t mask_for_elt_01 = 0x5250;
+    static constexpr uint32_t mask_for_elt_23 = 0x5351;
+    static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[0])
+                 : "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_01));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[1])
+                 : "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_23));
 
-        // Lastly, we subtract 1152 from our constructed number using fp16 math to get our signed integer as fp16.
-        static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(I8s_TO_F16s_MAGIC_NUM));
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(I8s_TO_F16s_MAGIC_NUM));
+    // Lastly, we subtract 1152 from our constructed number using fp16 math to
+    // get our signed integer as fp16.
+    static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
+    asm volatile("sub.f16x2 %0, %1, %2;\n"
+                 : "=r"(h[0])
+                 : "r"(h[0]), "r"(I8s_TO_F16s_MAGIC_NUM));
+    asm volatile("sub.f16x2 %0, %1, %2;\n"
+                 : "=r"(h[1])
+                 : "r"(h[1]), "r"(I8s_TO_F16s_MAGIC_NUM));
 
-        return result;
-    }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 template <int N>
-struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint8_t, N>
-{
-    static constexpr int VEC_WIDTH = 4;
-    static_assert(!(N % VEC_WIDTH), "N must be multiple of 4.");
+struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint8_t, N> {
+  static constexpr int VEC_WIDTH = 4;
+  static_assert(!(N % VEC_WIDTH), "N must be multiple of 4.");
 
-    using result_type = Array<half_t, N>;
-    using source_type = Array<uint8_t, N>;
+  using result_type = Array<half_t, N>;
+  using source_type = Array<uint8_t, N>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type, scalar_source_type, VEC_WIDTH>
-            convert_vector_;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type,
+                                                  scalar_source_type,
+                                                  VEC_WIDTH>
+        convert_vector_;
 
-        result_type result;
-        using vec_result = Array<scalar_result_type, VEC_WIDTH>;
-        using vec_source = Array<scalar_source_type, VEC_WIDTH>;
+    result_type result;
+    using vec_result = Array<scalar_result_type, VEC_WIDTH>;
+    using vec_source = Array<scalar_source_type, VEC_WIDTH>;
 
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / VEC_WIDTH; ++i)
-        {
-            result_ptr[i] = convert_vector_(source_ptr[i]);
-        }
-
-        return result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / VEC_WIDTH; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i]);
     }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 template <>
-struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint8_t, 4>
-{
-    using result_type = Array<bfloat16_t, 4>;
-    using source_type = Array<uint8_t, 4>;
+struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint8_t, 4> {
+  using result_type = Array<bfloat16_t, 4>;
+  using source_type = Array<uint8_t, 4>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
 
-        uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(&result);
-        uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+    uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(&result);
+    uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
 
-        static constexpr uint32_t fp32_base = 0x4B000000;
-        float fp32_intermediates[4];
+    static constexpr uint32_t fp32_base = 0x4B000000;
+    float fp32_intermediates[4];
 
-        // Construct FP32s, bfloat does not have enough mantissa for IADD trick
-        uint32_t* fp32_intermediates_casted = reinterpret_cast<uint32_t*>(fp32_intermediates);
-        fp32_intermediates_casted[0] = __byte_perm(i8s, fp32_base, 0x7650);
-        fp32_intermediates_casted[1] = __byte_perm(i8s, fp32_base, 0x7652);
-        fp32_intermediates_casted[2] = __byte_perm(i8s, fp32_base, 0x7651);
-        fp32_intermediates_casted[3] = __byte_perm(i8s, fp32_base, 0x7653);
+    // Construct FP32s, bfloat does not have enough mantissa for IADD trick
+    uint32_t* fp32_intermediates_casted =
+        reinterpret_cast<uint32_t*>(fp32_intermediates);
+    fp32_intermediates_casted[0] = __byte_perm(i8s, fp32_base, 0x7650);
+    fp32_intermediates_casted[1] = __byte_perm(i8s, fp32_base, 0x7652);
+    fp32_intermediates_casted[2] = __byte_perm(i8s, fp32_base, 0x7651);
+    fp32_intermediates_casted[3] = __byte_perm(i8s, fp32_base, 0x7653);
 
-        // Subtract out fp32_base + 128 to make the unsigned integer signed.
-        CUTLASS_PRAGMA_UNROLL
-        for (int ii = 0; ii < 4; ++ii)
-        {
-            fp32_intermediates[ii] -= 8388736.f;
-        }
+    // Subtract out fp32_base + 128 to make the unsigned integer signed.
+    CUTLASS_PRAGMA_UNROLL
+    for (int ii = 0; ii < 4; ++ii) {
+      fp32_intermediates[ii] -= 8388736.f;
+    }
 
-        // Truncate the fp32 representation and pack up as bfloat16s.
-        CUTLASS_PRAGMA_UNROLL
-        for (int ii = 0; ii < 2; ++ii)
-        {
-            bf16_result_ptr[ii]
-                = __byte_perm(fp32_intermediates_casted[2 * ii + 0], fp32_intermediates_casted[2 * ii + 1], 0x7632);
-        }
+    // Truncate the fp32 representation and pack up as bfloat16s.
+    CUTLASS_PRAGMA_UNROLL
+    for (int ii = 0; ii < 2; ++ii) {
+      bf16_result_ptr[ii] = __byte_perm(fp32_intermediates_casted[2 * ii + 0],
+                                        fp32_intermediates_casted[2 * ii + 1],
+                                        0x7632);
+    }
 #else
-        // Disable this on architectures older than Ampere since they lack hardware for bf16 mma. If one wishes to use
-        // HMMA on older hardware, they should Convert directly to FP16 using FP16 converters.
-        result.clear(); // Suppress compiler warning
-        arch::device_breakpoint();
+    // Disable this on architectures older than Ampere since they lack hardware
+    // for bf16 mma. If one wishes to use HMMA on older hardware, they should
+    // Convert directly to FP16 using FP16 converters.
+    result.clear();  // Suppress compiler warning
+    arch::device_breakpoint();
 #endif
-        return result;
-    }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 template <int N>
-struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint8_t, N>
-{
-    static constexpr int VEC_WIDTH = 4;
-    static_assert(!(N % VEC_WIDTH), "N must be multiple of 4.");
+struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint8_t, N> {
+  static constexpr int VEC_WIDTH = 4;
+  static_assert(!(N % VEC_WIDTH), "N must be multiple of 4.");
 
-    using result_type = Array<bfloat16_t, N>;
-    using source_type = Array<uint8_t, N>;
+  using result_type = Array<bfloat16_t, N>;
+  using source_type = Array<uint8_t, N>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type, scalar_source_type, VEC_WIDTH>
-            convert_vector_;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type,
+                                                  scalar_source_type,
+                                                  VEC_WIDTH>
+        convert_vector_;
 
-        result_type result;
-        using vec_result = Array<scalar_result_type, VEC_WIDTH>;
-        using vec_source = Array<scalar_source_type, VEC_WIDTH>;
+    result_type result;
+    using vec_result = Array<scalar_result_type, VEC_WIDTH>;
+    using vec_source = Array<scalar_source_type, VEC_WIDTH>;
 
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / VEC_WIDTH; ++i)
-        {
-            result_ptr[i] = convert_vector_(source_ptr[i]);
-        }
-
-        return result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / VEC_WIDTH; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i]);
     }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 template <>
-struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint4b_t, 8>
-{
-    using result_type = Array<half_t, 8>;
-    using source_type = Array<uint4b_t, 8>;
+struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint4b_t, 8> {
+  using result_type = Array<half_t, 8>;
+  using source_type = Array<uint4b_t, 8>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 
-        uint32_t* h = reinterpret_cast<uint32_t*>(&result);
-        uint32_t const i4s = reinterpret_cast<uint32_t const&>(source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    uint32_t const i4s = reinterpret_cast<uint32_t const&>(source);
 
-        // First, we extract the i4s and construct an intermediate fp16 number.
-        static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
-        static constexpr uint32_t BOTTOM_MASK = 0x000f000f;
-        static constexpr uint32_t TOP_MASK = 0x00f000f0;
-        static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400;
+    // First, we extract the i4s and construct an intermediate fp16 number.
+    static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
+    static constexpr uint32_t BOTTOM_MASK = 0x000f000f;
+    static constexpr uint32_t TOP_MASK = 0x00f000f0;
+    static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400;
 
-        // Note that the entire sequence only requires 1 shift instruction. This is thanks to the register packing
-        // format and the fact that we force our integers to be unsigned, and account for this in the fp16 subtractions.
-        // In addition, I exploit the fact that sub and fma have the same throughput in order to convert elt_23 and
-        // elt_67 to fp16 without having to shift them to the bottom bits before hand.
+    // Note that the entire sequence only requires 1 shift instruction. This is
+    // thanks to the register packing format and the fact that we force our
+    // integers to be unsigned, and account for this in the fp16 subtractions.
+    // In addition, I exploit the fact that sub and fma have the same throughput
+    // in order to convert elt_23 and elt_67 to fp16 without having to shift
+    // them to the bottom bits before hand.
 
-        // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW dependency if we issue
-        // immediately before required.
-        const uint32_t top_i4s = i4s >> 8;
-        // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400
-        asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
-                     : "=r"(h[0])
-                     : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
-        // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400
-        asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
-                     : "=r"(h[1])
-                     : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
-        // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400
-        asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
-                     : "=r"(h[2])
-                     : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
-        // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400
-        asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
-                     : "=r"(h[3])
-                     : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
+    // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide
+    // RAW dependency if we issue immediately before required.
+    const uint32_t top_i4s = i4s >> 8;
+    // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400
+    asm volatile(
+        "lop3.b32 %0, %1, %2, %3, %4;\n"
+        : "=r"(h[0])
+        : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
+    // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400
+    asm volatile(
+        "lop3.b32 %0, %1, %2, %3, %4;\n"
+        : "=r"(h[1])
+        : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
+    // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400
+    asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+                 : "=r"(h[2])
+                 : "r"(top_i4s),
+                   "n"(BOTTOM_MASK),
+                   "n"(I4s_TO_F16s_MAGIC_NUM),
+                   "n"(immLut));
+    // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400
+    asm volatile(
+        "lop3.b32 %0, %1, %2, %3, %4;\n"
+        : "=r"(h[3])
+        : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
 
-        // I use inline PTX below because I am not sure if the compiler will emit float2half instructions if I use the
-        // half2 ctor. In this case, I chose performance reliability over code readability.
+    // I use inline PTX below because I am not sure if the compiler will emit
+    // float2half instructions if I use the half2 ctor. In this case, I chose
+    // performance reliability over code readability.
 
-        // This is the half2 {1032, 1032} represented as an integer.
-        static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408;
-        // This is the half2 {1 / 16, 1 / 16} represented as an integer.
-        static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00;
-        // This is the half2 {-72, -72} represented as an integer.
-        static constexpr uint32_t NEG_72 = 0xd480d480;
+    // This is the half2 {1032, 1032} represented as an integer.
+    static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408;
+    // This is the half2 {1 / 16, 1 / 16} represented as an integer.
+    static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00;
+    // This is the half2 {-72, -72} represented as an integer.
+    static constexpr uint32_t NEG_72 = 0xd480d480;
 
-        // Finally, we construct the output numbers.
-        // Convert elt_01
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM));
-        // Convert elt_23
-        asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_72));
-        // Convert elt_45
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM));
-        // Convert elt_67
-        asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_72));
+    // Finally, we construct the output numbers.
+    // Convert elt_01
+    asm volatile("sub.f16x2 %0, %1, %2;\n"
+                 : "=r"(h[0])
+                 : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM));
+    // Convert elt_23
+    asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[1])
+                 : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_72));
+    // Convert elt_45
+    asm volatile("sub.f16x2 %0, %1, %2;\n"
+                 : "=r"(h[2])
+                 : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM));
+    // Convert elt_67
+    asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[3])
+                 : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_72));
 
-        return result;
-    }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 template <int N>
-struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint4b_t, N>
-{
-    static constexpr int VEC_WIDTH = 8;
-    static_assert(!(N % VEC_WIDTH), "N must be multiple of 8.");
+struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint4b_t, N> {
+  static constexpr int VEC_WIDTH = 8;
+  static_assert(!(N % VEC_WIDTH), "N must be multiple of 8.");
 
-    using result_type = Array<half_t, N>;
-    using source_type = Array<uint4b_t, N>;
+  using result_type = Array<half_t, N>;
+  using source_type = Array<uint4b_t, N>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type, scalar_source_type, VEC_WIDTH>
-            convert_vector_;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type,
+                                                  scalar_source_type,
+                                                  VEC_WIDTH>
+        convert_vector_;
 
-        result_type result;
-        using vec_result = Array<scalar_result_type, VEC_WIDTH>;
-        using vec_source = Array<scalar_source_type, VEC_WIDTH>;
+    result_type result;
+    using vec_result = Array<scalar_result_type, VEC_WIDTH>;
+    using vec_source = Array<scalar_source_type, VEC_WIDTH>;
 
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / VEC_WIDTH; ++i)
-        {
-            result_ptr[i] = convert_vector_(source_ptr[i]);
-        }
-
-        return result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / VEC_WIDTH; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i]);
     }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 template <>
-struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint4b_t, 8>
-{
-    using result_type = Array<bfloat16_t, 8>;
-    using source_type = Array<uint4b_t, 8>;
+struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint4b_t, 8> {
+  using result_type = Array<bfloat16_t, 8>;
+  using source_type = Array<uint4b_t, 8>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
 
-        uint32_t* h = reinterpret_cast<uint32_t*>(&result);
-        uint32_t const source_i4s = reinterpret_cast<uint32_t const&>(source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    uint32_t const source_i4s = reinterpret_cast<uint32_t const&>(source);
 
-        // First, we extract the i4s and construct an intermediate fp16 number.
-        static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
-        static constexpr uint32_t MASK = 0x000f000f;
-        static constexpr uint32_t I4s_TO_BF16s_MAGIC_NUM = 0x43004300;
+    // First, we extract the i4s and construct an intermediate fp16 number.
+    static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
+    static constexpr uint32_t MASK = 0x000f000f;
+    static constexpr uint32_t I4s_TO_BF16s_MAGIC_NUM = 0x43004300;
 
-        // We don't have enough mantissa to remove as much shift overhead as FP16, so we must loop.
-        // No shift needed for first item.
-        uint32_t i4s = source_i4s;
-        asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
-                     : "=r"(h[0])
-                     : "r"(i4s), "n"(MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut));
-        CUTLASS_PRAGMA_UNROLL
-        for (int ii = 1; ii < result_type::kElements / 2; ++ii)
-        {
-            i4s >>= sizeof_bits<typename source_type::Element>::value;
-            // (i4s & 0x000f000f) | 0x43004300
-            asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
-                         : "=r"(h[ii])
-                         : "r"(i4s), "n"(MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut));
-        }
+    // We don't have enough mantissa to remove as much shift overhead as FP16,
+    // so we must loop. No shift needed for first item.
+    uint32_t i4s = source_i4s;
+    asm volatile(
+        "lop3.b32 %0, %1, %2, %3, %4;\n"
+        : "=r"(h[0])
+        : "r"(i4s), "n"(MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut));
+    CUTLASS_PRAGMA_UNROLL
+    for (int ii = 1; ii < result_type::kElements / 2; ++ii) {
+      i4s >>= sizeof_bits<typename source_type::Element>::value;
+      // (i4s & 0x000f000f) | 0x43004300
+      asm volatile(
+          "lop3.b32 %0, %1, %2, %3, %4;\n"
+          : "=r"(h[ii])
+          : "r"(i4s), "n"(MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut));
+    }
 
-        // This is the BF16 {-136, -136} represented as an integer.
-        static constexpr uint32_t BF16_BIAS = 0xC308C308;
-        static constexpr uint32_t BF16_ONE = 0x3F803F80;
+    // This is the BF16 {-136, -136} represented as an integer.
+    static constexpr uint32_t BF16_BIAS = 0xC308C308;
+    static constexpr uint32_t BF16_ONE = 0x3F803F80;
 
-        // Finally, we construct the output numbers.
-        CUTLASS_PRAGMA_UNROLL
-        for (int ii = 0; ii < result_type::kElements / 2; ++ii)
-        {
-            // Since this section is for Ampere+, we use bf16 fma to do the bias subtraction
-            asm("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[ii]) : "r"(h[ii]), "r"(BF16_ONE), "r"(BF16_BIAS));
-        }
+    // Finally, we construct the output numbers.
+    CUTLASS_PRAGMA_UNROLL
+    for (int ii = 0; ii < result_type::kElements / 2; ++ii) {
+      // Since this section is for Ampere+, we use bf16 fma to do the bias
+      // subtraction
+      asm("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+          : "=r"(h[ii])
+          : "r"(h[ii]), "r"(BF16_ONE), "r"(BF16_BIAS));
+    }
 #else
-        // Disable this on architectures older than Ampere since they lack hardware for bf16 mma. If one wishes to use
-        // HMMA on older hardware, they should Convert directly to FP16 using FP16 converters.
-        arch::device_breakpoint();
-        result.clear(); // Suppress compiler warning.
+    // Disable this on architectures older than Ampere since they lack hardware
+    // for bf16 mma. If one wishes to use HMMA on older hardware, they should
+    // Convert directly to FP16 using FP16 converters.
+    arch::device_breakpoint();
+    result.clear();  // Suppress compiler warning.
 #endif
-        return result;
-    }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 template <int N>
-struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint4b_t, N>
-{
-    static constexpr int VEC_WIDTH = 8;
-    static_assert(!(N % VEC_WIDTH), "N must be multiple of 8.");
+struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint4b_t, N> {
+  static constexpr int VEC_WIDTH = 8;
+  static_assert(!(N % VEC_WIDTH), "N must be multiple of 8.");
 
-    using result_type = Array<bfloat16_t, N>;
-    using source_type = Array<uint4b_t, N>;
+  using result_type = Array<bfloat16_t, N>;
+  using source_type = Array<uint4b_t, N>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type, scalar_source_type, VEC_WIDTH>
-            convert_vector_;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type,
+                                                  scalar_source_type,
+                                                  VEC_WIDTH>
+        convert_vector_;
 
-        result_type result;
-        using vec_result = Array<scalar_result_type, VEC_WIDTH>;
-        using vec_source = Array<scalar_source_type, VEC_WIDTH>;
+    result_type result;
+    using vec_result = Array<scalar_result_type, VEC_WIDTH>;
+    using vec_source = Array<scalar_source_type, VEC_WIDTH>;
 
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / VEC_WIDTH; ++i)
-        {
-            result_ptr[i] = convert_vector_(source_ptr[i]);
-        }
-
-        return result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / VEC_WIDTH; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i]);
     }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 template <>
-struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint2b_t, 16>
-{
-    using result_type = Array<half_t, 16>;
-    using source_type = Array<uint2b_t, 16>;
+struct FastInterleavedAndBiasedNumericArrayConverter<half_t, uint2b_t, 16> {
+  using result_type = Array<half_t, 16>;
+  using source_type = Array<uint2b_t, 16>;
 
-    using ScaleComputeT = float;
-    using code_type = Array<ScaleComputeT, 4>;
+  using ScaleComputeT = float;
+  using code_type = Array<ScaleComputeT, 4>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source, ScaleComputeT code_scale, ScaleComputeT code_zp)
-    {
-        uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source,
+                             ScaleComputeT code_scale,
+                             ScaleComputeT code_zp) {
+    uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
 
-        // 2^23 = 8388608
-        static constexpr uint32_t FP32_BASE = 0x4B000000;
+    // 2^23 = 8388608
+    static constexpr uint32_t FP32_BASE = 0x4B000000;
 
-        float fp32_intermediates[4];
-        uint32_t* fp32_intermediates_casted = reinterpret_cast<uint32_t*>(fp32_intermediates);
-        fp32_intermediates_casted[0] = __byte_perm(i8s, FP32_BASE, 0x7650);
-        fp32_intermediates_casted[1] = __byte_perm(i8s, FP32_BASE, 0x7651);
-        fp32_intermediates_casted[2] = __byte_perm(i8s, FP32_BASE, 0x7652);
-        fp32_intermediates_casted[3] = __byte_perm(i8s, FP32_BASE, 0x7653);
+    float fp32_intermediates[4];
+    uint32_t* fp32_intermediates_casted =
+        reinterpret_cast<uint32_t*>(fp32_intermediates);
+    fp32_intermediates_casted[0] = __byte_perm(i8s, FP32_BASE, 0x7650);
+    fp32_intermediates_casted[1] = __byte_perm(i8s, FP32_BASE, 0x7651);
+    fp32_intermediates_casted[2] = __byte_perm(i8s, FP32_BASE, 0x7652);
+    fp32_intermediates_casted[3] = __byte_perm(i8s, FP32_BASE, 0x7653);
 
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[0]) : "r"(fp32_intermediates_casted[0]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[1]) : "r"(fp32_intermediates_casted[1]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[2]) : "r"(fp32_intermediates_casted[2]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[3]) : "r"(fp32_intermediates_casted[3]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[0])
+                 : "r"(fp32_intermediates_casted[0]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[1])
+                 : "r"(fp32_intermediates_casted[1]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[2])
+                 : "r"(fp32_intermediates_casted[2]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[3])
+                 : "r"(fp32_intermediates_casted[3]), "r"(FP32_BASE));
 
-        int32_t decode_value[4];
-        ScaleComputeT new_code_zp = code_zp + 0.5f;
+    int32_t decode_value[4];
+    ScaleComputeT new_code_zp = code_zp + 0.5f;
 
-        decode_value[0] = __float2int_rd(fmaf(fp32_intermediates[0], code_scale, new_code_zp));
-        decode_value[1] = __float2int_rd(fmaf(fp32_intermediates[1], code_scale, new_code_zp));
-        decode_value[2] = __float2int_rd(fmaf(fp32_intermediates[2], code_scale, new_code_zp));
-        decode_value[3] = __float2int_rd(fmaf(fp32_intermediates[3], code_scale, new_code_zp));
+    decode_value[0] =
+        __float2int_rd(fmaf(fp32_intermediates[0], code_scale, new_code_zp));
+    decode_value[1] =
+        __float2int_rd(fmaf(fp32_intermediates[1], code_scale, new_code_zp));
+    decode_value[2] =
+        __float2int_rd(fmaf(fp32_intermediates[2], code_scale, new_code_zp));
+    decode_value[3] =
+        __float2int_rd(fmaf(fp32_intermediates[3], code_scale, new_code_zp));
 
-        return convert_impl(decode_value);
-    }
+    return convert_impl(decode_value);
+  }
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source, code_type const& code_scale, code_type const& code_zp)
-    {
-        uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source,
+                             code_type const& code_scale,
+                             code_type const& code_zp) {
+    uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
 
-        // 2^23 = 8388608
-        static constexpr uint32_t FP32_BASE = 0x4B000000;
+    // 2^23 = 8388608
+    static constexpr uint32_t FP32_BASE = 0x4B000000;
 
-        float fp32_intermediates[4];
-        uint32_t* fp32_intermediates_casted = reinterpret_cast<uint32_t*>(fp32_intermediates);
-        fp32_intermediates_casted[0] = __byte_perm(i8s, FP32_BASE, 0x7650);
-        fp32_intermediates_casted[1] = __byte_perm(i8s, FP32_BASE, 0x7651);
-        fp32_intermediates_casted[2] = __byte_perm(i8s, FP32_BASE, 0x7652);
-        fp32_intermediates_casted[3] = __byte_perm(i8s, FP32_BASE, 0x7653);
+    float fp32_intermediates[4];
+    uint32_t* fp32_intermediates_casted =
+        reinterpret_cast<uint32_t*>(fp32_intermediates);
+    fp32_intermediates_casted[0] = __byte_perm(i8s, FP32_BASE, 0x7650);
+    fp32_intermediates_casted[1] = __byte_perm(i8s, FP32_BASE, 0x7651);
+    fp32_intermediates_casted[2] = __byte_perm(i8s, FP32_BASE, 0x7652);
+    fp32_intermediates_casted[3] = __byte_perm(i8s, FP32_BASE, 0x7653);
 
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[0]) : "r"(fp32_intermediates_casted[0]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[1]) : "r"(fp32_intermediates_casted[1]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[2]) : "r"(fp32_intermediates_casted[2]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[3]) : "r"(fp32_intermediates_casted[3]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[0])
+                 : "r"(fp32_intermediates_casted[0]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[1])
+                 : "r"(fp32_intermediates_casted[1]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[2])
+                 : "r"(fp32_intermediates_casted[2]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[3])
+                 : "r"(fp32_intermediates_casted[3]), "r"(FP32_BASE));
 
-        int32_t decode_value[4];
+    int32_t decode_value[4];
 
-        decode_value[0] = __float2int_rd(fmaf(fp32_intermediates[0], code_scale[0], code_zp[0] + 0.5f));
-        decode_value[1] = __float2int_rd(fmaf(fp32_intermediates[1], code_scale[1], code_zp[1] + 0.5f));
-        decode_value[2] = __float2int_rd(fmaf(fp32_intermediates[2], code_scale[2], code_zp[2] + 0.5f));
-        decode_value[3] = __float2int_rd(fmaf(fp32_intermediates[3], code_scale[3], code_zp[3] + 0.5f));
+    decode_value[0] = __float2int_rd(
+        fmaf(fp32_intermediates[0], code_scale[0], code_zp[0] + 0.5f));
+    decode_value[1] = __float2int_rd(
+        fmaf(fp32_intermediates[1], code_scale[1], code_zp[1] + 0.5f));
+    decode_value[2] = __float2int_rd(
+        fmaf(fp32_intermediates[2], code_scale[2], code_zp[2] + 0.5f));
+    decode_value[3] = __float2int_rd(
+        fmaf(fp32_intermediates[3], code_scale[3], code_zp[3] + 0.5f));
 
-        return convert_impl(decode_value);
-    }
+    return convert_impl(decode_value);
+  }
 
-    CUTLASS_DEVICE
-    static result_type convert_impl(int32_t* decode_value)
-    {
-        result_type result;
-        static constexpr uint32_t immLut = (0xF0 & 0xCC) | 0xAA;
+  CUTLASS_DEVICE
+  static result_type convert_impl(int32_t* decode_value) {
+    result_type result;
+    static constexpr uint32_t immLut = (0xF0 & 0xCC) | 0xAA;
 
-        static constexpr uint32_t MASK = 0x003F003F;
-        // 2^10 = 1024
-        static constexpr uint32_t EX = 0x64006400;
+    static constexpr uint32_t MASK = 0x003F003F;
+    // 2^10 = 1024
+    static constexpr uint32_t EX = 0x64006400;
 
-        uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
 
-        int32_t q0 = __byte_perm(decode_value[0], decode_value[1], 0x5410);
-        int32_t q1 = __byte_perm(decode_value[2], decode_value[3], 0x5410);
+    int32_t q0 = __byte_perm(decode_value[0], decode_value[1], 0x5410);
+    int32_t q1 = __byte_perm(decode_value[2], decode_value[3], 0x5410);
 
-        h[0] = lop3<immLut>(q0 >> 9, MASK, EX);
-        h[1] = lop3<immLut>(q0 >> 6, MASK, EX);
-        h[2] = lop3<immLut>(q0 >> 3, MASK, EX);
-        h[3] = lop3<immLut>(q0, MASK, EX);
+    h[0] = lop3<immLut>(q0 >> 9, MASK, EX);
+    h[1] = lop3<immLut>(q0 >> 6, MASK, EX);
+    h[2] = lop3<immLut>(q0 >> 3, MASK, EX);
+    h[3] = lop3<immLut>(q0, MASK, EX);
 
-        h[4] = lop3<immLut>(q1 >> 9, MASK, EX);
-        h[5] = lop3<immLut>(q1 >> 6, MASK, EX);
-        h[6] = lop3<immLut>(q1 >> 3, MASK, EX);
-        h[7] = lop3<immLut>(q1, MASK, EX);
+    h[4] = lop3<immLut>(q1 >> 9, MASK, EX);
+    h[5] = lop3<immLut>(q1 >> 6, MASK, EX);
+    h[6] = lop3<immLut>(q1 >> 3, MASK, EX);
+    h[7] = lop3<immLut>(q1, MASK, EX);
 
-        // 1024 + 32 = 1056
-        static constexpr uint32_t SUB = 0x64206420;
+    // 1024 + 32 = 1056
+    static constexpr uint32_t SUB = 0x64206420;
 
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(SUB));
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(SUB));
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(SUB));
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[3]) : "r"(h[3]), "r"(SUB));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(SUB));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(SUB));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(SUB));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[3]) : "r"(h[3]), "r"(SUB));
 
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[4]) : "r"(h[4]), "r"(SUB));
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[5]) : "r"(h[5]), "r"(SUB));
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[6]) : "r"(h[6]), "r"(SUB));
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[7]) : "r"(h[7]), "r"(SUB));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[4]) : "r"(h[4]), "r"(SUB));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[5]) : "r"(h[5]), "r"(SUB));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[6]) : "r"(h[6]), "r"(SUB));
+    asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[7]) : "r"(h[7]), "r"(SUB));
 
-        return result;
-    }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s, ScaleComputeT code_scale, ScaleComputeT code_zp)
-    {
-        return convert(s, code_scale, code_zp);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s,
+                         ScaleComputeT code_scale,
+                         ScaleComputeT code_zp) {
+    return convert(s, code_scale, code_zp);
+  }
 };
 
 template <>
-struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint2b_t, 16>
-{
-    using result_type = Array<bfloat16_t, 16>;
-    using source_type = Array<uint2b_t, 16>;
+struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint2b_t, 16> {
+  using result_type = Array<bfloat16_t, 16>;
+  using source_type = Array<uint2b_t, 16>;
 
-    using ScaleComputeT = float;
-    using code_type = Array<ScaleComputeT, 4>;
+  using ScaleComputeT = float;
+  using code_type = Array<ScaleComputeT, 4>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source, ScaleComputeT code_scale, ScaleComputeT code_zp)
-    {
-        uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source,
+                             ScaleComputeT code_scale,
+                             ScaleComputeT code_zp) {
+    uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
 
-        // 2^23 = 8388608
-        static constexpr uint32_t FP32_BASE = 0x4B000000;
+    // 2^23 = 8388608
+    static constexpr uint32_t FP32_BASE = 0x4B000000;
 
-        float fp32_intermediates[4];
-        uint32_t* fp32_intermediates_casted = reinterpret_cast<uint32_t*>(fp32_intermediates);
-        fp32_intermediates_casted[0] = __byte_perm(i8s, FP32_BASE, 0x7650);
-        fp32_intermediates_casted[1] = __byte_perm(i8s, FP32_BASE, 0x7651);
-        fp32_intermediates_casted[2] = __byte_perm(i8s, FP32_BASE, 0x7652);
-        fp32_intermediates_casted[3] = __byte_perm(i8s, FP32_BASE, 0x7653);
+    float fp32_intermediates[4];
+    uint32_t* fp32_intermediates_casted =
+        reinterpret_cast<uint32_t*>(fp32_intermediates);
+    fp32_intermediates_casted[0] = __byte_perm(i8s, FP32_BASE, 0x7650);
+    fp32_intermediates_casted[1] = __byte_perm(i8s, FP32_BASE, 0x7651);
+    fp32_intermediates_casted[2] = __byte_perm(i8s, FP32_BASE, 0x7652);
+    fp32_intermediates_casted[3] = __byte_perm(i8s, FP32_BASE, 0x7653);
 
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[0]) : "r"(fp32_intermediates_casted[0]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[1]) : "r"(fp32_intermediates_casted[1]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[2]) : "r"(fp32_intermediates_casted[2]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[3]) : "r"(fp32_intermediates_casted[3]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[0])
+                 : "r"(fp32_intermediates_casted[0]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[1])
+                 : "r"(fp32_intermediates_casted[1]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[2])
+                 : "r"(fp32_intermediates_casted[2]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[3])
+                 : "r"(fp32_intermediates_casted[3]), "r"(FP32_BASE));
 
-        int32_t decode_value[4];
-        ScaleComputeT new_code_zp = code_zp + 0.5f;
+    int32_t decode_value[4];
+    ScaleComputeT new_code_zp = code_zp + 0.5f;
 
-        decode_value[0] = __float2int_rd(fmaf(fp32_intermediates[0], code_scale, new_code_zp));
-        decode_value[1] = __float2int_rd(fmaf(fp32_intermediates[1], code_scale, new_code_zp));
-        decode_value[2] = __float2int_rd(fmaf(fp32_intermediates[2], code_scale, new_code_zp));
-        decode_value[3] = __float2int_rd(fmaf(fp32_intermediates[3], code_scale, new_code_zp));
+    decode_value[0] =
+        __float2int_rd(fmaf(fp32_intermediates[0], code_scale, new_code_zp));
+    decode_value[1] =
+        __float2int_rd(fmaf(fp32_intermediates[1], code_scale, new_code_zp));
+    decode_value[2] =
+        __float2int_rd(fmaf(fp32_intermediates[2], code_scale, new_code_zp));
+    decode_value[3] =
+        __float2int_rd(fmaf(fp32_intermediates[3], code_scale, new_code_zp));
 
-        return convert_impl(decode_value);
-    }
+    return convert_impl(decode_value);
+  }
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source, code_type const& code_scale, code_type const& code_zp)
-    {
-        uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source,
+                             code_type const& code_scale,
+                             code_type const& code_zp) {
+    uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
 
-        // 2^23 = 8388608
-        static constexpr uint32_t FP32_BASE = 0x4B000000;
+    // 2^23 = 8388608
+    static constexpr uint32_t FP32_BASE = 0x4B000000;
 
-        float fp32_intermediates[4];
-        uint32_t* fp32_intermediates_casted = reinterpret_cast<uint32_t*>(fp32_intermediates);
-        fp32_intermediates_casted[0] = __byte_perm(i8s, FP32_BASE, 0x7650);
-        fp32_intermediates_casted[1] = __byte_perm(i8s, FP32_BASE, 0x7651);
-        fp32_intermediates_casted[2] = __byte_perm(i8s, FP32_BASE, 0x7652);
-        fp32_intermediates_casted[3] = __byte_perm(i8s, FP32_BASE, 0x7653);
+    float fp32_intermediates[4];
+    uint32_t* fp32_intermediates_casted =
+        reinterpret_cast<uint32_t*>(fp32_intermediates);
+    fp32_intermediates_casted[0] = __byte_perm(i8s, FP32_BASE, 0x7650);
+    fp32_intermediates_casted[1] = __byte_perm(i8s, FP32_BASE, 0x7651);
+    fp32_intermediates_casted[2] = __byte_perm(i8s, FP32_BASE, 0x7652);
+    fp32_intermediates_casted[3] = __byte_perm(i8s, FP32_BASE, 0x7653);
 
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[0]) : "r"(fp32_intermediates_casted[0]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[1]) : "r"(fp32_intermediates_casted[1]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[2]) : "r"(fp32_intermediates_casted[2]), "r"(FP32_BASE));
-        asm volatile("sub.f32 %0, %1, %2;\n" : "=r"(fp32_intermediates_casted[3]) : "r"(fp32_intermediates_casted[3]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[0])
+                 : "r"(fp32_intermediates_casted[0]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[1])
+                 : "r"(fp32_intermediates_casted[1]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[2])
+                 : "r"(fp32_intermediates_casted[2]), "r"(FP32_BASE));
+    asm volatile("sub.f32 %0, %1, %2;\n"
+                 : "=r"(fp32_intermediates_casted[3])
+                 : "r"(fp32_intermediates_casted[3]), "r"(FP32_BASE));
 
-        int32_t decode_value[4];
+    int32_t decode_value[4];
 
-        decode_value[0] = __float2int_rd(fmaf(fp32_intermediates[0], code_scale[0], code_zp[0] + 0.5f));
-        decode_value[1] = __float2int_rd(fmaf(fp32_intermediates[1], code_scale[1], code_zp[1] + 0.5f));
-        decode_value[2] = __float2int_rd(fmaf(fp32_intermediates[2], code_scale[2], code_zp[2] + 0.5f));
-        decode_value[3] = __float2int_rd(fmaf(fp32_intermediates[3], code_scale[3], code_zp[3] + 0.5f));
+    decode_value[0] = __float2int_rd(
+        fmaf(fp32_intermediates[0], code_scale[0], code_zp[0] + 0.5f));
+    decode_value[1] = __float2int_rd(
+        fmaf(fp32_intermediates[1], code_scale[1], code_zp[1] + 0.5f));
+    decode_value[2] = __float2int_rd(
+        fmaf(fp32_intermediates[2], code_scale[2], code_zp[2] + 0.5f));
+    decode_value[3] = __float2int_rd(
+        fmaf(fp32_intermediates[3], code_scale[3], code_zp[3] + 0.5f));
 
-        return convert_impl(decode_value);
-    }
+    return convert_impl(decode_value);
+  }
 
-    CUTLASS_DEVICE
-    static result_type convert_impl(int32_t* decode_value)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert_impl(int32_t* decode_value) {
+    result_type result;
 
-        static constexpr uint32_t immLut = (0xF0 & 0xCC) | 0xAA;
-        static constexpr uint32_t MASK = 0x003F003F;
-        // 2^7 = 128
-        static constexpr uint32_t EX = 0x43004300;
+    static constexpr uint32_t immLut = (0xF0 & 0xCC) | 0xAA;
+    static constexpr uint32_t MASK = 0x003F003F;
+    // 2^7 = 128
+    static constexpr uint32_t EX = 0x43004300;
 
-        uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
 
-        int32_t q0 = __byte_perm(decode_value[0], decode_value[1], 0x5410);
-        int32_t q1 = __byte_perm(decode_value[2], decode_value[3], 0x5410);
+    int32_t q0 = __byte_perm(decode_value[0], decode_value[1], 0x5410);
+    int32_t q1 = __byte_perm(decode_value[2], decode_value[3], 0x5410);
 
-        h[0] = lop3<immLut>(q0 >> 9, MASK, EX);
-        h[1] = lop3<immLut>(q0 >> 6, MASK, EX);
-        h[2] = lop3<immLut>(q0 >> 3, MASK, EX);
-        h[3] = lop3<immLut>(q0, MASK, EX);
+    h[0] = lop3<immLut>(q0 >> 9, MASK, EX);
+    h[1] = lop3<immLut>(q0 >> 6, MASK, EX);
+    h[2] = lop3<immLut>(q0 >> 3, MASK, EX);
+    h[3] = lop3<immLut>(q0, MASK, EX);
 
-        h[4] = lop3<immLut>(q1 >> 9, MASK, EX);
-        h[5] = lop3<immLut>(q1 >> 6, MASK, EX);
-        h[6] = lop3<immLut>(q1 >> 3, MASK, EX);
-        h[7] = lop3<immLut>(q1, MASK, EX);
+    h[4] = lop3<immLut>(q1 >> 9, MASK, EX);
+    h[5] = lop3<immLut>(q1 >> 6, MASK, EX);
+    h[6] = lop3<immLut>(q1 >> 3, MASK, EX);
+    h[7] = lop3<immLut>(q1, MASK, EX);
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && defined(ENABLE_BF16))
-        // 128 + 32 = 160
-        static constexpr uint32_t SUB = 0x43204320;
+    // 128 + 32 = 160
+    static constexpr uint32_t SUB = 0x43204320;
 
-        asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(SUB));
-        asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(SUB));
-        asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(SUB));
-        asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[3]) : "r"(h[3]), "r"(SUB));
+    asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(SUB));
+    asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[1]) : "r"(h[1]), "r"(SUB));
+    asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(SUB));
+    asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[3]) : "r"(h[3]), "r"(SUB));
 
-        asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[4]) : "r"(h[4]), "r"(SUB));
-        asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[5]) : "r"(h[5]), "r"(SUB));
-        asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[6]) : "r"(h[6]), "r"(SUB));
-        asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[7]) : "r"(h[7]), "r"(SUB));
+    asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[4]) : "r"(h[4]), "r"(SUB));
+    asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[5]) : "r"(h[5]), "r"(SUB));
+    asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[6]) : "r"(h[6]), "r"(SUB));
+    asm volatile("sub.bf16x2 %0, %1, %2;\n" : "=r"(h[7]) : "r"(h[7]), "r"(SUB));
 #else
-        // 1.0
-        static constexpr uint32_t MUL = 0x3F803F80;
-        // -160
-        static constexpr uint32_t ADD = 0xC320C320;
+    // 1.0
+    static constexpr uint32_t MUL = 0x3F803F80;
+    // -160
+    static constexpr uint32_t ADD = 0xC320C320;
 
-        asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[0]) : "r"(h[0]), "r"(MUL), "r"(ADD));
-        asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(MUL), "r"(ADD));
-        asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[2]) : "r"(h[2]), "r"(MUL), "r"(ADD));
-        asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(MUL), "r"(ADD));
+    asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[0])
+                 : "r"(h[0]), "r"(MUL), "r"(ADD));
+    asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[1])
+                 : "r"(h[1]), "r"(MUL), "r"(ADD));
+    asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[2])
+                 : "r"(h[2]), "r"(MUL), "r"(ADD));
+    asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[3])
+                 : "r"(h[3]), "r"(MUL), "r"(ADD));
 
-        asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[4]) : "r"(h[4]), "r"(MUL), "r"(ADD));
-        asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[5]) : "r"(h[5]), "r"(MUL), "r"(ADD));
-        asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[6]) : "r"(h[6]), "r"(MUL), "r"(ADD));
-        asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n" : "=r"(h[7]) : "r"(h[7]), "r"(MUL), "r"(ADD));
+    asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[4])
+                 : "r"(h[4]), "r"(MUL), "r"(ADD));
+    asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[5])
+                 : "r"(h[5]), "r"(MUL), "r"(ADD));
+    asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[6])
+                 : "r"(h[6]), "r"(MUL), "r"(ADD));
+    asm volatile("fma.rn.bf16x2 %0, %1, %2, %3;\n"
+                 : "=r"(h[7])
+                 : "r"(h[7]), "r"(MUL), "r"(ADD));
 #endif
 
-        return result;
-    }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s, ScaleComputeT code_scale, ScaleComputeT code_zp)
-    {
-        return convert(s, code_scale, code_zp);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s,
+                         ScaleComputeT code_scale,
+                         ScaleComputeT code_zp) {
+    return convert(s, code_scale, code_zp);
+  }
 };
 
 template <typename T, int N>
-struct FastInterleavedAndBiasedNumericArrayConverter<T, uint2b_t, N>
-{
-    static_assert(platform::is_same<T, half_t>::value || platform::is_same<T, bfloat16_t>::value,
-        "T must be fp16 or bf16");
+struct FastInterleavedAndBiasedNumericArrayConverter<T, uint2b_t, N> {
+  static_assert(platform::is_same<T, half_t>::value ||
+                    platform::is_same<T, bfloat16_t>::value,
+                "T must be fp16 or bf16");
 
-    static constexpr int kVecWidth = 16;
-    static_assert(!(N % kVecWidth), "N must be multiple of 16.");
+  static constexpr int kVecWidth = 16;
+  static_assert(!(N % kVecWidth), "N must be multiple of 16.");
 
-    using result_type = Array<T, N>;
-    using source_type = Array<uint2b_t, N>;
-    using code_type = Array<float, N / kVecWidth>;
+  using result_type = Array<T, N>;
+  using source_type = Array<uint2b_t, N>;
+  using code_type = Array<float, N / kVecWidth>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source, code_type const& code_scale, code_type const& code_zp)
-    {
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type, scalar_source_type, kVecWidth>
-            convert_vector_;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source,
+                             code_type const& code_scale,
+                             code_type const& code_zp) {
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type,
+                                                  scalar_source_type,
+                                                  kVecWidth>
+        convert_vector_;
 
-        result_type result;
-        using vec_result = Array<scalar_result_type, kVecWidth>;
-        using vec_source = Array<scalar_source_type, kVecWidth>;
+    result_type result;
+    using vec_result = Array<scalar_result_type, kVecWidth>;
+    using vec_source = Array<scalar_source_type, kVecWidth>;
 
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / kVecWidth; ++i)
-        {
-            result_ptr[i] = convert_vector_(source_ptr[i], code_scale[i], code_zp[i]);
-        }
-
-        return result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / kVecWidth; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i], code_scale[i], code_zp[i]);
     }
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source, Array<float, N / 4> const& code_scale, Array<float, N / 4> const& code_zp)
-    {
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        using Converter = FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type, scalar_source_type, kVecWidth>;
+    return result;
+  }
 
-        result_type result;
-        using vec_result = typename Converter::result_type;
-        using vec_source = typename Converter::source_type;
-        using vec_code = typename Converter::code_type;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source,
+                             Array<float, N / 4> const& code_scale,
+                             Array<float, N / 4> const& code_zp) {
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    using Converter =
+        FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type,
+                                                      scalar_source_type,
+                                                      kVecWidth>;
 
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
-        vec_code const* code_scale_ptr = reinterpret_cast<vec_code const*>(&code_scale);
-        vec_code const* code_zp_ptr = reinterpret_cast<vec_code const*>(&code_zp);
+    result_type result;
+    using vec_result = typename Converter::result_type;
+    using vec_source = typename Converter::source_type;
+    using vec_code = typename Converter::code_type;
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / kVecWidth; ++i)
-        {
-            result_ptr[i] = Converter::convert(source_ptr[i], code_scale_ptr[i], code_zp_ptr[i]);
-        }
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    vec_code const* code_scale_ptr =
+        reinterpret_cast<vec_code const*>(&code_scale);
+    vec_code const* code_zp_ptr = reinterpret_cast<vec_code const*>(&code_zp);
 
-        return result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / kVecWidth; ++i) {
+      result_ptr[i] =
+          Converter::convert(source_ptr[i], code_scale_ptr[i], code_zp_ptr[i]);
     }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s, code_type const& code_scale, code_type const& code_zp)
-    {
-        return convert(s, code_scale, code_zp);
-    }
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s,
+                         code_type const& code_scale,
+                         code_type const& code_zp) {
+    return convert(s, code_scale, code_zp);
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_extensions/tile_interleaved_layout.h b/custom_ops/gpu_ops/cutlass_extensions/tile_interleaved_layout.h
index 5a0cd29570..928f2645a5 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/tile_interleaved_layout.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/tile_interleaved_layout.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -38,29 +39,24 @@
 #include "cutlass/matrix_coord.h"
 #include "cutlass/pitch_linear_coord.h"
 
-namespace cutlass
-{
-namespace layout
-{
+namespace cutlass {
+namespace layout {
 
 template <int RowsPerTile, int ColumnsInterleaved>
-struct ColumnMajorTileInterleave
-{
-    static constexpr int kRowsPerTile = RowsPerTile;
-    static constexpr int kColumnsInterleaved = ColumnsInterleaved;
+struct ColumnMajorTileInterleave {
+  static constexpr int kRowsPerTile = RowsPerTile;
+  static constexpr int kColumnsInterleaved = ColumnsInterleaved;
 };
 
 template <class T>
-struct IsColumnMajorTileInterleave
-{
-    static constexpr bool value = false;
+struct IsColumnMajorTileInterleave {
+  static constexpr bool value = false;
 };
 
 template <int U, int V>
-struct IsColumnMajorTileInterleave<ColumnMajorTileInterleave<U, V>>
-{
-    static constexpr bool value = true;
+struct IsColumnMajorTileInterleave<ColumnMajorTileInterleave<U, V>> {
+  static constexpr bool value = true;
 };
 
-} // namespace layout
-} // namespace cutlass
+}  // namespace layout
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/transform/threadblock/fine_grained_scale_zero_iterator.h b/custom_ops/gpu_ops/cutlass_extensions/transform/threadblock/fine_grained_scale_zero_iterator.h
index 6095925e37..6d45e5cb02 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/transform/threadblock/fine_grained_scale_zero_iterator.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/transform/threadblock/fine_grained_scale_zero_iterator.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,19 +18,20 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
-    \brief Templates for visiting scales to be used when dequantizing the weights for weight-only GEMM
-           quantization.
+    \brief Templates for visiting scales to be used when dequantizing the
+   weights for weight-only GEMM quantization.
 */
 
 #pragma once
@@ -50,201 +51,205 @@
 
 ////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace transform
-{
-namespace threadblock
-{
+namespace cutlass {
+namespace transform {
+namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template <typename Shape, typename Element, typename Layout, int AdvanceRank, int Alignment>
+template <typename Shape,
+          typename Element,
+          typename Layout,
+          int AdvanceRank,
+          int Alignment>
 class FineGrainedScaleZeroIterator;
 
 template <typename Shape_, typename Element_, int Alignment_>
-class FineGrainedScaleZeroIterator<Shape_, Element_, layout::RowMajor, 0, Alignment_>
-{
-public:
-    using Shape = Shape_;
-    using Element = Element_;
-    using Layout = layout::RowMajor;
-    static int const kAdvanceRank = 0;
-    static int const kAlignment = Alignment_;
+class FineGrainedScaleZeroIterator<Shape_,
+                                   Element_,
+                                   layout::RowMajor,
+                                   0,
+                                   Alignment_> {
+ public:
+  using Shape = Shape_;
+  using Element = Element_;
+  using Layout = layout::RowMajor;
+  static int const kAdvanceRank = 0;
+  static int const kAlignment = Alignment_;
 
-    static int const kAccessesPerVector = 1;
+  static int const kAccessesPerVector = 1;
 
-    /// Row index of scales corresponding to the groupsize of 64
-    int row_groupsize64_;
-    int group_size_;
+  /// Row index of scales corresponding to the groupsize of 64
+  int row_groupsize64_;
+  int group_size_;
 
-    using Index = typename Layout::Index;
-    using LongIndex = typename Layout::LongIndex;
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
 
-    using TensorRef = TensorRef<Element, Layout>;
-    using TensorView = TensorView<Element, Layout>;
-    using TensorCoord = typename Layout::TensorCoord;
-    using Pointer = Element*;
-    using NonConstPointer = typename platform::remove_const<Element>::type*;
+  using TensorRef = TensorRef<Element, Layout>;
+  using TensorView = TensorView<Element, Layout>;
+  using TensorCoord = typename Layout::TensorCoord;
+  using Pointer = Element*;
+  using NonConstPointer = typename platform::remove_const<Element>::type*;
 
-    using AccessType = AlignedArray<Element, kAlignment>;
+  using AccessType = AlignedArray<Element, kAlignment>;
 
-    using Fragment = cutlass::Array<Element, kAlignment>;
+  using Fragment = cutlass::Array<Element, kAlignment>;
 
-    // For compatibility with existing iterator interface
-    struct Params
-    {
-        LongIndex stride_ = 0;
+  // For compatibility with existing iterator interface
+  struct Params {
+    LongIndex stride_ = 0;
 
-        /// amount (in byte) to increment pointer from first access of current tile
-        /// to first access of next tile
-        LongIndex inc_advance_ = 0;
+    /// amount (in byte) to increment pointer from first access of current tile
+    /// to first access of next tile
+    LongIndex inc_advance_ = 0;
 
-        // Default ctor
-        CUTLASS_HOST_DEVICE
-        Params() {}
-
-        /// Construct the Params object given a pitch-linear tensor's layout
-        CUTLASS_HOST_DEVICE
-        Params(Layout const& layout)
-            : stride_(layout.stride(0))
-        {
-            inc_advance_ = Shape::kRow * stride_ * sizeof_bits<Element>::value / 8;
-        }
-    };
-
-private:
-    /// Internal pointer type permits fast address arithmetic
-    using BytePointer = char*;
-
-private:
-    //
-    // Data members
-    //
-
-    /// Parameters object with precomputed internal state
-    Params const params_;
-
-    /// Internal pointer to first access of tile
-    BytePointer pointer_scale_;
-    BytePointer pointer_zero_;
-
-    bool is_valid_ = false;
-
-public:
-    /// Constructs a TileIterator from its precomputed state, threadblock offset,
-    /// and thread ID
-    CUTLASS_DEVICE
-    FineGrainedScaleZeroIterator(
-        ///< Precomputed parameters object
-        Params const& params,
-        ///< Pointer to start of scale tensor
-        Pointer pointer_scale,
-        ///< Pointer to start of zero tensor
-        Pointer pointer_zero,
-        ///< Extent of the scale and bias
-        TensorCoord extent,
-        ///< ID of each participating thread
-        int thread_id,
-        ///< Initial offset of threadblock
-        TensorCoord const& threadblock_offset,
-        ///< Group size
-        int group_size)
-        : params_(params)
-        , pointer_scale_(reinterpret_cast<BytePointer>(const_cast<NonConstPointer>(pointer_scale)))
-        , pointer_zero_(reinterpret_cast<BytePointer>(const_cast<NonConstPointer>(pointer_zero)))
-    {
-        row_groupsize64_ = threadblock_offset.row();
-        group_size_ = group_size;
-
-        const LongIndex tb_row_byte_offset
-            = threadblock_offset.row() / (group_size / 64) * params_.stride_ * sizeof_bits<Element>::value / 8;
-        const LongIndex tb_col_byte_offset = threadblock_offset.column() * sizeof_bits<Element>::value / 8;
-        pointer_scale_ += (tb_row_byte_offset + tb_col_byte_offset);
-
-        if (pointer_zero_ != nullptr)
-        {
-            pointer_zero_ += (tb_row_byte_offset + tb_col_byte_offset);
-        }
-
-        static constexpr int THREADS_PER_ROW = Shape::kColumn / kAlignment;
-
-        int const thread_row = thread_id / THREADS_PER_ROW;
-        int const thread_col = thread_id % THREADS_PER_ROW;
-
-        const LongIndex thread_row_byte_offset = thread_row * params_.stride_ * sizeof_bits<Element>::value / 8;
-        const LongIndex thread_col_byte_offset = thread_col * kAlignment * sizeof_bits<Element>::value / 8;
-        pointer_scale_ += (thread_row_byte_offset + thread_col_byte_offset);
-        if (pointer_zero_ != nullptr)
-        {
-            pointer_zero_ += (thread_row_byte_offset + thread_col_byte_offset);
-        }
-
-        // For the rows, we must check that we are within the extent AND the tile to avoid extra reads on
-        // a given iteration. The same threads will be responsible for issues reads since the number of scales
-        // read in a given iteration is a constant. Therefore, we should never have to update is_valid_
-        // outside of the constructor.
-        int const global_row = threadblock_offset.row() + thread_row;
-        int const global_col = threadblock_offset.column() + thread_col * kAlignment;
-
-        bool const row_in_bounds = global_row < extent.row() && thread_row < Shape::kRow;
-        bool const col_in_bounds = global_col < extent.column();
-
-        is_valid_ = row_in_bounds && col_in_bounds;
-    }
-
-    /// Construct a PredicatedTileAccessIterator with zero threadblock offset
-    CUTLASS_HOST_DEVICE FineGrainedScaleZeroIterator(Params const& params, ///< Precomputed parameters object
-        Pointer pointer_scale,                                             ///< Pointer to start of scale tensor
-        Pointer pointer_zero,                                              ///< Pointer to start of zero tensor
-        TensorCoord extent,                                                ///< Extent of tensor
-        int thread_id,                                                     ///< ID of each participating thread
-        int group_size)
-        : FineGrainedScaleZeroIterator(
-            params, pointer_scale, pointer_zero, extent, thread_id, make_Coord(0, 0), group_size)
-    {
-    }
-
-    CUTLASS_DEVICE
-    void add_tile_offset(TensorCoord const& tile_offset)
-    {
-        const LongIndex row_byte_offset = tile_offset.row() * params_.inc_advance_;
-        const LongIndex col_byte_offset = tile_offset.column() * Shape::kColumn * sizeof_bits<Element>::value / 8;
-        pointer_scale_ += row_byte_offset + col_byte_offset;
-        if (pointer_zero_ != nullptr)
-        {
-            pointer_zero_ += row_byte_offset + col_byte_offset;
-        }
-    }
-
-    /// Clears the predicate set efficiently
-    CUTLASS_HOST_DEVICE void clear_mask(bool enable = true)
-    {
-        is_valid_ &= (!enable);
-    }
-
-    /// Returns whether access is valid or not
+    // Default ctor
     CUTLASS_HOST_DEVICE
-    bool valid() const
-    {
-        return is_valid_;
+    Params() {}
+
+    /// Construct the Params object given a pitch-linear tensor's layout
+    CUTLASS_HOST_DEVICE
+    Params(Layout const& layout) : stride_(layout.stride(0)) {
+      inc_advance_ = Shape::kRow * stride_ * sizeof_bits<Element>::value / 8;
+    }
+  };
+
+ private:
+  /// Internal pointer type permits fast address arithmetic
+  using BytePointer = char*;
+
+ private:
+  //
+  // Data members
+  //
+
+  /// Parameters object with precomputed internal state
+  Params const params_;
+
+  /// Internal pointer to first access of tile
+  BytePointer pointer_scale_;
+  BytePointer pointer_zero_;
+
+  bool is_valid_ = false;
+
+ public:
+  /// Constructs a TileIterator from its precomputed state, threadblock offset,
+  /// and thread ID
+  CUTLASS_DEVICE
+  FineGrainedScaleZeroIterator(
+      ///< Precomputed parameters object
+      Params const& params,
+      ///< Pointer to start of scale tensor
+      Pointer pointer_scale,
+      ///< Pointer to start of zero tensor
+      Pointer pointer_zero,
+      ///< Extent of the scale and bias
+      TensorCoord extent,
+      ///< ID of each participating thread
+      int thread_id,
+      ///< Initial offset of threadblock
+      TensorCoord const& threadblock_offset,
+      ///< Group size
+      int group_size)
+      : params_(params),
+        pointer_scale_(reinterpret_cast<BytePointer>(
+            const_cast<NonConstPointer>(pointer_scale))),
+        pointer_zero_(reinterpret_cast<BytePointer>(
+            const_cast<NonConstPointer>(pointer_zero))) {
+    row_groupsize64_ = threadblock_offset.row();
+    group_size_ = group_size;
+
+    const LongIndex tb_row_byte_offset = threadblock_offset.row() /
+                                         (group_size / 64) * params_.stride_ *
+                                         sizeof_bits<Element>::value / 8;
+    const LongIndex tb_col_byte_offset =
+        threadblock_offset.column() * sizeof_bits<Element>::value / 8;
+    pointer_scale_ += (tb_row_byte_offset + tb_col_byte_offset);
+
+    if (pointer_zero_ != nullptr) {
+      pointer_zero_ += (tb_row_byte_offset + tb_col_byte_offset);
     }
 
-    /// Returns a scale pointer
-    CUTLASS_HOST_DEVICE
-    AccessType* get_scale() const
-    {
-        return reinterpret_cast<AccessType*>(pointer_scale_);
+    static constexpr int THREADS_PER_ROW = Shape::kColumn / kAlignment;
+
+    int const thread_row = thread_id / THREADS_PER_ROW;
+    int const thread_col = thread_id % THREADS_PER_ROW;
+
+    const LongIndex thread_row_byte_offset =
+        thread_row * params_.stride_ * sizeof_bits<Element>::value / 8;
+    const LongIndex thread_col_byte_offset =
+        thread_col * kAlignment * sizeof_bits<Element>::value / 8;
+    pointer_scale_ += (thread_row_byte_offset + thread_col_byte_offset);
+    if (pointer_zero_ != nullptr) {
+      pointer_zero_ += (thread_row_byte_offset + thread_col_byte_offset);
     }
 
-    /// Returns a zero pointer
-    CUTLASS_HOST_DEVICE
-    AccessType* get_zero() const
-    {
-        return reinterpret_cast<AccessType*>(pointer_zero_);
+    // For the rows, we must check that we are within the extent AND the tile to
+    // avoid extra reads on a given iteration. The same threads will be
+    // responsible for issues reads since the number of scales read in a given
+    // iteration is a constant. Therefore, we should never have to update
+    // is_valid_ outside of the constructor.
+    int const global_row = threadblock_offset.row() + thread_row;
+    int const global_col =
+        threadblock_offset.column() + thread_col * kAlignment;
+
+    bool const row_in_bounds =
+        global_row < extent.row() && thread_row < Shape::kRow;
+    bool const col_in_bounds = global_col < extent.column();
+
+    is_valid_ = row_in_bounds && col_in_bounds;
+  }
+
+  /// Construct a PredicatedTileAccessIterator with zero threadblock offset
+  CUTLASS_HOST_DEVICE FineGrainedScaleZeroIterator(
+      Params const& params,   ///< Precomputed parameters object
+      Pointer pointer_scale,  ///< Pointer to start of scale tensor
+      Pointer pointer_zero,   ///< Pointer to start of zero tensor
+      TensorCoord extent,     ///< Extent of tensor
+      int thread_id,          ///< ID of each participating thread
+      int group_size)
+      : FineGrainedScaleZeroIterator(params,
+                                     pointer_scale,
+                                     pointer_zero,
+                                     extent,
+                                     thread_id,
+                                     make_Coord(0, 0),
+                                     group_size) {}
+
+  CUTLASS_DEVICE
+  void add_tile_offset(TensorCoord const& tile_offset) {
+    const LongIndex row_byte_offset = tile_offset.row() * params_.inc_advance_;
+    const LongIndex col_byte_offset =
+        tile_offset.column() * Shape::kColumn * sizeof_bits<Element>::value / 8;
+    pointer_scale_ += row_byte_offset + col_byte_offset;
+    if (pointer_zero_ != nullptr) {
+      pointer_zero_ += row_byte_offset + col_byte_offset;
     }
+  }
+
+  /// Clears the predicate set efficiently
+  CUTLASS_HOST_DEVICE void clear_mask(bool enable = true) {
+    is_valid_ &= (!enable);
+  }
+
+  /// Returns whether access is valid or not
+  CUTLASS_HOST_DEVICE
+  bool valid() const { return is_valid_; }
+
+  /// Returns a scale pointer
+  CUTLASS_HOST_DEVICE
+  AccessType* get_scale() const {
+    return reinterpret_cast<AccessType*>(pointer_scale_);
+  }
+
+  /// Returns a zero pointer
+  CUTLASS_HOST_DEVICE
+  AccessType* get_zero() const {
+    return reinterpret_cast<AccessType*>(pointer_zero_);
+  }
 };
 
-} // namespace threadblock
-} // namespace transform
-} // namespace cutlass
+}  // namespace threadblock
+}  // namespace transform
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/util/gather_tensor.hpp b/custom_ops/gpu_ops/cutlass_extensions/util/gather_tensor.hpp
index b430380b01..b29fc4db5f 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/util/gather_tensor.hpp
+++ b/custom_ops/gpu_ops/cutlass_extensions/util/gather_tensor.hpp
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 #pragma once
@@ -38,144 +39,148 @@ using namespace cute;
 
 /// Function object that applies an index to its argument
 template <class Iter>
-struct IndexedGather
-{
-    CUTE_HOST_DEVICE constexpr IndexedGather(Iter indices = {})
-        : indices_(indices)
-    {
-    }
+struct IndexedGather {
+  CUTE_HOST_DEVICE constexpr IndexedGather(Iter indices = {})
+      : indices_(indices) {}
 
-    template <typename I>
-    CUTE_HOST_DEVICE constexpr auto operator()(I i) const
-    {
-        return indices_[i];
-    }
+  template <typename I>
+  CUTE_HOST_DEVICE constexpr auto operator()(I i) const {
+    return indices_[i];
+  }
 
-    CUTE_HOST_DEVICE friend void print(IndexedGather const& s)
-    {
-        cute::print("Indexed{");
-        print(s.indices_);
-        print("}");
-    }
+  CUTE_HOST_DEVICE friend void print(IndexedGather const& s) {
+    cute::print("Indexed{");
+    print(s.indices_);
+    print("}");
+  }
 
-    Iter indices_;
+  Iter indices_;
 };
 
 /// Custom stride object that applies a function followed by a stride
 template <class Func, class Stride>
-struct CustomStride
-{
-    CUTE_HOST_DEVICE constexpr CustomStride(Func const& func, Stride const& stride)
-        : func_(func)
-        , stride_(stride)
-    {
-    }
+struct CustomStride {
+  CUTE_HOST_DEVICE constexpr CustomStride(Func const& func,
+                                          Stride const& stride)
+      : func_(func), stride_(stride) {}
 
-    template <class I>
-    CUTE_HOST_DEVICE constexpr friend auto operator*(I i, CustomStride const& s)
-    {
-        return s.func_(i) * s.stride_;
-    }
+  template <class I>
+  CUTE_HOST_DEVICE constexpr friend auto operator*(I i, CustomStride const& s) {
+    return s.func_(i) * s.stride_;
+  }
 
-    template <class I>
-    CUTE_HOST_DEVICE constexpr friend auto operator*(CustomStride const& s, I i)
-    {
-        return s.func_(i) * s.stride_;
-    }
+  template <class I>
+  CUTE_HOST_DEVICE constexpr friend auto operator*(CustomStride const& s, I i) {
+    return s.func_(i) * s.stride_;
+  }
 
-    CUTE_HOST_DEVICE friend void print(CustomStride const& s)
-    {
-        cute::print("Custom{");
-        print(s.func_);
-        cute::print(",");
-        print(s.stride_);
-        cute::print("}");
-    }
+  CUTE_HOST_DEVICE friend void print(CustomStride const& s) {
+    cute::print("Custom{");
+    print(s.func_);
+    cute::print(",");
+    print(s.stride_);
+    cute::print("}");
+  }
 
-    template <class Div>
-    CUTE_HOST_DEVICE constexpr friend auto safe_div(CustomStride const& s, Div const& div)
-    {
-        return CustomStride<Func, decltype(safe_div(s.stride_, div))>(s.func_, safe_div(s.stride_, div));
-    }
+  template <class Div>
+  CUTE_HOST_DEVICE constexpr friend auto safe_div(CustomStride const& s,
+                                                  Div const& div) {
+    return CustomStride<Func, decltype(safe_div(s.stride_, div))>(
+        s.func_, safe_div(s.stride_, div));
+  }
 
-    // Circumvent the requirement on make_layout that shape and stride are integral
-    template <class Shape>
-    CUTE_HOST_DEVICE constexpr friend auto make_layout(Shape const& shape, CustomStride const& stride)
-    {
-        return Layout<Shape, CustomStride>(shape, stride);
-    }
+  // Circumvent the requirement on make_layout that shape and stride are
+  // integral
+  template <class Shape>
+  CUTE_HOST_DEVICE constexpr friend auto make_layout(
+      Shape const& shape, CustomStride const& stride) {
+    return Layout<Shape, CustomStride>(shape, stride);
+  }
 
-    Func func_;
-    Stride stride_;
+  Func func_;
+  Stride stride_;
 };
 
 template <class Stride, class Func>
-CUTLASS_HOST_DEVICE auto make_custom_stride_layout(Stride const& stride, Func&& func)
-{
-    // Use a dummy shape and replace the first non-unit and non-zero stride with a custom gather stride
-    auto idx = find_if(stride, [](auto x) { return !is_constant<1, decltype(x)>{} && !is_constant<0, decltype(x)>{}; });
-    constexpr int I = decltype(idx)::value;
-    return make_layout(
-        repeat_like(stride, _1{}), replace<I>(stride, CustomStride{static_cast<Func&&>(func), get<I>(stride)}));
+CUTLASS_HOST_DEVICE auto make_custom_stride_layout(Stride const& stride,
+                                                   Func&& func) {
+  // Use a dummy shape and replace the first non-unit and non-zero stride with a
+  // custom gather stride
+  auto idx = find_if(stride, [](auto x) {
+    return !is_constant<1, decltype(x)>{} && !is_constant<0, decltype(x)>{};
+  });
+  constexpr int I = decltype(idx)::value;
+  return make_layout(
+      repeat_like(stride, _1{}),
+      replace<I>(stride,
+                 CustomStride{static_cast<Func&&>(func), get<I>(stride)}));
 }
 
 /// Helper function to optionally create a gather tensor
 template <class Iterator, class Shape, class Stride, class Func>
-CUTLASS_HOST_DEVICE auto make_gather_tensor(Iterator iter, Shape const& shape, Stride const& stride, Func&& func)
-{
-    Layout matrix_layout = make_identity_layout(shape);
-    auto offset = as_arithmetic_tuple(repeat_like(shape, _0{}));
-    Layout gather_layout = make_custom_stride_layout(stride, static_cast<Func&&>(func));
-    return make_tensor(iter, ComposedLayout{gather_layout, offset, matrix_layout});
+CUTLASS_HOST_DEVICE auto make_gather_tensor(Iterator iter,
+                                            Shape const& shape,
+                                            Stride const& stride,
+                                            Func&& func) {
+  Layout matrix_layout = make_identity_layout(shape);
+  auto offset = as_arithmetic_tuple(repeat_like(shape, _0{}));
+  Layout gather_layout =
+      make_custom_stride_layout(stride, static_cast<Func&&>(func));
+  return make_tensor(iter,
+                     ComposedLayout{gather_layout, offset, matrix_layout});
 }
 
-namespace cute
-{
+namespace cute {
 
 template <int N, int I, class Shape, class Stride>
-CUTE_HOST_DEVICE constexpr auto upcast(Shape const& shape, Stride const& stride)
-{
-    if constexpr (is_tuple<Shape>::value)
-    {
-        return transform_layout(shape, stride, [](auto const& s, auto const& d) { return upcast<N, I>(s, d); });
-    }
-    else if constexpr (is_scaled_basis<Stride>::value)
-    {
-        if constexpr (Stride::mode() == I)
-        {
-            return make_layout(shape_div(shape, Int<N>{}), shape_div(stride, Int<N>{}));
-        }
-        else
-        {
-            return make_layout(shape, stride);
-        }
-    }
-    else
-    {
-        return upcast<N>(shape, stride);
+CUTE_HOST_DEVICE constexpr auto upcast(Shape const& shape,
+                                       Stride const& stride) {
+  if constexpr (is_tuple<Shape>::value) {
+    return transform_layout(shape, stride, [](auto const& s, auto const& d) {
+      return upcast<N, I>(s, d);
+    });
+  } else if constexpr (is_scaled_basis<Stride>::value) {
+    if constexpr (Stride::mode() == I) {
+      return make_layout(shape_div(shape, Int<N>{}),
+                         shape_div(stride, Int<N>{}));
+    } else {
+      return make_layout(shape, stride);
     }
+  } else {
+    return upcast<N>(shape, stride);
+  }
 
-    CUTE_GCC_UNREACHABLE;
+  CUTE_GCC_UNREACHABLE;
 }
 
-template <int N, class OuterShape, class OuterStride, class Offset, class Shape, class Stride>
+template <int N,
+          class OuterShape,
+          class OuterStride,
+          class Offset,
+          class Shape,
+          class Stride>
 CUTE_HOST_DEVICE constexpr auto upcast(
-    ComposedLayout<Layout<OuterShape, OuterStride>, Offset, Layout<Shape, Stride>> const& layout)
-{
-    // Find index of the stride-1 mode - that is the only one that requires updating inner shape and offset
-    auto idx = find_if(layout.layout_a().stride(), [](auto x) { return is_constant<1, decltype(x)>{}; });
-    constexpr int I = decltype(idx)::value;
+    ComposedLayout<Layout<OuterShape, OuterStride>,
+                   Offset,
+                   Layout<Shape, Stride>> const& layout) {
+  // Find index of the stride-1 mode - that is the only one that requires
+  // updating inner shape and offset
+  auto idx = find_if(layout.layout_a().stride(),
+                     [](auto x) { return is_constant<1, decltype(x)>{}; });
+  constexpr int I = decltype(idx)::value;
 
-    // Upcast the outer layout (works as expected)
-    auto outer = upcast<N>(layout.layout_a());
+  // Upcast the outer layout (works as expected)
+  auto outer = upcast<N>(layout.layout_a());
 
-    // Upcast the accumulated offset along stride-1 mode
-    auto offset = as_arithmetic_tuple(replace<I>(layout.offset(), upcast<N>(get<I>(layout.offset()))));
+  // Upcast the accumulated offset along stride-1 mode
+  auto offset = as_arithmetic_tuple(
+      replace<I>(layout.offset(), upcast<N>(get<I>(layout.offset()))));
 
-    // Upcast the inner layout's shape along stride-1 mode
-    auto inner = upcast<N, I>(layout.layout_b().shape(), layout.layout_b().stride());
+  // Upcast the inner layout's shape along stride-1 mode
+  auto inner =
+      upcast<N, I>(layout.layout_b().shape(), layout.layout_b().stride());
 
-    return composition(outer, offset, inner);
+  return composition(outer, offset, inner);
 }
 
-} // namespace cute
+}  // namespace cute
diff --git a/custom_ops/gpu_ops/cutlass_extensions/weight_only_quant_op.h b/custom_ops/gpu_ops/cutlass_extensions/weight_only_quant_op.h
index 64774428e9..9f2d255234 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/weight_only_quant_op.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/weight_only_quant_op.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,41 +18,40 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
-  \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores.
+  \brief Defines iterators used by warp-level matrix multiply operations
+  targeting Tensor Cores.
 */
 
 #pragma once
 
-namespace cutlass
-{
+namespace cutlass {
 
-enum class WeightOnlyQuantOp
-{
-    UNDEFINED,
-    PER_COLUMN_SCALE_ONLY,
-    FINEGRAINED_SCALE_ONLY,
-    FINEGRAINED_SCALE_AND_ZEROS
+enum class WeightOnlyQuantOp {
+  UNDEFINED,
+  PER_COLUMN_SCALE_ONLY,
+  FINEGRAINED_SCALE_ONLY,
+  FINEGRAINED_SCALE_AND_ZEROS
 };
 
-constexpr bool isFinegrained(WeightOnlyQuantOp op)
-{
-    return op == WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS || op == WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY;
+constexpr bool isFinegrained(WeightOnlyQuantOp op) {
+  return op == WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS ||
+         op == WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY;
 }
 
-constexpr bool hasZero(WeightOnlyQuantOp op)
-{
-    return op == WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS;
+constexpr bool hasZero(WeightOnlyQuantOp op) {
+  return op == WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS;
 }
 
-} // namespace cutlass
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_extensions/wint_type_traits.h b/custom_ops/gpu_ops/cutlass_extensions/wint_type_traits.h
index fa28810697..da6fcf41a2 100644
--- a/custom_ops/gpu_ops/cutlass_extensions/wint_type_traits.h
+++ b/custom_ops/gpu_ops/cutlass_extensions/wint_type_traits.h
@@ -33,19 +33,23 @@ enum WintQuantMethod {
 };
 
 // Convert CUDA data type to cutlass data type
-template <typename T> struct CutlassDataType {
+template <typename T>
+struct CutlassDataType {
   using Type = T;
 };
 
-template <> struct CutlassDataType<half> {
+template <>
+struct CutlassDataType<half> {
   using Type = cutlass::half_t;
 };
 
-template <> struct CutlassDataType<__nv_bfloat16> {
+template <>
+struct CutlassDataType<__nv_bfloat16> {
   using Type = cutlass::bfloat16_t;
 };
 
-template <typename ElementT, WintQuantMethod Method> struct WintQuantTraits;
+template <typename ElementT, WintQuantMethod Method>
+struct WintQuantTraits;
 
 template <typename ElementT>
 struct WintQuantTraits<ElementT, WintQuantMethod::kNone> {
@@ -129,7 +133,7 @@ struct WintQuantTraits<ElementT, WintQuantMethod::kWeightOnlyInt2> {
   using CodeScaleZpType = float;
 
   struct Arguments {
-    uint8_t *local_scale_ptr; // quanted 4-bits
+    uint8_t *local_scale_ptr;  // quanted 4-bits
     float *code_scale_ptr;
     float *code_zp_ptr;
   };
@@ -140,4 +144,4 @@ struct WintQuantTraits<ElementT, WintQuantMethod::kWeightOnlyInt2> {
   }
 };
 
-} // namespace cutlass
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_kernels/cutlass_helper.h b/custom_ops/gpu_ops/cutlass_kernels/cutlass_helper.h
index 3ac548c625..f9e4aad2ff 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/cutlass_helper.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/cutlass_helper.h
@@ -24,11 +24,11 @@
 /**
  * Helper function for checking CUTLASS errors
  */
-#define CUTLASS_CHECK(status)                                                  \
-  {                                                                            \
-    cutlass::Status error = status;                                            \
-    PD_CHECK(error == cutlass::Status::kSuccess,                               \
-             cutlassGetStatusString(error));                                   \
+#define CUTLASS_CHECK(status)                    \
+  {                                              \
+    cutlass::Status error = status;              \
+    PD_CHECK(error == cutlass::Status::kSuccess, \
+             cutlassGetStatusString(error));     \
   }
 
 /**
@@ -38,44 +38,50 @@
  * __CUDA_ARCH__ is not defined in host code, so this lets us smuggle the ifdef
  * into code that will be executed on the device where it is defined.
  */
-template <typename Kernel> struct enable_sm90_or_later : Kernel {
-  template <typename... Args> CUTLASS_DEVICE void operator()(Args &&...args) {
+template <typename Kernel>
+struct enable_sm90_or_later : Kernel {
+  template <typename... Args>
+  CUTLASS_DEVICE void operator()(Args &&...args) {
 #if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 900
     Kernel::operator()(std::forward<Args>(args)...);
 #endif
   }
 };
 
-template <paddle::DataType D> class CutlassDtypeTraits;
+template <paddle::DataType D>
+class CutlassDtypeTraits;
 
-template <> class CutlassDtypeTraits<paddle::DataType::FLOAT32> {
-public:
+template <>
+class CutlassDtypeTraits<paddle::DataType::FLOAT32> {
+ public:
   typedef float DataType;
   typedef float data_t;
 };
 
-template <> class CutlassDtypeTraits<paddle::DataType::FLOAT16> {
-public:
+template <>
+class CutlassDtypeTraits<paddle::DataType::FLOAT16> {
+ public:
   typedef cutlass::half_t DataType;
   typedef paddle::float16 data_t;
 };
 
-template <> class CutlassDtypeTraits<paddle::DataType::BFLOAT16> {
-public:
+template <>
+class CutlassDtypeTraits<paddle::DataType::BFLOAT16> {
+ public:
   typedef cutlass::bfloat16_t DataType;
   typedef paddle::bfloat16 data_t;
 };
 
 class CutlassGemmConfigMannager {
-public:
+ public:
   static CutlassGemmConfigMannager &getInstance() {
     static CutlassGemmConfigMannager instance;
     return instance;
   }
 
   CutlassGemmConfigMannager(const CutlassGemmConfigMannager &) = delete;
-  CutlassGemmConfigMannager &
-  operator=(const CutlassGemmConfigMannager &) = delete;
+  CutlassGemmConfigMannager &operator=(const CutlassGemmConfigMannager &) =
+      delete;
 
   void up_date_configs(const nlohmann::json &j) {
     std::lock_guard<std::mutex> lock(mutex_);
@@ -102,7 +108,7 @@ public:
     return &json_;
   }
 
-private:
+ private:
   void save_gemm_best_configs_(const std::string &config_file_path) {
     std::ifstream file(config_file_path);
     if (!file.good()) {
diff --git a/custom_ops/gpu_ops/cutlass_kernels/cutlass_heuristic.cu b/custom_ops/gpu_ops/cutlass_kernels/cutlass_heuristic.cu
index 6db16981c6..6ea5a275ad 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/cutlass_heuristic.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/cutlass_heuristic.cu
@@ -19,14 +19,14 @@
 #ifndef _WIN32
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wstrict-aliasing"
-#endif // #ifndef _WIN32
+#endif  // #ifndef _WIN32
 
 #include "cutlass/gemm/gemm.h"
 #include "cutlass/numeric_types.h"
 
 #ifndef _WIN32
 #pragma GCC diagnostic pop
-#endif // #ifndef _WIN32
+#endif  // #ifndef _WIN32
 
 #include <cuda_runtime_api.h>
 #include <set>
@@ -35,491 +35,509 @@
 
 using namespace cutlass_extensions;
 
-namespace kernels
-{
-namespace cutlass_kernels
-{
+namespace kernels {
+namespace cutlass_kernels {
 
-struct TileShape
-{
-    int m;
-    int n;
+struct TileShape {
+  int m;
+  int n;
 };
 
-TileShape get_cta_shape_for_config(CutlassTileConfig tile_config)
-{
-    switch (tile_config)
-    {
-    case CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64: return TileShape{16, 128};
-    case CutlassTileConfig::CtaShape16x256x64_WarpShape16x64x64: return TileShape{16, 256};
-    case CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64: return TileShape{32, 128};
-    case CutlassTileConfig::CtaShape64x64x128_WarpShape32x64x64: return TileShape{64, 64};
+TileShape get_cta_shape_for_config(CutlassTileConfig tile_config) {
+  switch (tile_config) {
+    case CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64:
+      return TileShape{16, 128};
+    case CutlassTileConfig::CtaShape16x256x64_WarpShape16x64x64:
+      return TileShape{16, 256};
+    case CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64:
+      return TileShape{32, 128};
+    case CutlassTileConfig::CtaShape64x64x128_WarpShape32x64x64:
+      return TileShape{64, 64};
     case CutlassTileConfig::CtaShape64x128x64_WarpShape32x64x64:
-    case CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64: return TileShape{64, 128};
-    case CutlassTileConfig::CtaShape128x64x64_WarpShape64x32x64: return TileShape{128, 64};
+    case CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64:
+      return TileShape{64, 128};
+    case CutlassTileConfig::CtaShape128x64x64_WarpShape64x32x64:
+      return TileShape{128, 64};
     case CutlassTileConfig::CtaShape128x128x8_WarpShape64x64x8:
     case CutlassTileConfig::CtaShape128x128x64_WarpShape64x32x64:
     case CutlassTileConfig::CtaShape128x128x64_WarpShape64x64x64:
-    case CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64: return TileShape{128, 128};
-    case CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64: return TileShape{128, 256};
-    case CutlassTileConfig::CtaShape256x128x64_WarpShape64x64x64: return TileShape{256, 128};
-    case CutlassTileConfig::CtaShape16x256x128_WarpShape16x64x128: return TileShape{16, 256};
-    default: throw("[get_grid_shape_for_config] Invalid config");
-    }
+    case CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64:
+      return TileShape{128, 128};
+    case CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64:
+      return TileShape{128, 256};
+    case CutlassTileConfig::CtaShape256x128x64_WarpShape64x64x64:
+      return TileShape{256, 128};
+    case CutlassTileConfig::CtaShape16x256x128_WarpShape16x64x128:
+      return TileShape{16, 256};
+    default:
+      throw("[get_grid_shape_for_config] Invalid config");
+  }
 }
 
-bool is_valid_split_k_factor(int64_t const m, int64_t const n, int64_t const k, TileShape const tile_shape,
-    int const split_k_factor, size_t const workspace_bytes, bool const is_weight_only)
-{
+bool is_valid_split_k_factor(int64_t const m,
+                             int64_t const n,
+                             int64_t const k,
+                             TileShape const tile_shape,
+                             int const split_k_factor,
+                             size_t const workspace_bytes,
+                             bool const is_weight_only) {
+  // All tile sizes have a k_tile of 64.
+  static constexpr int k_tile = 128;
 
-    // All tile sizes have a k_tile of 64.
-    static constexpr int k_tile = 128;
-
-    // For weight-only quant, we need k and k_elements_per_split to be a multiple of cta_k
-    if (is_weight_only)
-    {
-        if ((k % k_tile) != 0)
-        {
-            return false;
-        }
-
-        if ((k % split_k_factor) != 0)
-        {
-            return false;
-        }
-
-        int const k_elements_per_split = k / split_k_factor;
-        if ((k_elements_per_split % k_tile) != 0)
-        {
-            return false;
-        }
+  // For weight-only quant, we need k and k_elements_per_split to be a multiple
+  // of cta_k
+  if (is_weight_only) {
+    if ((k % k_tile) != 0) {
+      return false;
     }
 
-    // Check that the workspace has sufficient space for this split-k factor
-    int const ctas_in_m_dim = (m + tile_shape.m - 1) / tile_shape.m;
-    int const ctas_in_n_dim = (n + tile_shape.n - 1) / tile_shape.n;
-    int const required_ws_bytes = split_k_factor == 1 ? 0 : sizeof(int) * ctas_in_m_dim * ctas_in_n_dim;
-
-    if (required_ws_bytes > workspace_bytes)
-    {
-        return false;
+    if ((k % split_k_factor) != 0) {
+      return false;
     }
 
-    return true;
+    int const k_elements_per_split = k / split_k_factor;
+    if ((k_elements_per_split % k_tile) != 0) {
+      return false;
+    }
+  }
+
+  // Check that the workspace has sufficient space for this split-k factor
+  int const ctas_in_m_dim = (m + tile_shape.m - 1) / tile_shape.m;
+  int const ctas_in_n_dim = (n + tile_shape.n - 1) / tile_shape.n;
+  int const required_ws_bytes =
+      split_k_factor == 1 ? 0 : sizeof(int) * ctas_in_m_dim * ctas_in_n_dim;
+
+  if (required_ws_bytes > workspace_bytes) {
+    return false;
+  }
+
+  return true;
 }
 
 std::vector<CutlassTileConfig> get_candidate_tiles(
-    int const sm, CutlassGemmConfig::CandidateConfigTypeParam const config_type_param)
-{
-    enum class CutlassGemmType : char
-    {
-        Default,
-        WeightOnly,
-        Simt,
-        Int8,
-        Fp8
-    };
+    int const sm,
+    CutlassGemmConfig::CandidateConfigTypeParam const config_type_param) {
+  enum class CutlassGemmType : char { Default, WeightOnly, Simt, Int8, Fp8 };
 
-    CutlassGemmType gemm_type = CutlassGemmType::Default;
-    if (config_type_param & CutlassGemmConfig::SIMT_ONLY)
-    {
-        gemm_type = CutlassGemmType::Simt;
-    }
-    else if (config_type_param & CutlassGemmConfig::WEIGHT_ONLY)
-    {
-        gemm_type = CutlassGemmType::WeightOnly;
-    }
-    else if (config_type_param & CutlassGemmConfig::INT8_ONLY)
-    {
-        gemm_type = CutlassGemmType::Int8;
-    }
-    else if (config_type_param & CutlassGemmConfig::FP8_ONLY)
-    {
-        gemm_type = CutlassGemmType::Fp8;
-    }
+  CutlassGemmType gemm_type = CutlassGemmType::Default;
+  if (config_type_param & CutlassGemmConfig::SIMT_ONLY) {
+    gemm_type = CutlassGemmType::Simt;
+  } else if (config_type_param & CutlassGemmConfig::WEIGHT_ONLY) {
+    gemm_type = CutlassGemmType::WeightOnly;
+  } else if (config_type_param & CutlassGemmConfig::INT8_ONLY) {
+    gemm_type = CutlassGemmType::Int8;
+  } else if (config_type_param & CutlassGemmConfig::FP8_ONLY) {
+    gemm_type = CutlassGemmType::Fp8;
+  }
 
-    std::vector<CutlassTileConfig> base_configs{
-        CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64, CutlassTileConfig::CtaShape64x128x64_WarpShape32x64x64};
-    if (sm >= 75)
-    {
-        base_configs.push_back(CutlassTileConfig::CtaShape128x128x64_WarpShape64x32x64);
-    }
+  std::vector<CutlassTileConfig> base_configs{
+      CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64,
+      CutlassTileConfig::CtaShape64x128x64_WarpShape32x64x64};
+  if (sm >= 75) {
+    base_configs.push_back(
+        CutlassTileConfig::CtaShape128x128x64_WarpShape64x32x64);
+  }
 
-    switch (gemm_type)
-    {
-    case CutlassGemmType::Simt: return {CutlassTileConfig::CtaShape128x128x8_WarpShape64x64x8};
+  switch (gemm_type) {
+    case CutlassGemmType::Simt:
+      return {CutlassTileConfig::CtaShape128x128x8_WarpShape64x64x8};
     case CutlassGemmType::WeightOnly:
-        if (sm >= 75)
-        {
-            return {CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64,
+      if (sm >= 75) {
+        return {CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64,
                 CutlassTileConfig::CtaShape16x256x64_WarpShape16x64x64,
                 CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64,
                 CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64,
                 CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64};
-        }
-        else
-        {
-            return {CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64,
-                CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64};
-        }
-    case CutlassGemmType::Int8:
+      } else {
         return {CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64,
-            CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64,
-            CutlassTileConfig::CtaShape128x64x64_WarpShape64x32x64,
-            CutlassTileConfig::CtaShape64x64x128_WarpShape32x64x64,
-            CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64,
-            CutlassTileConfig::CtaShape256x128x64_WarpShape64x64x64};
+                CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64};
+      }
+    case CutlassGemmType::Int8:
+      return {CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64,
+              CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64,
+              CutlassTileConfig::CtaShape128x64x64_WarpShape64x32x64,
+              CutlassTileConfig::CtaShape64x64x128_WarpShape32x64x64,
+              CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64,
+              CutlassTileConfig::CtaShape256x128x64_WarpShape64x64x64};
     case CutlassGemmType::Fp8:
-        if (config_type_param & CutlassGemmConfig::GROUPED_GEMM)
-        {
-            if (sm == 89)
-            {
-                return {CutlassTileConfig::CtaShape16x256x128_WarpShape16x64x128,
-                    CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64,
-                    CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64,
-                    CutlassTileConfig::CtaShape64x64x128_WarpShape32x64x64,
-                    CutlassTileConfig::CtaShape128x64x64_WarpShape64x32x64,
-                    CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64,
-                    CutlassTileConfig::CtaShape256x128x64_WarpShape64x64x64};
-            }
-            else
-            {
-                // no valid ampere style fp8 configs for sm90
-                return {};
-            }
+      if (config_type_param & CutlassGemmConfig::GROUPED_GEMM) {
+        if (sm == 89) {
+          return {CutlassTileConfig::CtaShape16x256x128_WarpShape16x64x128,
+                  CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64,
+                  CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64,
+                  CutlassTileConfig::CtaShape64x64x128_WarpShape32x64x64,
+                  CutlassTileConfig::CtaShape128x64x64_WarpShape64x32x64,
+                  CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64,
+                  CutlassTileConfig::CtaShape256x128x64_WarpShape64x64x64};
+        } else {
+          // no valid ampere style fp8 configs for sm90
+          return {};
         }
-    default: return base_configs;
-    }
+      }
+    default:
+      return base_configs;
+  }
 }
 
 std::vector<CutlassTileConfigSM90> get_candidate_tiles_sm90(
-    int const sm, CutlassGemmConfig::CandidateConfigTypeParam const config)
-{
+    int const sm, CutlassGemmConfig::CandidateConfigTypeParam const config) {
 #ifdef FAST_BUILD
-    // Fast build disables all configs except this one for SM90
-    return {CutlassTileConfigSM90::CtaShape128x128x128B};
+  // Fast build disables all configs except this one for SM90
+  return {CutlassTileConfigSM90::CtaShape128x128x128B};
 #else
-    if (config & CutlassGemmConfig::GROUPED_GEMM)
-    {
-        return {CutlassTileConfigSM90::CtaShape128x16x128B, CutlassTileConfigSM90::CtaShape128x32x128B,
-            CutlassTileConfigSM90::CtaShape128x64x128B, CutlassTileConfigSM90::CtaShape128x128x128B,
-            CutlassTileConfigSM90::CtaShape128x256x128B, CutlassTileConfigSM90::CtaShape256x128x128B};
-    }
-    else
-    {
-        return {CutlassTileConfigSM90::CtaShape64x16x128B, CutlassTileConfigSM90::CtaShape64x32x128B,
-            CutlassTileConfigSM90::CtaShape64x64x128B, CutlassTileConfigSM90::CtaShape64x128x128B,
-            CutlassTileConfigSM90::CtaShape64x256x128B, CutlassTileConfigSM90::CtaShape128x16x128B,
-            CutlassTileConfigSM90::CtaShape128x32x128B, CutlassTileConfigSM90::CtaShape128x64x128B,
-            CutlassTileConfigSM90::CtaShape128x128x128B, CutlassTileConfigSM90::CtaShape128x256x128B};
-    }
+  if (config & CutlassGemmConfig::GROUPED_GEMM) {
+    return {CutlassTileConfigSM90::CtaShape128x16x128B,
+            CutlassTileConfigSM90::CtaShape128x32x128B,
+            CutlassTileConfigSM90::CtaShape128x64x128B,
+            CutlassTileConfigSM90::CtaShape128x128x128B,
+            CutlassTileConfigSM90::CtaShape128x256x128B,
+            CutlassTileConfigSM90::CtaShape256x128x128B};
+  } else {
+    return {CutlassTileConfigSM90::CtaShape64x16x128B,
+            CutlassTileConfigSM90::CtaShape64x32x128B,
+            CutlassTileConfigSM90::CtaShape64x64x128B,
+            CutlassTileConfigSM90::CtaShape64x128x128B,
+            CutlassTileConfigSM90::CtaShape64x256x128B,
+            CutlassTileConfigSM90::CtaShape128x16x128B,
+            CutlassTileConfigSM90::CtaShape128x32x128B,
+            CutlassTileConfigSM90::CtaShape128x64x128B,
+            CutlassTileConfigSM90::CtaShape128x128x128B,
+            CutlassTileConfigSM90::CtaShape128x256x128B};
+  }
 #endif
 }
 
-// We only compile CUTLASS kernels with multi-cast along M if the M tile is >= 128. This is purely to improve
-// compilation speed.
-bool supports_mcast_along_m(CutlassTileConfigSM90 const tile)
-{
+// We only compile CUTLASS kernels with multi-cast along M if the M tile is >=
+// 128. This is purely to improve compilation speed.
+bool supports_mcast_along_m(CutlassTileConfigSM90 const tile) {
 #ifdef FAST_BUILD
-    return false;
+  return false;
 #else
-    std::set<CutlassTileConfigSM90> valid_tiles{CutlassTileConfigSM90::CtaShape128x16x128B,
-        CutlassTileConfigSM90::CtaShape128x32x128B, CutlassTileConfigSM90::CtaShape128x64x128B,
-        CutlassTileConfigSM90::CtaShape128x128x128B, CutlassTileConfigSM90::CtaShape128x256x128B,
-        CutlassTileConfigSM90::CtaShape256x128x128B};
-    return valid_tiles.count(tile) == 1;
+  std::set<CutlassTileConfigSM90> valid_tiles{
+      CutlassTileConfigSM90::CtaShape128x16x128B,
+      CutlassTileConfigSM90::CtaShape128x32x128B,
+      CutlassTileConfigSM90::CtaShape128x64x128B,
+      CutlassTileConfigSM90::CtaShape128x128x128B,
+      CutlassTileConfigSM90::CtaShape128x256x128B,
+      CutlassTileConfigSM90::CtaShape256x128x128B};
+  return valid_tiles.count(tile) == 1;
 #endif
 }
 
-// We only compile CUTLASS kernels with multi-cast along N if the N tile is >= 128. This is purely to improve
-// compilation speed.
-bool supports_mcast_along_n(CutlassTileConfigSM90 const tile)
-{
+// We only compile CUTLASS kernels with multi-cast along N if the N tile is >=
+// 128. This is purely to improve compilation speed.
+bool supports_mcast_along_n(CutlassTileConfigSM90 const tile) {
 #ifdef FAST_BUILD
-    return false;
+  return false;
 #else
-    std::set<CutlassTileConfigSM90> valid_tiles{CutlassTileConfigSM90::CtaShape64x128x128B,
-        CutlassTileConfigSM90::CtaShape64x256x128B, CutlassTileConfigSM90::CtaShape128x128x128B,
-        CutlassTileConfigSM90::CtaShape128x256x128B, CutlassTileConfigSM90::CtaShape256x128x128B};
-    return valid_tiles.count(tile) == 1;
+  std::set<CutlassTileConfigSM90> valid_tiles{
+      CutlassTileConfigSM90::CtaShape64x128x128B,
+      CutlassTileConfigSM90::CtaShape64x256x128B,
+      CutlassTileConfigSM90::CtaShape128x128x128B,
+      CutlassTileConfigSM90::CtaShape128x256x128B,
+      CutlassTileConfigSM90::CtaShape256x128x128B};
+  return valid_tiles.count(tile) == 1;
 #endif
 }
 
 // SM100 (Blackwell) candidate tile configurations
 std::vector<CutlassTileConfigSM100> get_candidate_tiles_sm100(
-    int /*sm*/, CutlassGemmConfig::CandidateConfigTypeParam const config)
-{
+    int /*sm*/, CutlassGemmConfig::CandidateConfigTypeParam const config) {
 #ifdef FAST_BUILD
-    return {CutlassTileConfigSM100::CtaShape128x128x128B};
+  return {CutlassTileConfigSM100::CtaShape128x128x128B};
 #else
-    /*  Grouped-GEMM path first (Blackwell uses 1-SM and 2-SM “cluster” kernels) */
-    if (config & CutlassGemmConfig::GROUPED_GEMM)
+  /*  Grouped-GEMM path first (Blackwell uses 1-SM and 2-SM “cluster” kernels)
+   */
+  if (config & CutlassGemmConfig::GROUPED_GEMM) {
+    if (config & CutlassGemmConfig::FP4_ONLY)  // nvfp4 / mx_fp4
     {
-        if (config & CutlassGemmConfig::FP4_ONLY)        // nvfp4 / mx_fp4
-        {
-            return {
-                /* 1 SM (M=128) */
-                CutlassTileConfigSM100::CtaShape128x128x128B,
-                CutlassTileConfigSM100::CtaShape128x256x128B,
-                /* 2 SM (M=256) */
-                CutlassTileConfigSM100::CtaShape256x128x128B,
-                CutlassTileConfigSM100::CtaShape256x256x128B,
-                /* slim tiles for very tall matrices */
-                CutlassTileConfigSM100::CtaShape128x64x128B,
-                CutlassTileConfigSM100::CtaShape256x64x128B};
-        }
+      return {/* 1 SM (M=128) */
+              CutlassTileConfigSM100::CtaShape128x128x128B,
+              CutlassTileConfigSM100::CtaShape128x256x128B,
+              /* 2 SM (M=256) */
+              CutlassTileConfigSM100::CtaShape256x128x128B,
+              CutlassTileConfigSM100::CtaShape256x256x128B,
+              /* slim tiles for very tall matrices */
+              CutlassTileConfigSM100::CtaShape128x64x128B,
+              CutlassTileConfigSM100::CtaShape256x64x128B};
+    }
 
-        /*  Fp8 / Fp16 grouped-GEMM      */
-        return {
-            CutlassTileConfigSM100::CtaShape128x128x128B,
+    /*  Fp8 / Fp16 grouped-GEMM      */
+    return {CutlassTileConfigSM100::CtaShape128x128x128B,
             CutlassTileConfigSM100::CtaShape128x256x128B,
             CutlassTileConfigSM100::CtaShape256x128x128B,
             CutlassTileConfigSM100::CtaShape256x256x128B};
-    }
+  }
 
-    /*  Non-grouped path (plain GEMM or weight-only)   */
-    return {
-        /* 1 SM tiles */
-        CutlassTileConfigSM100::CtaShape64x64x128B,
-        CutlassTileConfigSM100::CtaShape64x128x128B,
-        CutlassTileConfigSM100::CtaShape64x256x128B,
-        CutlassTileConfigSM100::CtaShape128x64x128B,
-        CutlassTileConfigSM100::CtaShape128x128x128B,
-        CutlassTileConfigSM100::CtaShape128x256x128B,
-        /* 2 SM tiles */
-        CutlassTileConfigSM100::CtaShape256x64x128B,
-        CutlassTileConfigSM100::CtaShape256x128x128B,
-        CutlassTileConfigSM100::CtaShape256x256x128B};
+  /*  Non-grouped path (plain GEMM or weight-only)   */
+  return {/* 1 SM tiles */
+          CutlassTileConfigSM100::CtaShape64x64x128B,
+          CutlassTileConfigSM100::CtaShape64x128x128B,
+          CutlassTileConfigSM100::CtaShape64x256x128B,
+          CutlassTileConfigSM100::CtaShape128x64x128B,
+          CutlassTileConfigSM100::CtaShape128x128x128B,
+          CutlassTileConfigSM100::CtaShape128x256x128B,
+          /* 2 SM tiles */
+          CutlassTileConfigSM100::CtaShape256x64x128B,
+          CutlassTileConfigSM100::CtaShape256x128x128B,
+          CutlassTileConfigSM100::CtaShape256x256x128B};
 #endif
 }
 
 // M-multicast support for SM100.
-bool supports_mcast_along_m_sm100(CutlassTileConfigSM100 tile)
-{
+bool supports_mcast_along_m_sm100(CutlassTileConfigSM100 tile) {
 #ifdef FAST_BUILD
-    return false;
+  return false;
 #else
-    std::set<CutlassTileConfigSM100> m_tiles{
-        CutlassTileConfigSM100::CtaShape128x64x128B,
-        CutlassTileConfigSM100::CtaShape128x128x128B,
-        CutlassTileConfigSM100::CtaShape128x256x128B,
-        CutlassTileConfigSM100::CtaShape256x64x128B,
-        CutlassTileConfigSM100::CtaShape256x128x128B,
-        CutlassTileConfigSM100::CtaShape256x256x128B};
-    return m_tiles.count(tile) == 1;
+  std::set<CutlassTileConfigSM100> m_tiles{
+      CutlassTileConfigSM100::CtaShape128x64x128B,
+      CutlassTileConfigSM100::CtaShape128x128x128B,
+      CutlassTileConfigSM100::CtaShape128x256x128B,
+      CutlassTileConfigSM100::CtaShape256x64x128B,
+      CutlassTileConfigSM100::CtaShape256x128x128B,
+      CutlassTileConfigSM100::CtaShape256x256x128B};
+  return m_tiles.count(tile) == 1;
 #endif
 }
 
 // N-multicast support for SM100.
-bool supports_mcast_along_n_sm100(CutlassTileConfigSM100 tile)
-{
+bool supports_mcast_along_n_sm100(CutlassTileConfigSM100 tile) {
 #ifdef FAST_BUILD
-    return false;
+  return false;
 #else
-    std::set<CutlassTileConfigSM100> n_tiles{
-        CutlassTileConfigSM100::CtaShape64x128x128B,
-        CutlassTileConfigSM100::CtaShape64x256x128B,
-        CutlassTileConfigSM100::CtaShape128x128x128B,
-        CutlassTileConfigSM100::CtaShape128x256x128B,
-        CutlassTileConfigSM100::CtaShape256x128x128B};
-    return n_tiles.count(tile) == 1;
+  std::set<CutlassTileConfigSM100> n_tiles{
+      CutlassTileConfigSM100::CtaShape64x128x128B,
+      CutlassTileConfigSM100::CtaShape64x256x128B,
+      CutlassTileConfigSM100::CtaShape128x128x128B,
+      CutlassTileConfigSM100::CtaShape128x256x128B,
+      CutlassTileConfigSM100::CtaShape256x128x128B};
+  return n_tiles.count(tile) == 1;
 #endif
 }
 
-
 std::vector<CutlassGemmConfig> get_candidate_configs(
-    int sm, int const max_split_k, CutlassGemmConfig::CandidateConfigTypeParam const config_type_param)
-{
-    if (sm == 90 && (config_type_param & CutlassGemmConfig::HOPPER))
-    {
-        std::vector<CutlassTileConfigSM90> tiles = get_candidate_tiles_sm90(sm, config_type_param);
+    int sm,
+    int const max_split_k,
+    CutlassGemmConfig::CandidateConfigTypeParam const config_type_param) {
+  if (sm == 90 && (config_type_param & CutlassGemmConfig::HOPPER)) {
+    std::vector<CutlassTileConfigSM90> tiles =
+        get_candidate_tiles_sm90(sm, config_type_param);
 
-        std::vector<CutlassGemmConfig> candidate_configs;
-        for (auto const& tile_config : tiles)
-        {
-            CutlassGemmConfig config(
-                tile_config, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO, ClusterShape::ClusterShape_1x1x1);
-            candidate_configs.push_back(config);
+    std::vector<CutlassGemmConfig> candidate_configs;
+    for (auto const& tile_config : tiles) {
+      CutlassGemmConfig config(tile_config,
+                               MainloopScheduleType::AUTO,
+                               EpilogueScheduleType::AUTO,
+                               ClusterShape::ClusterShape_1x1x1);
+      candidate_configs.push_back(config);
 
-            bool const has_m_mcast = supports_mcast_along_m(tile_config);
-            bool const has_n_mcast = supports_mcast_along_n(tile_config);
-            if (has_m_mcast)
-            {
-                CutlassGemmConfig config(tile_config, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO,
-                    ClusterShape::ClusterShape_2x1x1);
-                candidate_configs.push_back(config);
-            }
+      bool const has_m_mcast = supports_mcast_along_m(tile_config);
+      bool const has_n_mcast = supports_mcast_along_n(tile_config);
+      if (has_m_mcast) {
+        CutlassGemmConfig config(tile_config,
+                                 MainloopScheduleType::AUTO,
+                                 EpilogueScheduleType::AUTO,
+                                 ClusterShape::ClusterShape_2x1x1);
+        candidate_configs.push_back(config);
+      }
 
-            if (has_n_mcast)
-            {
-                CutlassGemmConfig config(tile_config, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO,
-                    ClusterShape::ClusterShape_1x2x1);
-                candidate_configs.push_back(config);
-            }
+      if (has_n_mcast) {
+        CutlassGemmConfig config(tile_config,
+                                 MainloopScheduleType::AUTO,
+                                 EpilogueScheduleType::AUTO,
+                                 ClusterShape::ClusterShape_1x2x1);
+        candidate_configs.push_back(config);
+      }
 
-            if (has_m_mcast && has_n_mcast)
-            {
-                CutlassGemmConfig config(tile_config, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO,
-                    ClusterShape::ClusterShape_2x2x1);
-                candidate_configs.push_back(config);
-            }
-        }
-        return candidate_configs;
+      if (has_m_mcast && has_n_mcast) {
+        CutlassGemmConfig config(tile_config,
+                                 MainloopScheduleType::AUTO,
+                                 EpilogueScheduleType::AUTO,
+                                 ClusterShape::ClusterShape_2x2x1);
+        candidate_configs.push_back(config);
+      }
     }
-    else if (sm == 100 && (config_type_param & CutlassGemmConfig::BLACKWELL)) // Assuming SM100 for Blackwell
-    {
-        std::vector<CutlassTileConfigSM100> tiles = get_candidate_tiles_sm100(sm, config_type_param);
-        std::vector<CutlassGemmConfig> candidate_configs;
-
-        for (auto const& tile_config_sm100 : tiles)
-        {
-            // SM100 uses MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO similar to SM90.
-            // Cluster shapes are also handled similarly.
-            CutlassGemmConfig config(
-                tile_config_sm100, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO, ClusterShape::ClusterShape_1x1x1);
-            candidate_configs.push_back(config);
-
-            bool const has_m_mcast = supports_mcast_along_m_sm100(tile_config_sm100);
-            bool const has_n_mcast = supports_mcast_along_n_sm100(tile_config_sm100);
-
-            if (has_m_mcast)
-            {
-                CutlassGemmConfig mcast_m_config(tile_config_sm100, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO,
-                    ClusterShape::ClusterShape_2x1x1);
-                candidate_configs.push_back(mcast_m_config);
-            }
-
-            if (has_n_mcast)
-            {
-                CutlassGemmConfig mcast_n_config(tile_config_sm100, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO,
-                    ClusterShape::ClusterShape_1x2x1);
-                candidate_configs.push_back(mcast_n_config);
-            }
-
-            if (has_m_mcast && has_n_mcast)
-            {
-                CutlassGemmConfig mcast_mn_config(tile_config_sm100, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO,
-                    ClusterShape::ClusterShape_2x2x1);
-                candidate_configs.push_back(mcast_mn_config);
-            }
-        }
-        return candidate_configs;
-    }
-
-    // Fallback to older architecture configurations
-    std::vector<CutlassTileConfig> tiles = get_candidate_tiles(sm, config_type_param);
-    std::vector<CutlassGemmConfig> candidate_configs; //Already declared above for SM90 path, ensure scope is correct or redeclare if necessary.
-                                                      // It's fine here as it's within an else if / else block.
-    bool const int8_configs_only = config_type_param & CutlassGemmConfig::INT8_ONLY;
-    int const min_stages = int8_configs_only ? 3 : 2;
-    int const max_stages = int8_configs_only ? 6 : (sm >= 80 ? 4 : 2);
-    for (auto const& tile_config : tiles)
-    {
-        for (int stages = min_stages; stages <= max_stages; ++stages)
-        {
-            CutlassGemmConfig config(tile_config, SplitKStyle::NO_SPLIT_K, 1, stages);
-            candidate_configs.push_back(config);
-            if (sm >= 75)
-            {
-                for (int split_k_factor = 2; split_k_factor <= max_split_k; ++split_k_factor)
-                {
-                    auto config = CutlassGemmConfig{tile_config, SplitKStyle::SPLIT_K_SERIAL, split_k_factor, stages};
-                    candidate_configs.push_back(config);
-                }
-            }
-        }
-    }
-
     return candidate_configs;
+  } else if (sm == 100 &&
+             (config_type_param &
+              CutlassGemmConfig::BLACKWELL))  // Assuming SM100 for Blackwell
+  {
+    std::vector<CutlassTileConfigSM100> tiles =
+        get_candidate_tiles_sm100(sm, config_type_param);
+    std::vector<CutlassGemmConfig> candidate_configs;
+
+    for (auto const& tile_config_sm100 : tiles) {
+      // SM100 uses MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO
+      // similar to SM90. Cluster shapes are also handled similarly.
+      CutlassGemmConfig config(tile_config_sm100,
+                               MainloopScheduleType::AUTO,
+                               EpilogueScheduleType::AUTO,
+                               ClusterShape::ClusterShape_1x1x1);
+      candidate_configs.push_back(config);
+
+      bool const has_m_mcast = supports_mcast_along_m_sm100(tile_config_sm100);
+      bool const has_n_mcast = supports_mcast_along_n_sm100(tile_config_sm100);
+
+      if (has_m_mcast) {
+        CutlassGemmConfig mcast_m_config(tile_config_sm100,
+                                         MainloopScheduleType::AUTO,
+                                         EpilogueScheduleType::AUTO,
+                                         ClusterShape::ClusterShape_2x1x1);
+        candidate_configs.push_back(mcast_m_config);
+      }
+
+      if (has_n_mcast) {
+        CutlassGemmConfig mcast_n_config(tile_config_sm100,
+                                         MainloopScheduleType::AUTO,
+                                         EpilogueScheduleType::AUTO,
+                                         ClusterShape::ClusterShape_1x2x1);
+        candidate_configs.push_back(mcast_n_config);
+      }
+
+      if (has_m_mcast && has_n_mcast) {
+        CutlassGemmConfig mcast_mn_config(tile_config_sm100,
+                                          MainloopScheduleType::AUTO,
+                                          EpilogueScheduleType::AUTO,
+                                          ClusterShape::ClusterShape_2x2x1);
+        candidate_configs.push_back(mcast_mn_config);
+      }
+    }
+    return candidate_configs;
+  }
+
+  // Fallback to older architecture configurations
+  std::vector<CutlassTileConfig> tiles =
+      get_candidate_tiles(sm, config_type_param);
+  std::vector<CutlassGemmConfig>
+      candidate_configs;  // Already declared above for SM90 path, ensure scope
+                          // is correct or redeclare if necessary.
+                          //  It's fine here as it's within an else if / else
+                          //  block.
+  bool const int8_configs_only =
+      config_type_param & CutlassGemmConfig::INT8_ONLY;
+  int const min_stages = int8_configs_only ? 3 : 2;
+  int const max_stages = int8_configs_only ? 6 : (sm >= 80 ? 4 : 2);
+  for (auto const& tile_config : tiles) {
+    for (int stages = min_stages; stages <= max_stages; ++stages) {
+      CutlassGemmConfig config(tile_config, SplitKStyle::NO_SPLIT_K, 1, stages);
+      candidate_configs.push_back(config);
+      if (sm >= 75) {
+        for (int split_k_factor = 2; split_k_factor <= max_split_k;
+             ++split_k_factor) {
+          auto config = CutlassGemmConfig{
+              tile_config, SplitKStyle::SPLIT_K_SERIAL, split_k_factor, stages};
+          candidate_configs.push_back(config);
+        }
+      }
+    }
+  }
+
+  return candidate_configs;
 }
 
-CutlassGemmConfig estimate_best_config_from_occupancies(std::vector<CutlassGemmConfig> const& candidate_configs,
-    std::vector<int> const& occupancies, int64_t const m, int64_t const n, int64_t const k, int64_t const num_experts,
-    int const split_k_limit, size_t const workspace_bytes, int const multi_processor_count, int const is_weight_only)
-{
+CutlassGemmConfig estimate_best_config_from_occupancies(
+    std::vector<CutlassGemmConfig> const& candidate_configs,
+    std::vector<int> const& occupancies,
+    int64_t const m,
+    int64_t const n,
+    int64_t const k,
+    int64_t const num_experts,
+    int const split_k_limit,
+    size_t const workspace_bytes,
+    int const multi_processor_count,
+    int const is_weight_only) {
+  if (occupancies.size() != candidate_configs.size()) {
+    throw(
+        "[estimate_best_config_from_occupancies] occpancies and "
+        "candidate configs vectors must have equal length.");
+  }
 
-    if (occupancies.size() != candidate_configs.size())
-    {
-        throw(
-            "[estimate_best_config_from_occupancies] occpancies and "
-            "candidate configs vectors must have equal length.");
+  CutlassGemmConfig best_config;
+  // Score will be [0, 1]. The objective is to minimize this score.
+  // It represents the fraction of SM resources unused in the last wave.
+  float config_score = 1.0f;
+  int config_waves = INT_MAX;
+  int current_m_tile = 0;
+
+  int const max_split_k = n >= multi_processor_count * 256 ? 1 : split_k_limit;
+  for (int ii = 0; ii < candidate_configs.size(); ++ii) {
+    CutlassGemmConfig candidate_config = candidate_configs[ii];
+    TileShape tile_shape =
+        get_cta_shape_for_config(candidate_config.tile_config);
+    int occupancy = occupancies[ii];
+
+    if (occupancy == 0) {
+      continue;
     }
 
-    CutlassGemmConfig best_config;
-    // Score will be [0, 1]. The objective is to minimize this score.
-    // It represents the fraction of SM resources unused in the last wave.
-    float config_score = 1.0f;
-    int config_waves = INT_MAX;
-    int current_m_tile = 0;
-
-    int const max_split_k = n >= multi_processor_count * 256 ? 1 : split_k_limit;
-    for (int ii = 0; ii < candidate_configs.size(); ++ii)
-    {
-        CutlassGemmConfig candidate_config = candidate_configs[ii];
-        TileShape tile_shape = get_cta_shape_for_config(candidate_config.tile_config);
-        int occupancy = occupancies[ii];
-
-        if (occupancy == 0)
-        {
-            continue;
-        }
-
-        // Keep small tile sizes when possible.
-        if (best_config.tile_config != CutlassTileConfig::ChooseWithHeuristic && m < current_m_tile
-            && current_m_tile < tile_shape.m)
-        {
-            continue;
-        }
-
-        int const ctas_in_m_dim = (m + tile_shape.m - 1) / tile_shape.m;
-        int const ctas_in_n_dim = (n + tile_shape.n - 1) / tile_shape.n;
-
-        for (int split_k_factor = 1; split_k_factor <= max_split_k; ++split_k_factor)
-        {
-            if (is_valid_split_k_factor(m, n, k, tile_shape, split_k_factor, workspace_bytes, is_weight_only))
-            {
-                int const ctas_per_wave = occupancy * multi_processor_count;
-                int const ctas_for_problem = ctas_in_m_dim * ctas_in_n_dim * split_k_factor;
-
-                int const num_waves_total = (ctas_for_problem + ctas_per_wave - 1) / ctas_per_wave;
-                float const num_waves_fractional = ctas_for_problem / float(ctas_per_wave);
-                float const current_score = float(num_waves_total) - num_waves_fractional;
-
-                float const score_slack = 0.1f;
-                if (current_score < config_score
-                    || ((config_waves > num_waves_total) && (current_score < config_score + score_slack)))
-                {
-                    config_score = current_score;
-                    config_waves = num_waves_total;
-                    SplitKStyle split_style
-                        = split_k_factor > 1 ? SplitKStyle::SPLIT_K_SERIAL : SplitKStyle::NO_SPLIT_K;
-                    best_config = CutlassGemmConfig(
-                        candidate_config.tile_config, split_style, split_k_factor, candidate_config.stages);
-                    current_m_tile = tile_shape.m;
-                }
-                else if (current_score == config_score
-                    && (best_config.stages < candidate_config.stages || split_k_factor < best_config.split_k_factor
-                        || current_m_tile < tile_shape.m))
-                {
-                    // Prefer deeper pipeline or smaller split-k
-                    SplitKStyle split_style
-                        = split_k_factor > 1 ? SplitKStyle::SPLIT_K_SERIAL : SplitKStyle::NO_SPLIT_K;
-                    best_config = CutlassGemmConfig(
-                        candidate_config.tile_config, split_style, split_k_factor, candidate_config.stages);
-                    current_m_tile = tile_shape.m;
-                    config_waves = num_waves_total;
-                }
-            }
-        }
+    // Keep small tile sizes when possible.
+    if (best_config.tile_config != CutlassTileConfig::ChooseWithHeuristic &&
+        m < current_m_tile && current_m_tile < tile_shape.m) {
+      continue;
     }
 
-    if (best_config.tile_config == CutlassTileConfig::ChooseWithHeuristic)
-    {
-        throw("Heurisitc failed to find a valid config.");
-    }
+    int const ctas_in_m_dim = (m + tile_shape.m - 1) / tile_shape.m;
+    int const ctas_in_n_dim = (n + tile_shape.n - 1) / tile_shape.n;
 
-    return best_config;
+    for (int split_k_factor = 1; split_k_factor <= max_split_k;
+         ++split_k_factor) {
+      if (is_valid_split_k_factor(m,
+                                  n,
+                                  k,
+                                  tile_shape,
+                                  split_k_factor,
+                                  workspace_bytes,
+                                  is_weight_only)) {
+        int const ctas_per_wave = occupancy * multi_processor_count;
+        int const ctas_for_problem =
+            ctas_in_m_dim * ctas_in_n_dim * split_k_factor;
+
+        int const num_waves_total =
+            (ctas_for_problem + ctas_per_wave - 1) / ctas_per_wave;
+        float const num_waves_fractional =
+            ctas_for_problem / float(ctas_per_wave);
+        float const current_score =
+            float(num_waves_total) - num_waves_fractional;
+
+        float const score_slack = 0.1f;
+        if (current_score < config_score ||
+            ((config_waves > num_waves_total) &&
+             (current_score < config_score + score_slack))) {
+          config_score = current_score;
+          config_waves = num_waves_total;
+          SplitKStyle split_style = split_k_factor > 1
+                                        ? SplitKStyle::SPLIT_K_SERIAL
+                                        : SplitKStyle::NO_SPLIT_K;
+          best_config = CutlassGemmConfig(candidate_config.tile_config,
+                                          split_style,
+                                          split_k_factor,
+                                          candidate_config.stages);
+          current_m_tile = tile_shape.m;
+        } else if (current_score == config_score &&
+                   (best_config.stages < candidate_config.stages ||
+                    split_k_factor < best_config.split_k_factor ||
+                    current_m_tile < tile_shape.m)) {
+          // Prefer deeper pipeline or smaller split-k
+          SplitKStyle split_style = split_k_factor > 1
+                                        ? SplitKStyle::SPLIT_K_SERIAL
+                                        : SplitKStyle::NO_SPLIT_K;
+          best_config = CutlassGemmConfig(candidate_config.tile_config,
+                                          split_style,
+                                          split_k_factor,
+                                          candidate_config.stages);
+          current_m_tile = tile_shape.m;
+          config_waves = num_waves_total;
+        }
+      }
+    }
+  }
+
+  if (best_config.tile_config == CutlassTileConfig::ChooseWithHeuristic) {
+    throw("Heurisitc failed to find a valid config.");
+  }
+
+  return best_config;
 }
 
-} // namespace cutlass_kernels
-} // namespace kernels
+}  // namespace cutlass_kernels
+}  // namespace kernels
diff --git a/custom_ops/gpu_ops/cutlass_kernels/cutlass_heuristic.h b/custom_ops/gpu_ops/cutlass_kernels/cutlass_heuristic.h
index 8165bc421c..b6839be7f3 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/cutlass_heuristic.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/cutlass_heuristic.h
@@ -20,36 +20,47 @@
 #include "cutlass_extensions/gemm_configs.h"
 #include "common/cudaUtils.h"
 
-namespace kernels
-{
-namespace cutlass_kernels
-{
+namespace kernels {
+namespace cutlass_kernels {
 
 template <class TileShape, class ClusterShape, class ActivationType>
-struct should_filter_sm90_gemm_problem_shape
-{
+struct should_filter_sm90_gemm_problem_shape {
 #ifdef FAST_BUILD
-    constexpr static int TILE_K = 128 * 8 / cutlass::sizeof_bits<ActivationType>::value;
-    using SupportedCtaShape = cute::Shape<cute::_128, cute::_128, cute::Int<TILE_K>>;
-    using SupportedCgaShape = cute::Shape<cute::_1, cute::_1, cute::_1>;
+  constexpr static int TILE_K =
+      128 * 8 / cutlass::sizeof_bits<ActivationType>::value;
+  using SupportedCtaShape =
+      cute::Shape<cute::_128, cute::_128, cute::Int<TILE_K>>;
+  using SupportedCgaShape = cute::Shape<cute::_1, cute::_1, cute::_1>;
 
-    constexpr static bool value
-        = !cute::is_same_v<SupportedCtaShape, TileShape> || !cute::is_same_v<SupportedCgaShape, ClusterShape>;
+  constexpr static bool value =
+      !cute::is_same_v<SupportedCtaShape, TileShape> ||
+      !cute::is_same_v<SupportedCgaShape, ClusterShape>;
 #else
-    constexpr static bool value = false;
+  constexpr static bool value = false;
 #endif
 };
 template <class TileShape, class ClusterShape, class ActivationType>
-constexpr static bool should_filter_sm90_gemm_problem_shape_v
-    = should_filter_sm90_gemm_problem_shape<TileShape, ClusterShape, ActivationType>::value;
+constexpr static bool should_filter_sm90_gemm_problem_shape_v =
+    should_filter_sm90_gemm_problem_shape<TileShape,
+                                          ClusterShape,
+                                          ActivationType>::value;
 
 std::vector<cutlass_extensions::CutlassGemmConfig> get_candidate_configs(
-    int sm, int const max_split_k, cutlass_extensions::CutlassGemmConfig::CandidateConfigTypeParam const);
+    int sm,
+    int const max_split_k,
+    cutlass_extensions::CutlassGemmConfig::CandidateConfigTypeParam const);
 
 cutlass_extensions::CutlassGemmConfig estimate_best_config_from_occupancies(
     std::vector<cutlass_extensions::CutlassGemmConfig> const& candidate_configs,
-    std::vector<int> const& occupancies, int64_t const m, int64_t const n, int64_t const k, int64_t const num_experts,
-    int const split_k_limit, size_t const workspace_bytes, int const multi_processor_count, int const is_weight_only);
+    std::vector<int> const& occupancies,
+    int64_t const m,
+    int64_t const n,
+    int64_t const k,
+    int64_t const num_experts,
+    int const split_k_limit,
+    size_t const workspace_bytes,
+    int const multi_processor_count,
+    int const is_weight_only);
 
-} // namespace cutlass_kernels
-} // namespace kernels
+}  // namespace cutlass_kernels
+}  // namespace kernels
diff --git a/custom_ops/gpu_ops/cutlass_kernels/cutlass_preprocessors.cu b/custom_ops/gpu_ops/cutlass_kernels/cutlass_preprocessors.cu
index 5c62b58085..41c0412cd9 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/cutlass_preprocessors.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/cutlass_preprocessors.cu
@@ -19,752 +19,781 @@
 #include "cutlass_kernels/cutlass_preprocessors.h"
 #include "cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h"
 
-namespace kernels
-{
-namespace cutlass_kernels
-{
+namespace kernels {
+namespace cutlass_kernels {
 
-struct LayoutDetails
-{
-    enum class Layout
-    {
-        UNKNOWN,
-        ROW_MAJOR,
-        COLUMN_MAJOR
-    };
+struct LayoutDetails {
+  enum class Layout { UNKNOWN, ROW_MAJOR, COLUMN_MAJOR };
 
-    Layout layoutB = Layout::UNKNOWN;
-    int rows_per_column_tile = 1;
-    int columns_interleaved = 1;
+  Layout layoutB = Layout::UNKNOWN;
+  int rows_per_column_tile = 1;
+  int columns_interleaved = 1;
 
-    bool uses_imma_ldsm = false;
+  bool uses_imma_ldsm = false;
 };
 
 template <typename Layout>
-struct getLayoutDetails
-{
+struct getLayoutDetails {};
+
+template <>
+struct getLayoutDetails<cutlass::layout::RowMajor> {
+  LayoutDetails operator()() {
+    LayoutDetails layout_details;
+    layout_details.layoutB = LayoutDetails::Layout::ROW_MAJOR;
+    return layout_details;
+  }
 };
 
 template <>
-struct getLayoutDetails<cutlass::layout::RowMajor>
-{
-    LayoutDetails operator()()
-    {
-        LayoutDetails layout_details;
-        layout_details.layoutB = LayoutDetails::Layout::ROW_MAJOR;
-        return layout_details;
-    }
-};
-
-template <>
-struct getLayoutDetails<cutlass::layout::ColumnMajor>
-{
-    LayoutDetails operator()()
-    {
-        LayoutDetails layout_details;
-        layout_details.layoutB = LayoutDetails::Layout::COLUMN_MAJOR;
-        return layout_details;
-    }
+struct getLayoutDetails<cutlass::layout::ColumnMajor> {
+  LayoutDetails operator()() {
+    LayoutDetails layout_details;
+    layout_details.layoutB = LayoutDetails::Layout::COLUMN_MAJOR;
+    return layout_details;
+  }
 };
 
 template <int RowsPerTile, int ColumnsInterleaved>
-struct getLayoutDetails<cutlass::layout::ColumnMajorTileInterleave<RowsPerTile, ColumnsInterleaved>>
-{
-    LayoutDetails operator()()
-    {
-        LayoutDetails layout_details;
-        layout_details.layoutB = LayoutDetails::Layout::COLUMN_MAJOR;
-        layout_details.rows_per_column_tile = RowsPerTile;
-        layout_details.columns_interleaved = ColumnsInterleaved;
-        return layout_details;
-    }
+struct getLayoutDetails<
+    cutlass::layout::ColumnMajorTileInterleave<RowsPerTile,
+                                               ColumnsInterleaved>> {
+  LayoutDetails operator()() {
+    LayoutDetails layout_details;
+    layout_details.layoutB = LayoutDetails::Layout::COLUMN_MAJOR;
+    layout_details.rows_per_column_tile = RowsPerTile;
+    layout_details.columns_interleaved = ColumnsInterleaved;
+    return layout_details;
+  }
 };
 
 template <typename cutlassArch, typename TypeA, typename TypeB>
-LayoutDetails getLayoutDetailsForArchAndQuantType()
-{
-
-    using CompileTraits = cutlass::gemm::kernel::LayoutDetailsB<TypeA, TypeB, cutlassArch>;
-    using LayoutB = typename CompileTraits::Layout;
-    using MmaOperator = typename CompileTraits::Operator;
-    LayoutDetails details = getLayoutDetails<LayoutB>()();
-    details.uses_imma_ldsm = std::is_same<MmaOperator, cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA>::value;
-    return details;
+LayoutDetails getLayoutDetailsForArchAndQuantType() {
+  using CompileTraits =
+      cutlass::gemm::kernel::LayoutDetailsB<TypeA, TypeB, cutlassArch>;
+  using LayoutB = typename CompileTraits::Layout;
+  using MmaOperator = typename CompileTraits::Operator;
+  LayoutDetails details = getLayoutDetails<LayoutB>()();
+  details.uses_imma_ldsm = std::is_same<
+      MmaOperator,
+      cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA>::value;
+  return details;
 }
 
 template <typename cutlassArch>
-LayoutDetails getLayoutDetailsForArch(QuantType quant_type)
-{
-    int const bits_per_weight_element = get_weight_quant_bits(quant_type);
-    LayoutDetails details;
-    switch (quant_type)
-    {
+LayoutDetails getLayoutDetailsForArch(QuantType quant_type) {
+  int const bits_per_weight_element = get_weight_quant_bits(quant_type);
+  LayoutDetails details;
+  switch (quant_type) {
     case QuantType::W8_A16:
-        details = getLayoutDetailsForArchAndQuantType<cutlassArch, cutlass::half_t, uint8_t>();
-        break;
+      details = getLayoutDetailsForArchAndQuantType<cutlassArch,
+                                                    cutlass::half_t,
+                                                    uint8_t>();
+      break;
     case QuantType::W4_A16:
-        details = getLayoutDetailsForArchAndQuantType<cutlassArch, cutlass::half_t, cutlass::uint4b_t>();
-        break;
+      details = getLayoutDetailsForArchAndQuantType<cutlassArch,
+                                                    cutlass::half_t,
+                                                    cutlass::uint4b_t>();
+      break;
     case QuantType::W4_AFP8:
-        details = getLayoutDetailsForArchAndQuantType<cutlassArch, cutlass::float_e4m3_t, cutlass::uint4b_t>();
-        break;
-    default: PADDLE_THROW("Unsupported quantization type");
-    }
-    return details;
+      details = getLayoutDetailsForArchAndQuantType<cutlassArch,
+                                                    cutlass::float_e4m3_t,
+                                                    cutlass::uint4b_t>();
+      break;
+    default:
+      PADDLE_THROW("Unsupported quantization type");
+  }
+  return details;
 }
 
-LayoutDetails getLayoutDetailsForTransform(QuantType quant_type, int arch)
-{
-    if (arch >= 70 && arch < 75)
-    {
-        return getLayoutDetailsForArch<cutlass::arch::Sm70>(quant_type);
-    }
-    else if (arch >= 75 && arch < 80)
-    {
-        return getLayoutDetailsForArch<cutlass::arch::Sm75>(quant_type);
-    }
-    else if (arch >= 80 && arch < 90)
-    {
-        return getLayoutDetailsForArch<cutlass::arch::Sm80>(quant_type);
-    }
-    else if (arch == 90)
-    {
-        return getLayoutDetailsForArch<cutlass::arch::Sm90>(quant_type);
-    }
-    else
-    {
-        PADDLE_ENFORCE(false, "Unsupported Arch");
-        return LayoutDetails();
-    }
+LayoutDetails getLayoutDetailsForTransform(QuantType quant_type, int arch) {
+  if (arch >= 70 && arch < 75) {
+    return getLayoutDetailsForArch<cutlass::arch::Sm70>(quant_type);
+  } else if (arch >= 75 && arch < 80) {
+    return getLayoutDetailsForArch<cutlass::arch::Sm75>(quant_type);
+  } else if (arch >= 80 && arch < 90) {
+    return getLayoutDetailsForArch<cutlass::arch::Sm80>(quant_type);
+  } else if (arch == 90) {
+    return getLayoutDetailsForArch<cutlass::arch::Sm90>(quant_type);
+  } else {
+    PADDLE_ENFORCE(false, "Unsupported Arch");
+    return LayoutDetails();
+  }
 }
 
-// Permutes the rows of B in a way that is compatible with Turing+ architectures.
+// Permutes the rows of B in a way that is compatible with Turing+
+// architectures.
 //
 // Throws an error for other architectures.
 // The data is permuted such that:
 // For W8_A16, each group of 16 rows is permuted using the map below:
 //  0 1 8 9 2 3 10 11 4 5 12 13 6 7 14 15
 // For W4_A16, each group of 32 rows is permuted using the map below:
-//  0 1 8 9 16 17 24 25 2 3 10 11 18 19 26 27 4 5 12 13 20 21 28 29 6 7 14 15 22 23 30 31
+//  0 1 8 9 16 17 24 25 2 3 10 11 18 19 26 27 4 5 12 13 20 21 28 29 6 7 14 15 22
+//  23 30 31
 // For W4_A8, see the map in the code. The idea is similar to above.
-// The goal of this permutation is to ensure data ends up in the correct threads after
-// we execute LDSM. It counteracts the effect of the data being of different widths.
-// For more information about the expected layouts, see the MMA section in the PTX docs.
-std::vector<int> get_permutation_map(QuantType quant_type)
-{
-
-    if (quant_type == QuantType::W8_A16)
-    {
-        return {0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15};
-    }
-    else if (quant_type == QuantType::W4_A16)
-    {
-        return {0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15,
-            22, 23, 30, 31};
-    }
-    else if (quant_type == QuantType::W4_AFP8)
-    {
-        return {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23, 8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15,
-            28, 29, 30, 31};
-    }
-    else
-    {
-        PADDLE_THROW("Invalid quantization type for LDSM permutation");
-    }
+// The goal of this permutation is to ensure data ends up in the correct threads
+// after we execute LDSM. It counteracts the effect of the data being of
+// different widths. For more information about the expected layouts, see the
+// MMA section in the PTX docs.
+std::vector<int> get_permutation_map(QuantType quant_type) {
+  if (quant_type == QuantType::W8_A16) {
+    return {0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15};
+  } else if (quant_type == QuantType::W4_A16) {
+    return {0, 1, 8,  9,  16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27,
+            4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31};
+  } else if (quant_type == QuantType::W4_AFP8) {
+    return {0, 1, 2,  3,  16, 17, 18, 19, 4,  5,  6,  7,  20, 21, 22, 23,
+            8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31};
+  } else {
+    PADDLE_THROW("Invalid quantization type for LDSM permutation");
+  }
 }
 
-void permute_B_rows_for_mixed_gemm(int8_t* permuted_quantized_tensor, int8_t const* quantized_tensor,
-    std::vector<size_t> const& shape, QuantType quant_type, int64_t const arch_version)
-{
-    // We only want to run this step for weight only quant.
-    std::vector<int> row_permutation = get_permutation_map(quant_type);
+void permute_B_rows_for_mixed_gemm(int8_t* permuted_quantized_tensor,
+                                   int8_t const* quantized_tensor,
+                                   std::vector<size_t> const& shape,
+                                   QuantType quant_type,
+                                   int64_t const arch_version) {
+  // We only want to run this step for weight only quant.
+  std::vector<int> row_permutation = get_permutation_map(quant_type);
 
-    PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3, "Shape must be 2-D or 3-D");
-    const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
-    const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
+  PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3,
+                 "Shape must be 2-D or 3-D");
+  const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    int const BITS_PER_ELT = get_weight_quant_bits(quant_type);
-    int const K = 16 / BITS_PER_ELT;
-    int const ELTS_PER_BYTE = 8 / BITS_PER_ELT;
-    int const ELTS_PER_REG = 32 / BITS_PER_ELT;
+  int const BITS_PER_ELT = get_weight_quant_bits(quant_type);
+  int const K = 16 / BITS_PER_ELT;
+  int const ELTS_PER_BYTE = 8 / BITS_PER_ELT;
+  int const ELTS_PER_REG = 32 / BITS_PER_ELT;
 
-    uint32_t const* input_byte_ptr = reinterpret_cast<uint32_t const*>(quantized_tensor);
-    uint32_t* output_byte_ptr = reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
+  uint32_t const* input_byte_ptr =
+      reinterpret_cast<uint32_t const*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
 
-    int MMA_SHAPE_N = 8;
-    int B_ROWS_PER_MMA = 8 * K;
-    int const elts_in_int32 = 32 / BITS_PER_ELT;
+  int MMA_SHAPE_N = 8;
+  int B_ROWS_PER_MMA = 8 * K;
+  int const elts_in_int32 = 32 / BITS_PER_ELT;
 
-    int const num_vec_cols = num_cols / elts_in_int32;
+  int const num_vec_cols = num_cols / elts_in_int32;
 
-    PADDLE_ENFORCE(
-        arch_version >= 75, "Unsupported Arch. Pre-volta not supported. Column interleave not needed on Volta.");
+  PADDLE_ENFORCE(arch_version >= 75,
+                 "Unsupported Arch. Pre-volta not supported. Column interleave "
+                 "not needed on Volta.");
 
-    PADDLE_ENFORCE(num_rows % B_ROWS_PER_MMA == 0,
-        "Invalid shape for quantized tensor. Number of rows of quantized matrix must be a multiple of %d",
-            B_ROWS_PER_MMA);
-    PADDLE_ENFORCE(num_cols % MMA_SHAPE_N == 0,
-        "Invalid shape for quantized tensor. On turing/Ampere, the number of cols must be a multiple of %d.",
-            MMA_SHAPE_N);
+  PADDLE_ENFORCE(num_rows % B_ROWS_PER_MMA == 0,
+                 "Invalid shape for quantized tensor. Number of rows of "
+                 "quantized matrix must be a multiple of %d",
+                 B_ROWS_PER_MMA);
+  PADDLE_ENFORCE(num_cols % MMA_SHAPE_N == 0,
+                 "Invalid shape for quantized tensor. On turing/Ampere, the "
+                 "number of cols must be a multiple of %d.",
+                 MMA_SHAPE_N);
 
-    PADDLE_ENFORCE(size_t(B_ROWS_PER_MMA) == row_permutation.size(), "Unexpected number of LDSM rows permuted.");
+  PADDLE_ENFORCE(size_t(B_ROWS_PER_MMA) == row_permutation.size(),
+                 "Unexpected number of LDSM rows permuted.");
 
-    for (int expert = 0; expert < num_experts; ++expert)
-    {
-        const int64_t matrix_offset = expert * int64_t(num_rows) * int64_t(num_vec_cols);
-        for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA)
-        {
-            for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row)
-            {
+  for (int expert = 0; expert < num_experts; ++expert) {
+    const int64_t matrix_offset =
+        expert * int64_t(num_rows) * int64_t(num_vec_cols);
+    for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
+      for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+        for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
+          int const write_row = base_row + tile_row;
+          int const tile_read_row = row_permutation[tile_row];
+          int const read_row = base_row + tile_read_row;
+          int const read_col = write_col;
 
-                for (int write_col = 0; write_col < num_vec_cols; ++write_col)
-                {
-                    int const write_row = base_row + tile_row;
-                    int const tile_read_row = row_permutation[tile_row];
-                    int const read_row = base_row + tile_read_row;
-                    int const read_col = write_col;
+          const int64_t read_offset =
+              matrix_offset + int64_t(read_row) * num_vec_cols + read_col;
+          const int64_t write_offset =
+              matrix_offset + int64_t(write_row) * num_vec_cols + write_col;
 
-                    const int64_t read_offset = matrix_offset + int64_t(read_row) * num_vec_cols + read_col;
-                    const int64_t write_offset = matrix_offset + int64_t(write_row) * num_vec_cols + write_col;
-
-                    output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-                }
-            }
+          output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
         }
+      }
     }
+  }
 }
 
 // We need to use this transpose to correctly handle packed int4 and int8 data
-// The reason this code is relatively complex is that the "trivial" loops took a substantial
-// amount of time to transpose leading to long preprocessing times. This seemed to be a big
-// issue for relatively large models.
+// The reason this code is relatively complex is that the "trivial" loops took a
+// substantial amount of time to transpose leading to long preprocessing times.
+// This seemed to be a big issue for relatively large models.
 template <QuantType quant_type>
-void subbyte_transpose_impl(
-    int8_t* transposed_quantized_tensor, int8_t const* quantized_tensor, std::vector<size_t> const& shape)
-{
-    constexpr int bits_per_elt = get_weight_quant_bits(quant_type);
+void subbyte_transpose_impl(int8_t* transposed_quantized_tensor,
+                            int8_t const* quantized_tensor,
+                            std::vector<size_t> const& shape) {
+  constexpr int bits_per_elt = get_weight_quant_bits(quant_type);
 
-    PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3, "Shape must be 2-D or 3-D");
-    const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
-    const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
+  PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3,
+                 "Shape must be 2-D or 3-D");
+  const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    const size_t col_bytes = num_cols * bits_per_elt / 8;
-    const size_t col_bytes_trans = num_rows * bits_per_elt / 8;
-    const size_t num_bytes = size_t(num_experts) * num_rows * col_bytes;
+  const size_t col_bytes = num_cols * bits_per_elt / 8;
+  const size_t col_bytes_trans = num_rows * bits_per_elt / 8;
+  const size_t num_bytes = size_t(num_experts) * num_rows * col_bytes;
 
-    uint8_t const* input_byte_ptr = reinterpret_cast<uint8_t const*>(quantized_tensor);
-    uint8_t* output_byte_ptr = reinterpret_cast<uint8_t*>(transposed_quantized_tensor);
+  uint8_t const* input_byte_ptr =
+      reinterpret_cast<uint8_t const*>(quantized_tensor);
+  uint8_t* output_byte_ptr =
+      reinterpret_cast<uint8_t*>(transposed_quantized_tensor);
 
-    static constexpr int ELTS_PER_BYTE = 8 / bits_per_elt;
+  static constexpr int ELTS_PER_BYTE = 8 / bits_per_elt;
 
-    static constexpr int M_TILE_L1 = 64;
-    static constexpr int N_TILE_L1 = M_TILE_L1 / ELTS_PER_BYTE;
-    uint8_t cache_buf[M_TILE_L1][N_TILE_L1];
+  static constexpr int M_TILE_L1 = 64;
+  static constexpr int N_TILE_L1 = M_TILE_L1 / ELTS_PER_BYTE;
+  uint8_t cache_buf[M_TILE_L1][N_TILE_L1];
 
-    static constexpr int VECTOR_WIDTH = std::min(32, N_TILE_L1);
+  static constexpr int VECTOR_WIDTH = std::min(32, N_TILE_L1);
 
-    // We assume the dims are a multiple of vector width. Our kernels only handle dims which are multiples
-    // of 64 for weight-only quantization. As a result, this seemed like a reasonable tradeoff because it
-    // allows GCC to emit vector instructions.
-    PADDLE_ENFORCE(!(col_bytes_trans % VECTOR_WIDTH) && !(col_bytes % VECTOR_WIDTH),
-        "Number of bytes for rows and cols must be a multiple of %d. However, num_rows_bytes = %ld and "
-               "num_col_bytes = %ld.",
-            VECTOR_WIDTH, col_bytes_trans, col_bytes);
+  // We assume the dims are a multiple of vector width. Our kernels only handle
+  // dims which are multiples of 64 for weight-only quantization. As a result,
+  // this seemed like a reasonable tradeoff because it allows GCC to emit vector
+  // instructions.
+  PADDLE_ENFORCE(
+      !(col_bytes_trans % VECTOR_WIDTH) && !(col_bytes % VECTOR_WIDTH),
+      "Number of bytes for rows and cols must be a multiple of %d. However, "
+      "num_rows_bytes = %ld and "
+      "num_col_bytes = %ld.",
+      VECTOR_WIDTH,
+      col_bytes_trans,
+      col_bytes);
 
-    int const num_m_tiles = (num_rows + M_TILE_L1 - 1) / M_TILE_L1;
-    int const num_n_tiles = (col_bytes + N_TILE_L1 - 1) / N_TILE_L1;
+  int const num_m_tiles = (num_rows + M_TILE_L1 - 1) / M_TILE_L1;
+  int const num_n_tiles = (col_bytes + N_TILE_L1 - 1) / N_TILE_L1;
 
-    for (size_t expert = 0; expert < num_experts; ++expert)
-    {
-        const size_t matrix_offset = expert * num_rows * col_bytes;
-        for (size_t row_tile_start = 0; row_tile_start < num_rows; row_tile_start += M_TILE_L1)
-        {
-            for (size_t col_tile_start_byte = 0; col_tile_start_byte < col_bytes; col_tile_start_byte += N_TILE_L1)
-            {
+  for (size_t expert = 0; expert < num_experts; ++expert) {
+    const size_t matrix_offset = expert * num_rows * col_bytes;
+    for (size_t row_tile_start = 0; row_tile_start < num_rows;
+         row_tile_start += M_TILE_L1) {
+      for (size_t col_tile_start_byte = 0; col_tile_start_byte < col_bytes;
+           col_tile_start_byte += N_TILE_L1) {
+        int const row_limit = std::min(row_tile_start + M_TILE_L1, num_rows);
+        int const col_limit =
+            std::min(col_tile_start_byte + N_TILE_L1, col_bytes);
 
-                int const row_limit = std::min(row_tile_start + M_TILE_L1, num_rows);
-                int const col_limit = std::min(col_tile_start_byte + N_TILE_L1, col_bytes);
+        for (int ii = 0; ii < M_TILE_L1; ++ii) {
+          int const row = row_tile_start + ii;
 
-                for (int ii = 0; ii < M_TILE_L1; ++ii)
-                {
-                    int const row = row_tile_start + ii;
+          for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
+            int const col = col_tile_start_byte + jj;
 
-                    for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH)
-                    {
-                        int const col = col_tile_start_byte + jj;
+            const size_t logical_src_offset =
+                matrix_offset + row * col_bytes + col;
 
-                        const size_t logical_src_offset = matrix_offset + row * col_bytes + col;
-
-                        if (row < row_limit && col < col_limit)
-                        {
-                            for (int v = 0; v < VECTOR_WIDTH; ++v)
-                            {
-                                cache_buf[ii][jj + v] = input_byte_ptr[logical_src_offset + v];
-                            }
-                        }
-                    }
-                }
-
-                if constexpr (bits_per_elt == 8)
-                {
-                    for (int ii = 0; ii < M_TILE_L1; ++ii)
-                    {
-                        for (int jj = ii + 1; jj < N_TILE_L1; ++jj)
-                        {
-                            std::swap(cache_buf[ii][jj], cache_buf[jj][ii]);
-                        }
-                    }
-                }
-                else if constexpr (bits_per_elt == 4)
-                {
-
-                    for (int ii = 0; ii < M_TILE_L1; ++ii)
-                    {
-                        // Using M_TILE_L1 here is deliberate since we assume that the cache tile
-                        // is square in the number of elements (not necessarily the number of bytes).
-                        for (int jj = ii + 1; jj < M_TILE_L1; ++jj)
-                        {
-                            int const ii_byte = ii / ELTS_PER_BYTE;
-                            int const ii_bit_offset = ii % ELTS_PER_BYTE;
-
-                            int const jj_byte = jj / ELTS_PER_BYTE;
-                            int const jj_bit_offset = jj % ELTS_PER_BYTE;
-
-                            uint8_t src_elt = 0xF & (cache_buf[ii][jj_byte] >> (4 * jj_bit_offset));
-                            uint8_t tgt_elt = 0xF & (cache_buf[jj][ii_byte] >> (4 * ii_bit_offset));
-
-                            cache_buf[ii][jj_byte] &= (0xF0 >> (4 * jj_bit_offset));
-                            cache_buf[jj][ii_byte] &= (0xF0 >> (4 * ii_bit_offset));
-
-                            cache_buf[ii][jj_byte] |= (tgt_elt << (4 * jj_bit_offset));
-                            cache_buf[jj][ii_byte] |= (src_elt << (4 * ii_bit_offset));
-                        }
-                    }
-                }
-                else
-                {
-                    PADDLE_ENFORCE(false, "Unsupported quantization type.");
-                }
-
-                const size_t row_tile_start_trans = col_tile_start_byte * ELTS_PER_BYTE;
-                const size_t col_tile_start_byte_trans = row_tile_start / ELTS_PER_BYTE;
-
-                int const row_limit_trans = std::min(row_tile_start_trans + M_TILE_L1, num_cols);
-                int const col_limit_trans = std::min(col_tile_start_byte_trans + N_TILE_L1, col_bytes_trans);
-
-                for (int ii = 0; ii < M_TILE_L1; ++ii)
-                {
-                    int const row = row_tile_start_trans + ii;
-                    for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH)
-                    {
-                        int const col = col_tile_start_byte_trans + jj;
-
-                        const size_t logical_tgt_offset = matrix_offset + row * col_bytes_trans + col;
-
-                        if (row < row_limit_trans && col < col_limit_trans)
-                        {
-                            for (int v = 0; v < VECTOR_WIDTH; ++v)
-                            {
-                                output_byte_ptr[logical_tgt_offset + v] = cache_buf[ii][jj + v];
-                            }
-                        }
-                    }
-                }
+            if (row < row_limit && col < col_limit) {
+              for (int v = 0; v < VECTOR_WIDTH; ++v) {
+                cache_buf[ii][jj + v] = input_byte_ptr[logical_src_offset + v];
+              }
             }
+          }
         }
-    }
-}
 
-void subbyte_transpose(int8_t* transposed_quantized_tensor, int8_t const* quantized_tensor,
-    std::vector<size_t> const& shape, QuantType quant_type)
-{
-    if (quant_type == QuantType::W8_A16)
-    {
-        subbyte_transpose_impl<QuantType::W8_A16>(transposed_quantized_tensor, quantized_tensor, shape);
-    }
-    else if (quant_type == QuantType::W4_A16)
-    {
-        subbyte_transpose_impl<QuantType::W4_A16>(transposed_quantized_tensor, quantized_tensor, shape);
-    }
-    else if (quant_type == QuantType::W4_AFP8)
-    {
-        subbyte_transpose_impl<QuantType::W4_AFP8>(transposed_quantized_tensor, quantized_tensor, shape);
-    }
-    else
-    {
-        PADDLE_ENFORCE(false, "Invalid quant_type");
-    }
-}
-
-void add_bias_and_interleave_int8s_inplace(int8_t* int8_tensor, const size_t num_elts)
-{
-    for (int ii = 0; ii < num_elts; ++ii)
-    {
-        int8_tensor[ii] = int8_t(int(int8_tensor[ii]) + 128);
-    }
-
-    // Step 2 will transform the layout of a 32-bit register in CUDA in order to match the int4 layout. This has no
-    // performance benefit and is purely so that int4 and int8 have the same layout.
-    // Pictorially, this does the following:
-    // bit 32                                                      0
-    //      [elt_3  elt_2  elt_1  elt_0] (each elt occupies 8 bits)
-    //
-    // And it will rearrange the output 32 bit register to be the following:
-    // bit 32                                                      0
-    //      [elt_3  elt_1  elt_2  elt_0] (each elt occupies 8 bits)
-
-    PADDLE_ENFORCE(num_elts % 4 == 0, "Dimensions of int8 tensor must be a multiple of 4 for register relayout");
-    for (size_t base = 0; base < num_elts; base += 4)
-    {
-        std::swap(int8_tensor[base + 1], int8_tensor[base + 2]);
-    }
-}
-
-void add_bias_and_interleave_int4s_inplace(int8_t* packed_int4_tensor, const size_t num_elts)
-{
-    int const num_bytes = num_elts / 2;
-
-    // Step 1 will be to transform all the int4s to unsigned in order to make the dequantize take as little
-    // instructions as possible in the CUDA code.
-    for (size_t ii = 0; ii < num_bytes; ++ii)
-    {
-        int8_t transformed_packed_int4s = 0;
-        int8_t transformed_first_elt
-            = (int8_t(packed_int4_tensor[ii] << 4) >> 4) + 8; // The double shift here is to ensure sign extension
-        int8_t transformed_second_elt = (packed_int4_tensor[ii] >> 4) + 8;
-
-        PADDLE_ENFORCE(
-            transformed_first_elt >= 0 && transformed_first_elt <= 15, "Illegal result for int4 transform (first elt)");
-        PADDLE_ENFORCE(transformed_second_elt >= 0 && transformed_second_elt <= 15,
-            "Illegal result for int4 transform (second elt)");
-
-        // We don't need to mask in these ops since everything should be in the range 0-15
-        transformed_packed_int4s |= transformed_first_elt;
-        transformed_packed_int4s |= (transformed_second_elt << 4);
-        packed_int4_tensor[ii] = transformed_packed_int4s;
-    }
-
-    // Step 2 will transform the layout of a 32-bit register in CUDA in order to minimize the number of shift & logical
-    // instructions That are needed to extract the int4s in the GEMM main loop. Pictorially, the loop below will do the
-    // following: Take as input a 32 bit register with layout: bit 32 0
-    //      [elt_7  elt_6  elt_5  elt_4  elt_3  elt_2  elt_1  elt_0] (each elt occupies 4 bits)
-    //
-    // And it will rearrange the output 32 bit register to be the following:
-    // bit 32                                                      0
-    //      [elt_7  elt_5  elt_3  elt_1  elt_6  elt_4  elt_2  elt_0] (each elt occupies 4 bits)
-
-    PADDLE_ENFORCE(num_bytes % 4 == 0, "Dimensions of int4 tensor must be a multiple of 8 for register relayout");
-    const size_t num_registers = num_bytes / 4;
-
-    uint32_t* register_ptr = reinterpret_cast<uint32_t*>(packed_int4_tensor);
-    for (size_t ii = 0; ii < num_registers; ++ii)
-    {
-        const uint32_t current_register = register_ptr[ii];
-        uint32_t transformed_register = 0;
-
-        for (int dest_idx = 0; dest_idx < 8; ++dest_idx)
-        {
-            int const src_idx = dest_idx < 4 ? 2 * dest_idx : 2 * (dest_idx - 4) + 1;
-            int const src_shift = 4 * src_idx;
-            int const dest_shift = 4 * dest_idx;
-
-            const uint32_t src_bits = (current_register >> src_shift) & 0xF;
-            transformed_register |= (src_bits << dest_shift);
-        }
-        register_ptr[ii] = transformed_register;
-    }
-}
-
-void add_bias_and_interleave_quantized_tensor_inplace(int8_t* tensor, const size_t num_elts, QuantType quant_type)
-{
-    if (quant_type == QuantType::W8_A16)
-    {
-        add_bias_and_interleave_int8s_inplace(tensor, num_elts);
-    }
-    else if (quant_type == QuantType::W4_A16 || quant_type == QuantType::W4_AFP8)
-    {
-        // W4_AFP8 uses the same preprocessor as W4_A16 because the FP8 data must
-        // be converted to FP16 before the scales can be applied using CUDA cores.
-        // As a result, we still want permute the data so that it is well aligned
-        // for conversion to FP16.
-        add_bias_and_interleave_int4s_inplace(tensor, num_elts);
-    }
-    else
-    {
-        PADDLE_ENFORCE(false, "Invalid quantization type for interleaving.");
-    }
-}
-
-void interleave_column_major_tensor(int8_t* interleaved_quantized_tensor, int8_t const* quantized_tensor,
-    std::vector<size_t> const& shape, QuantType quant_type, LayoutDetails details)
-{
-    PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3, "Shape must be 2-D or 3-D");
-    const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
-    const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
-
-    int const BITS_PER_ELT = get_weight_quant_bits(quant_type);
-    int const elts_in_int32 = 32 / BITS_PER_ELT;
-
-    int const rows_per_tile = details.rows_per_column_tile;
-
-    PADDLE_ENFORCE(!(num_rows % elts_in_int32),
-        "The number of rows must be a multiple of %d but the number of rows is %ld.", elts_in_int32, num_rows);
-
-    uint32_t const* input_byte_ptr = reinterpret_cast<uint32_t const*>(quantized_tensor);
-    uint32_t* output_byte_ptr = reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
-
-    PADDLE_ENFORCE(!(num_rows % rows_per_tile),
-        "The number of rows must be a multiple of %d but the number of rows is %ld.", rows_per_tile, num_rows);
-
-    int const num_vec_rows = num_rows / elts_in_int32;
-    int const vec_rows_per_tile = rows_per_tile / elts_in_int32;
-    int const interleave = details.columns_interleaved;
-
-    for (int expert = 0; expert < num_experts; ++expert)
-    {
-        const int64_t matrix_offset = expert * int64_t(num_vec_rows) * int64_t(num_cols);
-        for (int read_col = 0; read_col < num_cols; ++read_col)
-        {
-            const int64_t write_col = read_col / interleave;
-            for (int base_vec_row = 0; base_vec_row < num_vec_rows; base_vec_row += vec_rows_per_tile)
-            {
-                for (int vec_read_row = base_vec_row;
-                     vec_read_row < std::min(num_vec_rows, base_vec_row + vec_rows_per_tile); ++vec_read_row)
-                {
-                    const int64_t vec_write_row = interleave * base_vec_row
-                        + vec_rows_per_tile * (read_col % interleave) + vec_read_row % vec_rows_per_tile;
-
-                    const int64_t read_offset = matrix_offset + int64_t(read_col) * num_vec_rows + vec_read_row;
-                    const int64_t write_offset
-                        = matrix_offset + int64_t(write_col) * num_vec_rows * interleave + vec_write_row;
-                    output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-                }
+        if constexpr (bits_per_elt == 8) {
+          for (int ii = 0; ii < M_TILE_L1; ++ii) {
+            for (int jj = ii + 1; jj < N_TILE_L1; ++jj) {
+              std::swap(cache_buf[ii][jj], cache_buf[jj][ii]);
             }
+          }
+        } else if constexpr (bits_per_elt == 4) {
+          for (int ii = 0; ii < M_TILE_L1; ++ii) {
+            // Using M_TILE_L1 here is deliberate since we assume that the cache
+            // tile is square in the number of elements (not necessarily the
+            // number of bytes).
+            for (int jj = ii + 1; jj < M_TILE_L1; ++jj) {
+              int const ii_byte = ii / ELTS_PER_BYTE;
+              int const ii_bit_offset = ii % ELTS_PER_BYTE;
+
+              int const jj_byte = jj / ELTS_PER_BYTE;
+              int const jj_bit_offset = jj % ELTS_PER_BYTE;
+
+              uint8_t src_elt =
+                  0xF & (cache_buf[ii][jj_byte] >> (4 * jj_bit_offset));
+              uint8_t tgt_elt =
+                  0xF & (cache_buf[jj][ii_byte] >> (4 * ii_bit_offset));
+
+              cache_buf[ii][jj_byte] &= (0xF0 >> (4 * jj_bit_offset));
+              cache_buf[jj][ii_byte] &= (0xF0 >> (4 * ii_bit_offset));
+
+              cache_buf[ii][jj_byte] |= (tgt_elt << (4 * jj_bit_offset));
+              cache_buf[jj][ii_byte] |= (src_elt << (4 * ii_bit_offset));
+            }
+          }
+        } else {
+          PADDLE_ENFORCE(false, "Unsupported quantization type.");
         }
+
+        const size_t row_tile_start_trans = col_tile_start_byte * ELTS_PER_BYTE;
+        const size_t col_tile_start_byte_trans = row_tile_start / ELTS_PER_BYTE;
+
+        int const row_limit_trans =
+            std::min(row_tile_start_trans + M_TILE_L1, num_cols);
+        int const col_limit_trans =
+            std::min(col_tile_start_byte_trans + N_TILE_L1, col_bytes_trans);
+
+        for (int ii = 0; ii < M_TILE_L1; ++ii) {
+          int const row = row_tile_start_trans + ii;
+          for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
+            int const col = col_tile_start_byte_trans + jj;
+
+            const size_t logical_tgt_offset =
+                matrix_offset + row * col_bytes_trans + col;
+
+            if (row < row_limit_trans && col < col_limit_trans) {
+              for (int v = 0; v < VECTOR_WIDTH; ++v) {
+                output_byte_ptr[logical_tgt_offset + v] = cache_buf[ii][jj + v];
+              }
+            }
+          }
+        }
+      }
     }
+  }
 }
 
-void preprocess_weights_for_mixed_gemm(int8_t* preprocessed_quantized_weight, int8_t const* row_major_quantized_weight,
-    std::vector<size_t> const& shape, QuantType quant_type, bool force_interleave)
-{
-    int arch = 89;
-    if (force_interleave && arch == 90)
-    {
-        // Workaround for MOE which doesn't have specialised Hopper kernels yet
-        arch = 80;
+void subbyte_transpose(int8_t* transposed_quantized_tensor,
+                       int8_t const* quantized_tensor,
+                       std::vector<size_t> const& shape,
+                       QuantType quant_type) {
+  if (quant_type == QuantType::W8_A16) {
+    subbyte_transpose_impl<QuantType::W8_A16>(
+        transposed_quantized_tensor, quantized_tensor, shape);
+  } else if (quant_type == QuantType::W4_A16) {
+    subbyte_transpose_impl<QuantType::W4_A16>(
+        transposed_quantized_tensor, quantized_tensor, shape);
+  } else if (quant_type == QuantType::W4_AFP8) {
+    subbyte_transpose_impl<QuantType::W4_AFP8>(
+        transposed_quantized_tensor, quantized_tensor, shape);
+  } else {
+    PADDLE_ENFORCE(false, "Invalid quant_type");
+  }
+}
+
+void add_bias_and_interleave_int8s_inplace(int8_t* int8_tensor,
+                                           const size_t num_elts) {
+  for (int ii = 0; ii < num_elts; ++ii) {
+    int8_tensor[ii] = int8_t(int(int8_tensor[ii]) + 128);
+  }
+
+  // Step 2 will transform the layout of a 32-bit register in CUDA in order to
+  // match the int4 layout. This has no performance benefit and is purely so
+  // that int4 and int8 have the same layout. Pictorially, this does the
+  // following: bit 32                                                      0
+  //      [elt_3  elt_2  elt_1  elt_0] (each elt occupies 8 bits)
+  //
+  // And it will rearrange the output 32 bit register to be the following:
+  // bit 32                                                      0
+  //      [elt_3  elt_1  elt_2  elt_0] (each elt occupies 8 bits)
+
+  PADDLE_ENFORCE(num_elts % 4 == 0,
+                 "Dimensions of int8 tensor must be a multiple of 4 for "
+                 "register relayout");
+  for (size_t base = 0; base < num_elts; base += 4) {
+    std::swap(int8_tensor[base + 1], int8_tensor[base + 2]);
+  }
+}
+
+void add_bias_and_interleave_int4s_inplace(int8_t* packed_int4_tensor,
+                                           const size_t num_elts) {
+  int const num_bytes = num_elts / 2;
+
+  // Step 1 will be to transform all the int4s to unsigned in order to make the
+  // dequantize take as little instructions as possible in the CUDA code.
+  for (size_t ii = 0; ii < num_bytes; ++ii) {
+    int8_t transformed_packed_int4s = 0;
+    int8_t transformed_first_elt =
+        (int8_t(packed_int4_tensor[ii] << 4) >> 4) +
+        8;  // The double shift here is to ensure sign extension
+    int8_t transformed_second_elt = (packed_int4_tensor[ii] >> 4) + 8;
+
+    PADDLE_ENFORCE(transformed_first_elt >= 0 && transformed_first_elt <= 15,
+                   "Illegal result for int4 transform (first elt)");
+    PADDLE_ENFORCE(transformed_second_elt >= 0 && transformed_second_elt <= 15,
+                   "Illegal result for int4 transform (second elt)");
+
+    // We don't need to mask in these ops since everything should be in the
+    // range 0-15
+    transformed_packed_int4s |= transformed_first_elt;
+    transformed_packed_int4s |= (transformed_second_elt << 4);
+    packed_int4_tensor[ii] = transformed_packed_int4s;
+  }
+
+  // Step 2 will transform the layout of a 32-bit register in CUDA in order to
+  // minimize the number of shift & logical instructions That are needed to
+  // extract the int4s in the GEMM main loop. Pictorially, the loop below will
+  // do the following: Take as input a 32 bit register with layout: bit 32 0
+  //      [elt_7  elt_6  elt_5  elt_4  elt_3  elt_2  elt_1  elt_0] (each elt
+  //      occupies 4 bits)
+  //
+  // And it will rearrange the output 32 bit register to be the following:
+  // bit 32                                                      0
+  //      [elt_7  elt_5  elt_3  elt_1  elt_6  elt_4  elt_2  elt_0] (each elt
+  //      occupies 4 bits)
+
+  PADDLE_ENFORCE(num_bytes % 4 == 0,
+                 "Dimensions of int4 tensor must be a multiple of 8 for "
+                 "register relayout");
+  const size_t num_registers = num_bytes / 4;
+
+  uint32_t* register_ptr = reinterpret_cast<uint32_t*>(packed_int4_tensor);
+  for (size_t ii = 0; ii < num_registers; ++ii) {
+    const uint32_t current_register = register_ptr[ii];
+    uint32_t transformed_register = 0;
+
+    for (int dest_idx = 0; dest_idx < 8; ++dest_idx) {
+      int const src_idx = dest_idx < 4 ? 2 * dest_idx : 2 * (dest_idx - 4) + 1;
+      int const src_shift = 4 * src_idx;
+      int const dest_shift = 4 * dest_idx;
+
+      const uint32_t src_bits = (current_register >> src_shift) & 0xF;
+      transformed_register |= (src_bits << dest_shift);
     }
-    LayoutDetails details = getLayoutDetailsForTransform(quant_type, arch);
+    register_ptr[ii] = transformed_register;
+  }
+}
 
-    PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3, "Shape must be 2-D or 3-D");
+void add_bias_and_interleave_quantized_tensor_inplace(int8_t* tensor,
+                                                      const size_t num_elts,
+                                                      QuantType quant_type) {
+  if (quant_type == QuantType::W8_A16) {
+    add_bias_and_interleave_int8s_inplace(tensor, num_elts);
+  } else if (quant_type == QuantType::W4_A16 ||
+             quant_type == QuantType::W4_AFP8) {
+    // W4_AFP8 uses the same preprocessor as W4_A16 because the FP8 data must
+    // be converted to FP16 before the scales can be applied using CUDA cores.
+    // As a result, we still want permute the data so that it is well aligned
+    // for conversion to FP16.
+    add_bias_and_interleave_int4s_inplace(tensor, num_elts);
+  } else {
+    PADDLE_ENFORCE(false, "Invalid quantization type for interleaving.");
+  }
+}
 
-    size_t num_elts = 1;
-    for (auto const& dim : shape)
-    {
-        num_elts *= dim;
+void interleave_column_major_tensor(int8_t* interleaved_quantized_tensor,
+                                    int8_t const* quantized_tensor,
+                                    std::vector<size_t> const& shape,
+                                    QuantType quant_type,
+                                    LayoutDetails details) {
+  PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3,
+                 "Shape must be 2-D or 3-D");
+  const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
+
+  int const BITS_PER_ELT = get_weight_quant_bits(quant_type);
+  int const elts_in_int32 = 32 / BITS_PER_ELT;
+
+  int const rows_per_tile = details.rows_per_column_tile;
+
+  PADDLE_ENFORCE(!(num_rows % elts_in_int32),
+                 "The number of rows must be a multiple of %d but the number "
+                 "of rows is %ld.",
+                 elts_in_int32,
+                 num_rows);
+
+  uint32_t const* input_byte_ptr =
+      reinterpret_cast<uint32_t const*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
+
+  PADDLE_ENFORCE(!(num_rows % rows_per_tile),
+                 "The number of rows must be a multiple of %d but the number "
+                 "of rows is %ld.",
+                 rows_per_tile,
+                 num_rows);
+
+  int const num_vec_rows = num_rows / elts_in_int32;
+  int const vec_rows_per_tile = rows_per_tile / elts_in_int32;
+  int const interleave = details.columns_interleaved;
+
+  for (int expert = 0; expert < num_experts; ++expert) {
+    const int64_t matrix_offset =
+        expert * int64_t(num_vec_rows) * int64_t(num_cols);
+    for (int read_col = 0; read_col < num_cols; ++read_col) {
+      const int64_t write_col = read_col / interleave;
+      for (int base_vec_row = 0; base_vec_row < num_vec_rows;
+           base_vec_row += vec_rows_per_tile) {
+        for (int vec_read_row = base_vec_row;
+             vec_read_row <
+             std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
+             ++vec_read_row) {
+          const int64_t vec_write_row =
+              interleave * base_vec_row +
+              vec_rows_per_tile * (read_col % interleave) +
+              vec_read_row % vec_rows_per_tile;
+
+          const int64_t read_offset =
+              matrix_offset + int64_t(read_col) * num_vec_rows + vec_read_row;
+          const int64_t write_offset =
+              matrix_offset + int64_t(write_col) * num_vec_rows * interleave +
+              vec_write_row;
+          output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+        }
+      }
     }
+  }
+}
 
-    const size_t num_bytes = num_elts * get_weight_quant_bits(quant_type) / 8;
+void preprocess_weights_for_mixed_gemm(int8_t* preprocessed_quantized_weight,
+                                       int8_t const* row_major_quantized_weight,
+                                       std::vector<size_t> const& shape,
+                                       QuantType quant_type,
+                                       bool force_interleave) {
+  int arch = 89;
+  if (force_interleave && arch == 90) {
+    // Workaround for MOE which doesn't have specialised Hopper kernels yet
+    arch = 80;
+  }
+  LayoutDetails details = getLayoutDetailsForTransform(quant_type, arch);
 
-    std::vector<int8_t> src_buf(num_bytes);
-    std::vector<int8_t> dst_buf(num_bytes);
-    std::copy(row_major_quantized_weight, row_major_quantized_weight + num_bytes, src_buf.begin());
+  PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3,
+                 "Shape must be 2-D or 3-D");
 
-    // Works on row major data, so issue this permutation first.
-    if (details.uses_imma_ldsm)
-    {
-        permute_B_rows_for_mixed_gemm(dst_buf.data(), src_buf.data(), shape, quant_type, arch);
-        src_buf.swap(dst_buf);
-    }
+  size_t num_elts = 1;
+  for (auto const& dim : shape) {
+    num_elts *= dim;
+  }
 
-    if (details.layoutB == LayoutDetails::Layout::COLUMN_MAJOR)
-    {
-        subbyte_transpose(dst_buf.data(), src_buf.data(), shape, quant_type);
-        src_buf.swap(dst_buf);
-    }
+  const size_t num_bytes = num_elts * get_weight_quant_bits(quant_type) / 8;
 
-    if (details.columns_interleaved > 1)
-    {
-        interleave_column_major_tensor(dst_buf.data(), src_buf.data(), shape, quant_type, details);
-        src_buf.swap(dst_buf);
-    }
+  std::vector<int8_t> src_buf(num_bytes);
+  std::vector<int8_t> dst_buf(num_bytes);
+  std::copy(row_major_quantized_weight,
+            row_major_quantized_weight + num_bytes,
+            src_buf.begin());
 
-    if (arch >= 70 && arch < 90)
-    {
-        add_bias_and_interleave_quantized_tensor_inplace(src_buf.data(), num_elts, quant_type);
-    }
-    std::copy(src_buf.begin(), src_buf.end(), preprocessed_quantized_weight);
+  // Works on row major data, so issue this permutation first.
+  if (details.uses_imma_ldsm) {
+    permute_B_rows_for_mixed_gemm(
+        dst_buf.data(), src_buf.data(), shape, quant_type, arch);
+    src_buf.swap(dst_buf);
+  }
+
+  if (details.layoutB == LayoutDetails::Layout::COLUMN_MAJOR) {
+    subbyte_transpose(dst_buf.data(), src_buf.data(), shape, quant_type);
+    src_buf.swap(dst_buf);
+  }
+
+  if (details.columns_interleaved > 1) {
+    interleave_column_major_tensor(
+        dst_buf.data(), src_buf.data(), shape, quant_type, details);
+    src_buf.swap(dst_buf);
+  }
+
+  if (arch >= 70 && arch < 90) {
+    add_bias_and_interleave_quantized_tensor_inplace(
+        src_buf.data(), num_elts, quant_type);
+  }
+  std::copy(src_buf.begin(), src_buf.end(), preprocessed_quantized_weight);
 }
 
 /*
     Arguments:
-      input_weight_ptr - the weight tensor to be quantized. Must be 2-D or 3-D and of type FP16.
+      input_weight_ptr - the weight tensor to be quantized. Must be 2-D or 3-D
+and of type FP16.
 
       quant_type - the type of the output quantization weight.
 
-    This function does symmetric quantization on 2-D or 3-D tensors. It uses the full int range and assumes the
-    zero-point is zero and will automatically construct the scales.
+    This function does symmetric quantization on 2-D or 3-D tensors. It uses the
+full int range and assumes the zero-point is zero and will automatically
+construct the scales.
 
-    It always quantizes the last axis of the tensor. For 3-D tensors, it operates in "batched" mode where the tensor is
-    viewed as a stack of matrices and a scale is produced for each column of every matrix.
+    It always quantizes the last axis of the tensor. For 3-D tensors, it
+operates in "batched" mode where the tensor is viewed as a stack of matrices and
+a scale is produced for each column of every matrix.
 
 Outputs
-    processed_quantized_weight - quantized AND processed weight for GEMM. This MUST be used with the CUTLASS GEMM
-    unprocessed_quantized_weight - quantized but unprocessed weights. Useful for reference checking.
-    scale_ptr - scales for the quantized weight.
+    processed_quantized_weight - quantized AND processed weight for GEMM. This
+MUST be used with the CUTLASS GEMM unprocessed_quantized_weight - quantized but
+unprocessed weights. Useful for reference checking. scale_ptr - scales for the
+quantized weight.
 
-    Note that the returned quantized_weights will be preprocessed in a way to accelerate the mixed type GEMM. The data
-    layout may not make sense if printed.
+    Note that the returned quantized_weights will be preprocessed in a way to
+accelerate the mixed type GEMM. The data layout may not make sense if printed.
 
     Shapes:
       quant_type == int8:
-        If weight is a [m,n] matrix, quantized_weights will have shape [m,n] and scales of shape [n]
-        If weight is a [b,m,n] tensor, unprocessed_quantized_weight will have shape [b,m,n] and scales of shape [b,n]
-      quant_type == int4:
-        If weight is a [m,n] matrix, quantized_weights will have shape [m, ceil(n/2)] and scales of shape [n]
-        If weight is a [b,m,n] tensor, unprocessed_quantized_weight will have shape [b,m, ceil(n/2)] and scales of shape
-          [b,n]
+        If weight is a [m,n] matrix, quantized_weights will have shape [m,n] and
+scales of shape [n] If weight is a [b,m,n] tensor, unprocessed_quantized_weight
+will have shape [b,m,n] and scales of shape [b,n] quant_type == int4: If weight
+is a [m,n] matrix, quantized_weights will have shape [m, ceil(n/2)] and scales
+of shape [n] If weight is a [b,m,n] tensor, unprocessed_quantized_weight will
+have shape [b,m, ceil(n/2)] and scales of shape [b,n]
 
-      The quantized_weight will be of type torch.int8 and have two int4 values packed in a single byte. This is the
-      reason for halving the shape. At the time of writing this code, there was not an elegant way to handle this kind
-      of batched quantization using torch's quantized tensors (to the best of the author's knowledge). Scale tensors
-      must have a dimension of 1, which breaks the semantics we need for batched weights.
+      The quantized_weight will be of type torch.int8 and have two int4 values
+packed in a single byte. This is the reason for halving the shape. At the time
+of writing this code, there was not an elegant way to handle this kind of
+batched quantization using torch's quantized tensors (to the best of the
+author's knowledge). Scale tensors must have a dimension of 1, which breaks the
+semantics we need for batched weights.
   */
 
 template <typename ComputeType, typename WeightType>
-void symmetric_quantize(int8_t* processed_quantized_weight, int8_t* unprocessed_quantized_weight,
-    ComputeType* scale_ptr, WeightType const* input_weight_ptr, std::vector<size_t> const& shape, QuantType quant_type,
-    bool force_interleave)
-{
+void symmetric_quantize(int8_t* processed_quantized_weight,
+                        int8_t* unprocessed_quantized_weight,
+                        ComputeType* scale_ptr,
+                        WeightType const* input_weight_ptr,
+                        std::vector<size_t> const& shape,
+                        QuantType quant_type,
+                        bool force_interleave) {
+  PADDLE_ENFORCE(processed_quantized_weight,
+                 "Processed quantized tensor is NULL");
+  PADDLE_ENFORCE(scale_ptr, "Scale output pointer is NULL");
+  PADDLE_ENFORCE(input_weight_ptr, "Input weight pointer is NULL");
 
-    PADDLE_ENFORCE(processed_quantized_weight, "Processed quantized tensor is NULL");
-    PADDLE_ENFORCE(scale_ptr, "Scale output pointer is NULL");
-    PADDLE_ENFORCE(input_weight_ptr, "Input weight pointer is NULL");
+  PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3,
+                 "Shape must be 2-D or 3-D");
+  const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    PADDLE_ENFORCE(shape.size() == 2 || shape.size() == 3, "Shape must be 2-D or 3-D");
-    const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
-    const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
+  int const bits_in_type = get_weight_quant_bits(quant_type);
+  int const bytes_per_out_col = num_cols * bits_in_type / 8;
 
-    int const bits_in_type = get_weight_quant_bits(quant_type);
-    int const bytes_per_out_col = num_cols * bits_in_type / 8;
+  int const bits_per_weigtht_element = get_weight_quant_bits(quant_type);
 
-    int const bits_per_weigtht_element = get_weight_quant_bits(quant_type);
+  std::vector<int8_t> weight_buf;
+  if (unprocessed_quantized_weight == nullptr) {
+    weight_buf.resize(num_experts * num_rows * num_cols);
+    unprocessed_quantized_weight = weight_buf.data();
+  }
 
-    std::vector<int8_t> weight_buf;
-    if (unprocessed_quantized_weight == nullptr)
-    {
-        weight_buf.resize(num_experts * num_rows * num_cols);
-        unprocessed_quantized_weight = weight_buf.data();
+  int const input_mat_size = num_rows * num_cols;
+  int const quantized_mat_size = num_rows * bytes_per_out_col;
+  float const quant_range_scale = 1.f / float(1 << (bits_in_type - 1));
+
+  std::vector<float> per_col_max(num_cols);
+
+  for (int expert = 0; expert < num_experts; ++expert) {
+    WeightType const* current_weight =
+        input_weight_ptr + expert * input_mat_size;
+    int8_t* current_quantized_weight =
+        unprocessed_quantized_weight + expert * quantized_mat_size;
+
+    // First we find the per column max for this expert weight.
+    for (int jj = 0; jj < num_cols; ++jj) {
+      per_col_max[jj] = 0.f;
     }
 
-    int const input_mat_size = num_rows * num_cols;
-    int const quantized_mat_size = num_rows * bytes_per_out_col;
-    float const quant_range_scale = 1.f / float(1 << (bits_in_type - 1));
-
-    std::vector<float> per_col_max(num_cols);
-
-    for (int expert = 0; expert < num_experts; ++expert)
-    {
-        WeightType const* current_weight = input_weight_ptr + expert * input_mat_size;
-        int8_t* current_quantized_weight = unprocessed_quantized_weight + expert * quantized_mat_size;
-
-        // First we find the per column max for this expert weight.
-        for (int jj = 0; jj < num_cols; ++jj)
-        {
-            per_col_max[jj] = 0.f;
-        }
-
-        for (int ii = 0; ii < num_rows; ++ii)
-        {
-            WeightType const* current_weight_row = current_weight + ii * num_cols;
-            for (int jj = 0; jj < num_cols; ++jj)
-            {
-                per_col_max[jj] = std::max(per_col_max[jj], std::abs(float(current_weight_row[jj])));
-            }
-        }
-
-        // Then, we construct the scales
-        ComputeType* current_scales = scale_ptr + expert * num_cols;
-        for (int jj = 0; jj < num_cols; ++jj)
-        {
-            per_col_max[jj] *= quant_range_scale;
-            current_scales[jj] = ComputeType(per_col_max[jj]);
-        }
-
-        // Finally, construct the weights.
-        for (int ii = 0; ii < num_rows; ++ii)
-        {
-            int8_t* current_quantized_weight_row = current_quantized_weight + ii * bytes_per_out_col;
-            WeightType const* current_weight_row = current_weight + ii * num_cols;
-            for (int jj = 0; jj < bytes_per_out_col; ++jj)
-            {
-
-                if (bits_per_weigtht_element == 8)
-                {
-                    float const col_scale = per_col_max[jj];
-                    float const weight_elt = float(current_weight_row[jj]);
-                    float const scaled_weight = (col_scale != 0.0f) ? round(weight_elt / col_scale) : 0.0f;
-                    const int8_t clipped_weight = int8_t(std::max(-128.f, std::min(127.f, scaled_weight)));
-                    current_quantized_weight_row[jj] = clipped_weight;
-                }
-                else if (bits_per_weigtht_element == 4)
-                {
-
-                    // We will pack two int4 elements per iteration of the inner loop.
-                    int8_t packed_int4s = 0;
-                    for (int packed_idx = 0; packed_idx < 2; ++packed_idx)
-                    {
-                        int const input_idx = 2 * jj + packed_idx;
-                        if (input_idx < num_cols)
-                        {
-                            float const col_scale = per_col_max[input_idx];
-                            float const weight_elt = float(current_weight_row[input_idx]);
-                            float const scaled_weight = (col_scale != 0.0f) ? round(weight_elt / col_scale) : 0.0f;
-                            int int_weight = int(scaled_weight);
-                            const int8_t clipped_weight = std::max(-8, std::min(7, int_weight));
-
-                            // Kill the sign extension bits (hence 0x0F mask) then shift to upper bits
-                            // if packing the second int4 and or the bits into the final result.
-                            packed_int4s |= ((clipped_weight & 0x0F) << (4 * packed_idx));
-                        }
-                    }
-                    current_quantized_weight_row[jj] = packed_int4s;
-                }
-                else
-                {
-                    PADDLE_ENFORCE(false, "Unsupported quantization type");
-                }
-            }
-        }
+    for (int ii = 0; ii < num_rows; ++ii) {
+      WeightType const* current_weight_row = current_weight + ii * num_cols;
+      for (int jj = 0; jj < num_cols; ++jj) {
+        per_col_max[jj] =
+            std::max(per_col_max[jj], std::abs(float(current_weight_row[jj])));
+      }
     }
 
-    preprocess_weights_for_mixed_gemm(
-        processed_quantized_weight, unprocessed_quantized_weight, shape, quant_type, force_interleave);
+    // Then, we construct the scales
+    ComputeType* current_scales = scale_ptr + expert * num_cols;
+    for (int jj = 0; jj < num_cols; ++jj) {
+      per_col_max[jj] *= quant_range_scale;
+      current_scales[jj] = ComputeType(per_col_max[jj]);
+    }
+
+    // Finally, construct the weights.
+    for (int ii = 0; ii < num_rows; ++ii) {
+      int8_t* current_quantized_weight_row =
+          current_quantized_weight + ii * bytes_per_out_col;
+      WeightType const* current_weight_row = current_weight + ii * num_cols;
+      for (int jj = 0; jj < bytes_per_out_col; ++jj) {
+        if (bits_per_weigtht_element == 8) {
+          float const col_scale = per_col_max[jj];
+          float const weight_elt = float(current_weight_row[jj]);
+          float const scaled_weight =
+              (col_scale != 0.0f) ? round(weight_elt / col_scale) : 0.0f;
+          const int8_t clipped_weight =
+              int8_t(std::max(-128.f, std::min(127.f, scaled_weight)));
+          current_quantized_weight_row[jj] = clipped_weight;
+        } else if (bits_per_weigtht_element == 4) {
+          // We will pack two int4 elements per iteration of the inner loop.
+          int8_t packed_int4s = 0;
+          for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
+            int const input_idx = 2 * jj + packed_idx;
+            if (input_idx < num_cols) {
+              float const col_scale = per_col_max[input_idx];
+              float const weight_elt = float(current_weight_row[input_idx]);
+              float const scaled_weight =
+                  (col_scale != 0.0f) ? round(weight_elt / col_scale) : 0.0f;
+              int int_weight = int(scaled_weight);
+              const int8_t clipped_weight =
+                  std::max(-8, std::min(7, int_weight));
+
+              // Kill the sign extension bits (hence 0x0F mask) then shift to
+              // upper bits if packing the second int4 and or the bits into the
+              // final result.
+              packed_int4s |= ((clipped_weight & 0x0F) << (4 * packed_idx));
+            }
+          }
+          current_quantized_weight_row[jj] = packed_int4s;
+        } else {
+          PADDLE_ENFORCE(false, "Unsupported quantization type");
+        }
+      }
+    }
+  }
+
+  preprocess_weights_for_mixed_gemm(processed_quantized_weight,
+                                    unprocessed_quantized_weight,
+                                    shape,
+                                    quant_type,
+                                    force_interleave);
 }
 
-template void symmetric_quantize<half, float>(
-    int8_t*, int8_t*, half*, float const*, std::vector<size_t> const&, QuantType, bool);
+template void symmetric_quantize<half, float>(int8_t*,
+                                              int8_t*,
+                                              half*,
+                                              float const*,
+                                              std::vector<size_t> const&,
+                                              QuantType,
+                                              bool);
 
-template void symmetric_quantize<half, half>(
-    int8_t*, int8_t*, half*, half const*, std::vector<size_t> const&, QuantType, bool);
+template void symmetric_quantize<half, half>(int8_t*,
+                                             int8_t*,
+                                             half*,
+                                             half const*,
+                                             std::vector<size_t> const&,
+                                             QuantType,
+                                             bool);
 
 #ifdef ENABLE_BF16
 template void symmetric_quantize<__nv_bfloat16, __nv_bfloat16>(
-    int8_t*, int8_t*, __nv_bfloat16*, __nv_bfloat16 const*, std::vector<size_t> const&, QuantType, bool);
+    int8_t*,
+    int8_t*,
+    __nv_bfloat16*,
+    __nv_bfloat16 const*,
+    std::vector<size_t> const&,
+    QuantType,
+    bool);
 
 template void symmetric_quantize<__nv_bfloat16, float>(
-    int8_t*, int8_t*, __nv_bfloat16*, float const*, std::vector<size_t> const&, QuantType, bool);
+    int8_t*,
+    int8_t*,
+    __nv_bfloat16*,
+    float const*,
+    std::vector<size_t> const&,
+    QuantType,
+    bool);
 #endif
 
 template <typename ComputeType, typename WeightType>
-void symmetric_quantize(int8_t* processed_quantized_weight, ComputeType* scale_ptr, WeightType const* input_weight_ptr,
-    std::vector<size_t> const& shape, QuantType quant_type, bool force_interleave)
-{
-    symmetric_quantize(
-        processed_quantized_weight, nullptr, scale_ptr, input_weight_ptr, shape, quant_type, force_interleave);
+void symmetric_quantize(int8_t* processed_quantized_weight,
+                        ComputeType* scale_ptr,
+                        WeightType const* input_weight_ptr,
+                        std::vector<size_t> const& shape,
+                        QuantType quant_type,
+                        bool force_interleave) {
+  symmetric_quantize(processed_quantized_weight,
+                     nullptr,
+                     scale_ptr,
+                     input_weight_ptr,
+                     shape,
+                     quant_type,
+                     force_interleave);
 }
 
 template void symmetric_quantize<float, float>(
@@ -773,21 +802,42 @@ template void symmetric_quantize<float, float>(
 template void symmetric_quantize<half, float>(
     int8_t*, half*, float const*, std::vector<size_t> const&, QuantType, bool);
 
-template void symmetric_quantize<half, half>(int8_t*, half*, half const*, std::vector<size_t> const&, QuantType, bool);
+template void symmetric_quantize<half, half>(
+    int8_t*, half*, half const*, std::vector<size_t> const&, QuantType, bool);
 
 #ifdef ENABLE_BF16
 template void symmetric_quantize<__nv_bfloat16, __nv_bfloat16>(
-    int8_t*, __nv_bfloat16*, __nv_bfloat16 const*, std::vector<size_t> const&, QuantType, bool);
+    int8_t*,
+    __nv_bfloat16*,
+    __nv_bfloat16 const*,
+    std::vector<size_t> const&,
+    QuantType,
+    bool);
 
 template void symmetric_quantize<__nv_bfloat16, half>(
-    int8_t*, __nv_bfloat16*, half const*, std::vector<size_t> const&, QuantType, bool);
+    int8_t*,
+    __nv_bfloat16*,
+    half const*,
+    std::vector<size_t> const&,
+    QuantType,
+    bool);
 
 template void symmetric_quantize<half, __nv_bfloat16>(
-    int8_t*, half*, __nv_bfloat16 const*, std::vector<size_t> const&, QuantType, bool);
+    int8_t*,
+    half*,
+    __nv_bfloat16 const*,
+    std::vector<size_t> const&,
+    QuantType,
+    bool);
 
 template void symmetric_quantize<__nv_bfloat16, float>(
-    int8_t*, __nv_bfloat16*, float const*, std::vector<size_t> const&, QuantType, bool);
+    int8_t*,
+    __nv_bfloat16*,
+    float const*,
+    std::vector<size_t> const&,
+    QuantType,
+    bool);
 #endif
 
-} // namespace cutlass_kernels
-} // namespace kernels
+}  // namespace cutlass_kernels
+}  // namespace kernels
diff --git a/custom_ops/gpu_ops/cutlass_kernels/cutlass_preprocessors.h b/custom_ops/gpu_ops/cutlass_kernels/cutlass_preprocessors.h
index 8d025c1289..292d7d8cbe 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/cutlass_preprocessors.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/cutlass_preprocessors.h
@@ -20,52 +20,67 @@
 #include <stdint.h>
 #include <vector>
 
-namespace kernels
-{
-namespace cutlass_kernels
-{
+namespace kernels {
+namespace cutlass_kernels {
 
-enum class QuantType
-{
-    W8_A16,
-    W4_A16,
-    W4_AFP8
-};
+enum class QuantType { W8_A16, W4_A16, W4_AFP8 };
 
-constexpr int get_weight_quant_bits(QuantType quant_type)
-{
-    switch (quant_type)
-    {
-    case QuantType::W8_A16: return 8;
-    case QuantType::W4_A16: return 4;
-    case QuantType::W4_AFP8: return 4;
-    default: PADDLE_THROW("Invalid quant_type"); return -1;
-    }
+constexpr int get_weight_quant_bits(QuantType quant_type) {
+  switch (quant_type) {
+    case QuantType::W8_A16:
+      return 8;
+    case QuantType::W4_A16:
+      return 4;
+    case QuantType::W4_AFP8:
+      return 4;
+    default:
+      PADDLE_THROW("Invalid quant_type");
+      return -1;
+  }
 }
 
 // Shapes here can be 2 or 3D. 2-D shapes are [num_rows, num_cols]
 // 3-D shapes are [num_experts, num_rows, num_cols]
-void permute_B_rows_for_mixed_gemm(int8_t* permuted_quantized_tensor, int8_t const* quantized_tensor,
-    std::vector<size_t> const& shape, QuantType quant_type, const int64_t arch_version);
+void permute_B_rows_for_mixed_gemm(int8_t* permuted_quantized_tensor,
+                                   int8_t const* quantized_tensor,
+                                   std::vector<size_t> const& shape,
+                                   QuantType quant_type,
+                                   const int64_t arch_version);
 
-void subbyte_transpose(int8_t* transposed_quantized_tensor, int8_t const* quantized_tensor,
-    std::vector<size_t> const& shape, QuantType quant_type);
+void subbyte_transpose(int8_t* transposed_quantized_tensor,
+                       int8_t const* quantized_tensor,
+                       std::vector<size_t> const& shape,
+                       QuantType quant_type);
 
-void add_bias_and_interleave_quantized_tensor_inplace(int8_t* tensor, const size_t num_elts, QuantType quant_type);
+void add_bias_and_interleave_quantized_tensor_inplace(int8_t* tensor,
+                                                      const size_t num_elts,
+                                                      QuantType quant_type);
 
-void preprocess_weights_for_mixed_gemm(int8_t* preprocessed_quantized_weight, int8_t const* row_major_quantized_weight,
-    std::vector<size_t> const& shape, QuantType quant_type, bool force_interleave = false);
+void preprocess_weights_for_mixed_gemm(int8_t* preprocessed_quantized_weight,
+                                       int8_t const* row_major_quantized_weight,
+                                       std::vector<size_t> const& shape,
+                                       QuantType quant_type,
+                                       bool force_interleave = false);
 
 template <typename ComputeType, typename WeightType>
-void symmetric_quantize(int8_t* processed_quantized_weight, ComputeType* scale_ptr, WeightType const* input_weight_ptr,
-    std::vector<size_t> const& shape, QuantType quant_type, bool force_interleave);
+void symmetric_quantize(int8_t* processed_quantized_weight,
+                        ComputeType* scale_ptr,
+                        WeightType const* input_weight_ptr,
+                        std::vector<size_t> const& shape,
+                        QuantType quant_type,
+                        bool force_interleave);
 
-// This is exposed so that we can write tests that use the processed weights for CUTLASS but the unprocessed weight
-// to implement a simple reference implementation.
+// This is exposed so that we can write tests that use the processed weights for
+// CUTLASS but the unprocessed weight to implement a simple reference
+// implementation.
 template <typename ComputeType, typename WeightType>
-void symmetric_quantize(int8_t* processed_quantized_weight, int8_t* unprocessed_quantized_weight,
-    ComputeType* scale_ptr, WeightType const* input_weight_ptr, std::vector<size_t> const& shape, QuantType quant_type,
-    bool force_interleave);
+void symmetric_quantize(int8_t* processed_quantized_weight,
+                        int8_t* unprocessed_quantized_weight,
+                        ComputeType* scale_ptr,
+                        WeightType const* input_weight_ptr,
+                        std::vector<size_t> const& shape,
+                        QuantType quant_type,
+                        bool force_interleave);
 
-} // namespace cutlass_kernels
-} // namespace kernels
+}  // namespace cutlass_kernels
+}  // namespace kernels
diff --git a/custom_ops/gpu_ops/cutlass_kernels/cutlass_type_conversion.h b/custom_ops/gpu_ops/cutlass_kernels/cutlass_type_conversion.h
index cf344772a6..a10f49cc77 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/cutlass_type_conversion.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/cutlass_type_conversion.h
@@ -24,45 +24,38 @@
 #include "cutlass/float8.h"
 #include "cutlass/half.h"
 
-namespace kernels
-{
-namespace cutlass_kernels
-{
+namespace kernels {
+namespace cutlass_kernels {
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 // Cuda to Cutlass
 
 template <typename T>
-struct CudaToCutlassTypeAdapter
-{
-    using type = T;
+struct CudaToCutlassTypeAdapter {
+  using type = T;
 };
 
 template <>
-struct CudaToCutlassTypeAdapter<half>
-{
-    using type = cutlass::half_t;
+struct CudaToCutlassTypeAdapter<half> {
+  using type = cutlass::half_t;
 };
 
 #if defined(ENABLE_BF16)
 template <>
-struct CudaToCutlassTypeAdapter<__nv_bfloat16>
-{
-    using type = cutlass::bfloat16_t;
+struct CudaToCutlassTypeAdapter<__nv_bfloat16> {
+  using type = cutlass::bfloat16_t;
 };
 #endif
 
 #if defined(ENABLE_FP8)
 template <>
-struct CudaToCutlassTypeAdapter<__nv_fp8_e4m3>
-{
-    using type = cutlass::float_e4m3_t;
+struct CudaToCutlassTypeAdapter<__nv_fp8_e4m3> {
+  using type = cutlass::float_e4m3_t;
 };
 
 template <>
-struct CudaToCutlassTypeAdapter<__nv_fp8_e5m2>
-{
-    using type = cutlass::float_e5m2_t;
+struct CudaToCutlassTypeAdapter<__nv_fp8_e5m2> {
+  using type = cutlass::float_e5m2_t;
 };
 #endif
 
@@ -70,40 +63,35 @@ struct CudaToCutlassTypeAdapter<__nv_fp8_e5m2>
 // Cutlass to Cuda
 
 template <typename T>
-struct CutlassToCudaTypeAdapter
-{
-    using type = T;
+struct CutlassToCudaTypeAdapter {
+  using type = T;
 };
 
 template <>
-struct CutlassToCudaTypeAdapter<cutlass::half_t>
-{
-    using type = half;
+struct CutlassToCudaTypeAdapter<cutlass::half_t> {
+  using type = half;
 };
 
 #if defined(ENABLE_BF16)
 template <>
-struct CutlassToCudaTypeAdapter<cutlass::bfloat16_t>
-{
-    using type = __nv_bfloat16;
+struct CutlassToCudaTypeAdapter<cutlass::bfloat16_t> {
+  using type = __nv_bfloat16;
 };
 #endif
 
 #if defined(ENABLE_FP8)
 template <>
-struct CutlassToCudaTypeAdapter<cutlass::float_e4m3_t>
-{
-    using type = __nv_fp8_e4m3;
+struct CutlassToCudaTypeAdapter<cutlass::float_e4m3_t> {
+  using type = __nv_fp8_e4m3;
 };
 
 template <>
-struct CutlassToCudaTypeAdapter<cutlass::float_e5m2_t>
-{
-    using type = __nv_fp8_e5m2;
+struct CutlassToCudaTypeAdapter<cutlass::float_e5m2_t> {
+  using type = __nv_fp8_e5m2;
 };
 #endif
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass_kernels
-} // namespace kernels
+}  // namespace cutlass_kernels
+}  // namespace kernels
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_gelu_and_mul.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_gelu_and_mul.h
index 9dfddf83fc..a4f421a01f 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_gelu_and_mul.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_gelu_and_mul.h
@@ -67,87 +67,87 @@ template <
         ElementOutput_,  ///< Data type used to compute linear combination
     FloatRoundStyle Round = FloatRoundStyle::round_to_nearest>
 class LeftGELUAndMul {
-    public:
-    using ElementOutput = ElementOutput_;
-    using ElementAccumulator = ElementAccumulator_;
-    using ElementCompute = ElementCompute_;
+ public:
+  using ElementOutput = ElementOutput_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
 
-    static int const kCount = Count;
-    using FragmentOutput = Array<ElementOutput, kCount>;
-    using FragmentAccumulator = Array<ElementAccumulator, kCount>;
-    using ComputeFragment = Array<ElementCompute, kCount>;
+  static int const kCount = Count;
+  using FragmentOutput = Array<ElementOutput, kCount>;
+  using FragmentAccumulator = Array<ElementAccumulator, kCount>;
+  using ComputeFragment = Array<ElementCompute, kCount>;
 
-    static FloatRoundStyle const kRound = Round;
+  static FloatRoundStyle const kRound = Round;
 
-    struct Params {
-        ElementCompute alpha;
-
-      CUTLASS_HOST_DEVICE
-      Params() : alpha(ElementCompute(1)) {}
-
-      CUTLASS_HOST_DEVICE
-      Params(ElementCompute alpha) : alpha(alpha) {}  // NOLINT
-    };
-
-    private:
-    //
-    // Data members
-    //
-
-    ElementCompute alpha_;
-    ElementCompute beta_;
-
-    public:
-    /// Constructs the function object, possibly loading from pointers in host
-    /// memory
-    CUTLASS_HOST_DEVICE
-    LeftGELUAndMul(Params const &params) { alpha_ = params.alpha; }  // NOLINT
-
-    /// Returns true if source is needed
-    CUTLASS_HOST_DEVICE
-    bool is_source_needed() const { return true; }
-
-    /// Functionally required for serial reduction in the epilogue
-    CUTLASS_HOST_DEVICE
-    void set_k_partition(int k_partition, int k_partition_count) {
-      assert(false);
-    }
-
-    /// Computes linear scaling: D = alpha * accumulator + beta * source
-    CUTLASS_HOST_DEVICE
-    FragmentOutput operator()(FragmentAccumulator const &lhs,
-                              FragmentAccumulator const &rhs) const {
-        // Convert source to internal compute numeric type
-        NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
-            accumulator_to_compute;
-
-        // Convert to destination numeric type
-        NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
-            compute_to_output;
-
-        ComputeFragment converted_lhs = accumulator_to_compute(lhs);
-        ComputeFragment converted_rhs = accumulator_to_compute(rhs);
-
-        cutlass::epilogue::thread::GELU_taylor<ComputeFragment> gelu;
-        cutlass::multiplies<ComputeFragment> mul;
-        auto gelu_lhs = gelu(converted_lhs);
-        // return compute_to_output(mul(gelu_lhs, converted_rhs));
-        auto tmp = mul(gelu_lhs, converted_rhs);
-        return compute_to_output(mul(alpha_, tmp));
-    }
+  struct Params {
+    ElementCompute alpha;
 
     CUTLASS_HOST_DEVICE
-    ElementOutput operator()(ElementAccumulator const &lhs,
-                            ElementAccumulator const &rhs) const {
-        ElementCompute convert_lhs(lhs);
-        ElementCompute convert_rhs(rhs);
-        cutlass::epilogue::thread::GELU_taylor<ElementCompute> gelu;
-        cutlass::multiplies<ElementCompute> mul;
-        auto gelu_lhs = gelu(convert_lhs);
-        // return ElementOutput(mul(gelu_lhs, convert_rhs));
-        auto tmp = mul(gelu_lhs, convert_rhs);
-        return compute_to_output(mul(alpha_, tmp));
-    }
+    Params() : alpha(ElementCompute(1)) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(ElementCompute alpha) : alpha(alpha) {}  // NOLINT
+  };
+
+ private:
+  //
+  // Data members
+  //
+
+  ElementCompute alpha_;
+  ElementCompute beta_;
+
+ public:
+  /// Constructs the function object, possibly loading from pointers in host
+  /// memory
+  CUTLASS_HOST_DEVICE
+  LeftGELUAndMul(Params const &params) { alpha_ = params.alpha; }  // NOLINT
+
+  /// Returns true if source is needed
+  CUTLASS_HOST_DEVICE
+  bool is_source_needed() const { return true; }
+
+  /// Functionally required for serial reduction in the epilogue
+  CUTLASS_HOST_DEVICE
+  void set_k_partition(int k_partition, int k_partition_count) {
+    assert(false);
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator + beta * source
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &lhs,
+                            FragmentAccumulator const &rhs) const {
+    // Convert source to internal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
+        accumulator_to_compute;
+
+    // Convert to destination numeric type
+    NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
+        compute_to_output;
+
+    ComputeFragment converted_lhs = accumulator_to_compute(lhs);
+    ComputeFragment converted_rhs = accumulator_to_compute(rhs);
+
+    cutlass::epilogue::thread::GELU_taylor<ComputeFragment> gelu;
+    cutlass::multiplies<ComputeFragment> mul;
+    auto gelu_lhs = gelu(converted_lhs);
+    // return compute_to_output(mul(gelu_lhs, converted_rhs));
+    auto tmp = mul(gelu_lhs, converted_rhs);
+    return compute_to_output(mul(alpha_, tmp));
+  }
+
+  CUTLASS_HOST_DEVICE
+  ElementOutput operator()(ElementAccumulator const &lhs,
+                           ElementAccumulator const &rhs) const {
+    ElementCompute convert_lhs(lhs);
+    ElementCompute convert_rhs(rhs);
+    cutlass::epilogue::thread::GELU_taylor<ElementCompute> gelu;
+    cutlass::multiplies<ElementCompute> mul;
+    auto gelu_lhs = gelu(convert_lhs);
+    // return ElementOutput(mul(gelu_lhs, convert_rhs));
+    auto tmp = mul(gelu_lhs, convert_rhs);
+    return compute_to_output(mul(alpha_, tmp));
+  }
 };
 
 }  // namespace thread
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_silu_and_mul.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_silu_and_mul.h
index 7a433bccd6..51da87e898 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_silu_and_mul.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_silu_and_mul.h
@@ -67,87 +67,87 @@ template <
         ElementOutput_,  ///< Data type used to compute linear combination
     FloatRoundStyle Round = FloatRoundStyle::round_to_nearest>
 class LeftSiLUAndMul {
-    public:
-    using ElementOutput = ElementOutput_;
-    using ElementAccumulator = ElementAccumulator_;
-    using ElementCompute = ElementCompute_;
+ public:
+  using ElementOutput = ElementOutput_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
 
-    static int const kCount = Count;
-    using FragmentOutput = Array<ElementOutput, kCount>;
-    using FragmentAccumulator = Array<ElementAccumulator, kCount>;
-    using ComputeFragment = Array<ElementCompute, kCount>;
+  static int const kCount = Count;
+  using FragmentOutput = Array<ElementOutput, kCount>;
+  using FragmentAccumulator = Array<ElementAccumulator, kCount>;
+  using ComputeFragment = Array<ElementCompute, kCount>;
 
-    static FloatRoundStyle const kRound = Round;
+  static FloatRoundStyle const kRound = Round;
 
-    struct Params {
-        ElementCompute alpha;
-
-        CUTLASS_HOST_DEVICE
-        Params() : alpha(ElementCompute(1)) {}
-
-        CUTLASS_HOST_DEVICE
-        Params(ElementCompute alpha) : alpha(alpha) {}  // NOLINT
-    };
-
-    private:
-    //
-    // Data members
-    //
-
-    ElementCompute alpha_;
-    ElementCompute beta_;
-
-    public:
-    /// Constructs the function object, possibly loading from pointers in host
-    /// memory
-    CUTLASS_HOST_DEVICE
-    LeftSiLUAndMul(Params const &params) { alpha_ = params.alpha; }  // NOLINT
-
-    /// Returns true if source is needed
-    CUTLASS_HOST_DEVICE
-    bool is_source_needed() const { return true; }
-
-    /// Functionally required for serial reduction in the epilogue
-    CUTLASS_HOST_DEVICE
-    void set_k_partition(int k_partition, int k_partition_count) {
-        assert(false);
-    }
-
-    /// Computes linear scaling: D = alpha * accumulator + beta * source
-    CUTLASS_HOST_DEVICE
-    FragmentOutput operator()(FragmentAccumulator const &lhs,
-                              FragmentAccumulator const &rhs) const {
-        // Convert source to internal compute numeric type
-        NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
-            accumulator_to_compute;
-
-        // Convert to destination numeric type
-        NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
-            compute_to_output;
-
-        ComputeFragment converted_lhs = accumulator_to_compute(lhs);
-        ComputeFragment converted_rhs = accumulator_to_compute(rhs);
-
-        cutlass::epilogue::thread::SiLu<ComputeFragment> silu;
-        cutlass::multiplies<ComputeFragment> mul;
-        auto silu_lhs = silu(converted_lhs);
-        // return compute_to_output(mul(silu_lhs, converted_rhs));
-        auto tmp = mul(silu_lhs, converted_rhs);
-        return compute_to_output(mul(alpha_, tmp));
-    }
+  struct Params {
+    ElementCompute alpha;
 
     CUTLASS_HOST_DEVICE
-    ElementOutput operator()(ElementAccumulator const &lhs,
-                            ElementAccumulator const &rhs) const {
-        ElementCompute convert_lhs(lhs);
-        ElementCompute convert_rhs(rhs);
-        cutlass::epilogue::thread::SiLu<ElementCompute> silu;
-        cutlass::multiplies<ElementCompute> mul;
-        auto silu_lhs = silu(convert_lhs);
-        // return ElementOutput(mul(silu_lhs, convert_rhs));
-        auto tmp = mul(silu_lhs, convert_rhs);
-        return ElementOutput(mul(alpha_, tmp));
-    }
+    Params() : alpha(ElementCompute(1)) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(ElementCompute alpha) : alpha(alpha) {}  // NOLINT
+  };
+
+ private:
+  //
+  // Data members
+  //
+
+  ElementCompute alpha_;
+  ElementCompute beta_;
+
+ public:
+  /// Constructs the function object, possibly loading from pointers in host
+  /// memory
+  CUTLASS_HOST_DEVICE
+  LeftSiLUAndMul(Params const &params) { alpha_ = params.alpha; }  // NOLINT
+
+  /// Returns true if source is needed
+  CUTLASS_HOST_DEVICE
+  bool is_source_needed() const { return true; }
+
+  /// Functionally required for serial reduction in the epilogue
+  CUTLASS_HOST_DEVICE
+  void set_k_partition(int k_partition, int k_partition_count) {
+    assert(false);
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator + beta * source
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &lhs,
+                            FragmentAccumulator const &rhs) const {
+    // Convert source to internal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
+        accumulator_to_compute;
+
+    // Convert to destination numeric type
+    NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
+        compute_to_output;
+
+    ComputeFragment converted_lhs = accumulator_to_compute(lhs);
+    ComputeFragment converted_rhs = accumulator_to_compute(rhs);
+
+    cutlass::epilogue::thread::SiLu<ComputeFragment> silu;
+    cutlass::multiplies<ComputeFragment> mul;
+    auto silu_lhs = silu(converted_lhs);
+    // return compute_to_output(mul(silu_lhs, converted_rhs));
+    auto tmp = mul(silu_lhs, converted_rhs);
+    return compute_to_output(mul(alpha_, tmp));
+  }
+
+  CUTLASS_HOST_DEVICE
+  ElementOutput operator()(ElementAccumulator const &lhs,
+                           ElementAccumulator const &rhs) const {
+    ElementCompute convert_lhs(lhs);
+    ElementCompute convert_rhs(rhs);
+    cutlass::epilogue::thread::SiLu<ElementCompute> silu;
+    cutlass::multiplies<ElementCompute> mul;
+    auto silu_lhs = silu(convert_lhs);
+    // return ElementOutput(mul(silu_lhs, convert_rhs));
+    auto tmp = mul(silu_lhs, convert_rhs);
+    return ElementOutput(mul(alpha_, tmp));
+  }
 };
 
 }  // namespace thread
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/threadblock/dual_epilogue.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/threadblock/dual_epilogue.h
index 7d679c341d..2be6408013 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/threadblock/dual_epilogue.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/dual_gemm/threadblock/dual_epilogue.h
@@ -98,357 +98,357 @@ template <typename Shape_,  // Shape of threadblock tile (concept: GemmShape)
                                   // op is large
           (!IsEpilogueFunctorHeavy<OutputOp0_>::value)>
 class DualEpilogue {
-    public:
-    using Base = EpilogueBase<Shape_,
-                              typename WarpMmaOperator_::Shape,
-                              PartitionsK,
-                              AccumulatorFragmentIterator_,
-                              WarpTileIterator_,
-                              Padding_,
-                              FragmentsPerPartition>;
+ public:
+  using Base = EpilogueBase<Shape_,
+                            typename WarpMmaOperator_::Shape,
+                            PartitionsK,
+                            AccumulatorFragmentIterator_,
+                            WarpTileIterator_,
+                            Padding_,
+                            FragmentsPerPartition>;
 
-    using Shape = Shape_;
-    using WarpMmaOperator = WarpMmaOperator_;
-    static int const kPartitionsK = PartitionsK;
-    static bool constexpr kStoreD0 = StoreD0;
-    static bool constexpr kStoreD1 = StoreD1;
-    using OutputTileIterator = OutputTileIterator_;
-    using OutputTileIterator2 = OutputTileIterator2_;
-    using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
-    using WarpTileIterator = WarpTileIterator_;
-    using SharedLoadIterator = SharedLoadIterator_;
-    using OutputOp0 = OutputOp0_;
-    using OutputOp1 = OutputOp1_;
-    using OutputOp2 = OutputOp2_;
-    using Padding = Padding_;
+  using Shape = Shape_;
+  using WarpMmaOperator = WarpMmaOperator_;
+  static int const kPartitionsK = PartitionsK;
+  static bool constexpr kStoreD0 = StoreD0;
+  static bool constexpr kStoreD1 = StoreD1;
+  using OutputTileIterator = OutputTileIterator_;
+  using OutputTileIterator2 = OutputTileIterator2_;
+  using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
+  using WarpTileIterator = WarpTileIterator_;
+  using SharedLoadIterator = SharedLoadIterator_;
+  using OutputOp0 = OutputOp0_;
+  using OutputOp1 = OutputOp1_;
+  using OutputOp2 = OutputOp2_;
+  using Padding = Padding_;
 
-    using Layout = layout::RowMajor;
-    using LongIndex = typename Layout::LongIndex;
+  using Layout = layout::RowMajor;
+  using LongIndex = typename Layout::LongIndex;
 
-    // The complete warp-level accumulator tile
-    using AccumulatorTile = typename Base::AccumulatorTile;
+  // The complete warp-level accumulator tile
+  using AccumulatorTile = typename Base::AccumulatorTile;
 
-    // Accumulator element
-    using ElementAccumulator = typename WarpTileIterator::Element;
+  // Accumulator element
+  using ElementAccumulator = typename WarpTileIterator::Element;
 
-    // Output element
-    using ElementOutput = typename OutputTileIterator::Element;
+  // Output element
+  using ElementOutput = typename OutputTileIterator::Element;
 
-    // Output access size
-    static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
+  // Output access size
+  static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
 
-    // Tensor reference to destination tensor
-    using TensorRef = typename OutputTileIterator::TensorRef;
+  // Tensor reference to destination tensor
+  using TensorRef = typename OutputTileIterator::TensorRef;
 
-    // Tensor reference to sync tensor
-    using SyncTensorRef =
-        typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
+  // Tensor reference to sync tensor
+  using SyncTensorRef =
+      typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
 
-    // Const tensor reference to source tensor
-    using ConstTensorRef = typename OutputTileIterator::ConstTensorRef;
+  // Const tensor reference to source tensor
+  using ConstTensorRef = typename OutputTileIterator::ConstTensorRef;
 
-    // Array type used to output
-    using OutputAccessType = Array<typename OutputTileIterator::Element,
-                                  OutputTileIterator::kElementsPerAccess>;
+  // Array type used to output
+  using OutputAccessType = Array<typename OutputTileIterator::Element,
+                                 OutputTileIterator::kElementsPerAccess>;
 
-    // Array type used to output
-    using OutputAccessType2 = Array<typename OutputTileIterator2::Element,
-                                    OutputTileIterator2::kElementsPerAccess>;
+  // Array type used to output
+  using OutputAccessType2 = Array<typename OutputTileIterator2::Element,
+                                  OutputTileIterator2::kElementsPerAccess>;
 
-    // Array type used by output functor
-    using AccumulatorAccessType = Array<typename WarpTileIterator::Element,
-                                        OutputTileIterator::kElementsPerAccess>;
+  // Array type used by output functor
+  using AccumulatorAccessType = Array<typename WarpTileIterator::Element,
+                                      OutputTileIterator::kElementsPerAccess>;
 
-    // Number of warps
-    using WarpCount = typename Base::WarpCount;
+  // Number of warps
+  using WarpCount = typename Base::WarpCount;
 
-    struct SharedStorage {
-        using Element = typename WarpTileIterator::Element;
+  struct SharedStorage {
+    using Element = typename WarpTileIterator::Element;
 
-        // Tensor reference to shared memory allocation
-        using TensorRef = typename WarpTileIterator::TensorRef;
+    // Tensor reference to shared memory allocation
+    using TensorRef = typename WarpTileIterator::TensorRef;
 
-        // Logical shape of the shared memory tile written to by all warps.
-        using Shape = typename Base::Shape;
+    // Logical shape of the shared memory tile written to by all warps.
+    using Shape = typename Base::Shape;
 
-        // Shape of the shared memory allocation for the epilogue
-        using StorageShape = typename Base::SharedStorage::StorageShape;
+    // Shape of the shared memory allocation for the epilogue
+    using StorageShape = typename Base::SharedStorage::StorageShape;
 
-        //
-        // Data members
-        //
+    //
+    // Data members
+    //
 
-        AlignedBuffer<Element, StorageShape::kCount> storage[2];
+    AlignedBuffer<Element, StorageShape::kCount> storage[2];
 
-        //
-        // Methods
-        //
+    //
+    // Methods
+    //
 
-        // Returns a tensor reference to the shared memory buffer
-        CUTLASS_DEVICE
-        TensorRef reference(int i) {
-            return TensorRef(
-                storage[i].data(),
-                Layout::packed({StorageShape::kRow, StorageShape::kColumn}));
-        }
-    };
-
-    static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1
-                                          ? Base::kFragmentsPerIteration
-                                          : kPartitionsK;
-    static int constexpr kSmemPointerOffset =
-        SharedStorage::StorageShape::kCount / kSmemTiles;
-
-    public:
-    static_assert(
-        SharedLoadIterator::Fragment::kElements ==
-            OutputTileIterator::Fragment::kElements,
-        "Mismatch between shared load iterator and output tile iterator.");
-
-    static_assert(OutputTileIterator::kElementsPerAccess,
-                  "OutputTileIterator::kElementsPerAccess must not be zero.");
-
-    static_assert(!(OutputTileIterator::Fragment::kElements %
-                    OutputTileIterator::kElementsPerAccess),
-                  "Divisibility");
-
-    private:
-    // Loads fragment from shared memory aligned with output tensor
-    SharedLoadIterator shared_load_iterator0_;
-    SharedLoadIterator shared_load_iterator1_;
-
-    // Stores a warp's fragment of accumulators to SMEM
-    WarpTileIterator warp_tile_iterator0_;
-    WarpTileIterator warp_tile_iterator1_;
-
-    public:
-    // Constructor
+    // Returns a tensor reference to the shared memory buffer
     CUTLASS_DEVICE
-    DualEpilogue(
-        SharedStorage &shared_storage,  // Shared storage object  // NOLINT
-        int thread_idx,                 // ID of a thread within the threadblock
-        int warp_idx,                   // ID of warp within threadblock
-        int lane_idx                    // Id of thread within warp
-        )
-        : shared_load_iterator0_(shared_storage.reference(0), thread_idx),
-          shared_load_iterator1_(shared_storage.reference(1), thread_idx),
-          warp_tile_iterator0_(shared_storage.reference(0), lane_idx),
-          warp_tile_iterator1_(shared_storage.reference(1), lane_idx) {
-        int warp_k = warp_idx / (WarpCount::kM * WarpCount::kN);
-        int warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN);
-        int warp_m = warp_mn % WarpCount::kM;
-        int warp_n = warp_mn / WarpCount::kM;
+    TensorRef reference(int i) {
+      return TensorRef(
+          storage[i].data(),
+          Layout::packed({StorageShape::kRow, StorageShape::kColumn}));
+    }
+  };
 
-        MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, warp_n};
+  static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1
+                                        ? Base::kFragmentsPerIteration
+                                        : kPartitionsK;
+  static int constexpr kSmemPointerOffset =
+      SharedStorage::StorageShape::kCount / kSmemTiles;
 
-        warp_tile_iterator0_.add_tile_offset(warp_offset);
-        warp_tile_iterator1_.add_tile_offset(warp_offset);
+ public:
+  static_assert(
+      SharedLoadIterator::Fragment::kElements ==
+          OutputTileIterator::Fragment::kElements,
+      "Mismatch between shared load iterator and output tile iterator.");
+
+  static_assert(OutputTileIterator::kElementsPerAccess,
+                "OutputTileIterator::kElementsPerAccess must not be zero.");
+
+  static_assert(!(OutputTileIterator::Fragment::kElements %
+                  OutputTileIterator::kElementsPerAccess),
+                "Divisibility");
+
+ private:
+  // Loads fragment from shared memory aligned with output tensor
+  SharedLoadIterator shared_load_iterator0_;
+  SharedLoadIterator shared_load_iterator1_;
+
+  // Stores a warp's fragment of accumulators to SMEM
+  WarpTileIterator warp_tile_iterator0_;
+  WarpTileIterator warp_tile_iterator1_;
+
+ public:
+  // Constructor
+  CUTLASS_DEVICE
+  DualEpilogue(
+      SharedStorage &shared_storage,  // Shared storage object  // NOLINT
+      int thread_idx,                 // ID of a thread within the threadblock
+      int warp_idx,                   // ID of warp within threadblock
+      int lane_idx                    // Id of thread within warp
+      )
+      : shared_load_iterator0_(shared_storage.reference(0), thread_idx),
+        shared_load_iterator1_(shared_storage.reference(1), thread_idx),
+        warp_tile_iterator0_(shared_storage.reference(0), lane_idx),
+        warp_tile_iterator1_(shared_storage.reference(1), lane_idx) {
+    int warp_k = warp_idx / (WarpCount::kM * WarpCount::kN);
+    int warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN);
+    int warp_m = warp_mn % WarpCount::kM;
+    int warp_n = warp_mn / WarpCount::kM;
+
+    MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, warp_n};
+
+    warp_tile_iterator0_.add_tile_offset(warp_offset);
+    warp_tile_iterator1_.add_tile_offset(warp_offset);
+  }
+
+  // Streams the result to global memory
+  CUTLASS_DEVICE
+  void operator()(OutputOp0 const &output_op0,
+                  OutputOp1 const &output_op1,
+                  OutputOp2 const &output_op2,
+                  OutputTileIterator dest0,
+                  OutputTileIterator dest1,
+                  OutputTileIterator2 dest2,
+                  AccumulatorTile const &accumulator0,
+                  AccumulatorTile const &accumulator1,
+                  OutputTileIterator source_iterator[2],
+                  bool writeToD2  // true if it's the final split-k
+  ) {
+    // TODO: Implement when no source is needed  // NOLINT
+    typename OutputTileIterator::Fragment source_fragment[2];
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < 2; ++i) {
+      source_fragment[i].clear();
     }
 
-    // Streams the result to global memory
-    CUTLASS_DEVICE
-    void operator()(OutputOp0 const &output_op0,
-                    OutputOp1 const &output_op1,
-                    OutputOp2 const &output_op2,
-                    OutputTileIterator dest0,
-                    OutputTileIterator dest1,
-                    OutputTileIterator2 dest2,
-                    AccumulatorTile const &accumulator0,
-                    AccumulatorTile const &accumulator1,
-                    OutputTileIterator source_iterator[2],
-                    bool writeToD2  // true if it's the final split-k
-    ) {
-      // TODO: Implement when no source is needed  // NOLINT
-      typename OutputTileIterator::Fragment source_fragment[2];
+    //
+    // Iterator over warp-level accumulator fragment
+    //
+
+    AccumulatorFragmentIterator accum_fragment_iterator[2] = {accumulator0,
+                                                              accumulator1};
+
+    //
+    // Iterate over accumulator tile
+    //
+
+#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1)
+    for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
+      //
+      // Load the source
+      //
+
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < 2; ++i) {
-          source_fragment[i].clear();
+        source_iterator[i].load(source_fragment[i]);
+        ++source_iterator[i];
       }
 
       //
-      // Iterator over warp-level accumulator fragment
+      // Convert and store fragment
       //
 
-      AccumulatorFragmentIterator accum_fragment_iterator[2] = {accumulator0,
-                                                                accumulator1};
+      __syncthreads();
+
+      acc2smem_source_needed<cutlass::make_index_sequence<
+          OutputTileIterator::kIterations>>::push(iter,
+                                                  accum_fragment_iterator[0],
+                                                  this->warp_tile_iterator0_);
+      acc2smem_source_needed<cutlass::make_index_sequence<
+          OutputTileIterator::kIterations>>::push(iter,
+                                                  accum_fragment_iterator[1],
+                                                  this->warp_tile_iterator1_);
+
+      __syncthreads();
 
       //
-      // Iterate over accumulator tile
+      // Load fragments from shared memory
       //
 
-      #pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1)
-      for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
-          //
-          // Load the source
-          //
+      typename SharedLoadIterator::Fragment
+          aligned_accum_fragment0[kPartitionsK];
+      typename SharedLoadIterator::Fragment
+          aligned_accum_fragment1[kPartitionsK];
 
-          CUTLASS_PRAGMA_UNROLL
-          for (int i = 0; i < 2; ++i) {
-              source_iterator[i].load(source_fragment[i]);
-              ++source_iterator[i];
-          }
+      shared_load_iterator0_.load(aligned_accum_fragment0[0]);
+      shared_load_iterator1_.load(aligned_accum_fragment1[0]);
 
-          //
-          // Convert and store fragment
-          //
-
-          __syncthreads();
-
-          acc2smem_source_needed<cutlass::make_index_sequence<
-              OutputTileIterator::kIterations>>::push(iter,
-                                                      accum_fragment_iterator[0],
-                                                      this->warp_tile_iterator0_);
-          acc2smem_source_needed<cutlass::make_index_sequence<
-              OutputTileIterator::kIterations>>::push(iter,
-                                                      accum_fragment_iterator[1],
-                                                      this->warp_tile_iterator1_);
-
-          __syncthreads();
-
-          //
-          // Load fragments from shared memory
-          //
-
-          typename SharedLoadIterator::Fragment
-              aligned_accum_fragment0[kPartitionsK];
-          typename SharedLoadIterator::Fragment
-              aligned_accum_fragment1[kPartitionsK];
-
-          shared_load_iterator0_.load(aligned_accum_fragment0[0]);
-          shared_load_iterator1_.load(aligned_accum_fragment1[0]);
-
-          // If the number of k-slices is > 1 - perform a reduction amongst the
-          // k-slices
-          if (kPartitionsK > 1) {
-            plus<typename SharedLoadIterator::Fragment> add_fragments;
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int i = 1; i < kPartitionsK; ++i) {
-              shared_load_iterator0_.add_pointer_offset(kSmemPointerOffset);
-              shared_load_iterator1_.add_pointer_offset(kSmemPointerOffset);
-              shared_load_iterator0_.load(aligned_accum_fragment0[i]);
-              shared_load_iterator1_.load(aligned_accum_fragment1[i]);
-              aligned_accum_fragment0[0] = add_fragments(
-                  aligned_accum_fragment0[0], aligned_accum_fragment0[i]);
-              aligned_accum_fragment1[0] = add_fragments(
-                  aligned_accum_fragment1[0], aligned_accum_fragment1[i]);
-            }
-
-            shared_load_iterator0_.add_pointer_offset((1 - kPartitionsK) *
-                                                      kSmemPointerOffset);
-            shared_load_iterator1_.add_pointer_offset((1 - kPartitionsK) *
-                                                      kSmemPointerOffset);
-          }
-
-          //
-          // Compute the output result
-          //
-
-          typename OutputTileIterator::Fragment output_fragment[2];
-          typename OutputTileIterator2::Fragment output_fragment_final;
-
-          apply_output_operator_(output_fragment,
-                                output_fragment_final,
-                                output_op0,
-                                output_op1,
-                                output_op2,
-                                aligned_accum_fragment0[0],
-                                aligned_accum_fragment1[0],
-                                source_fragment);
-
-          //
-          // Store the final result
-          //
-
-          if (kStoreD0) {
-              dest0.store(output_fragment[0]);
-              ++dest0;
-          }
-          if (kStoreD1) {
-              dest1.store(output_fragment[1]);
-              ++dest1;
-          }
-          if (writeToD2) {
-              dest2.store(output_fragment_final);
-              ++dest2;
-          }
-      }
-    }
-
-    private:
-    static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1,
-                  "One of these must be exactly 1.");
-
-    template <class Seq>
-    struct acc2smem_source_needed;
-
-    template <size_t... Seq>
-    struct acc2smem_source_needed<cutlass::index_sequence<Seq...>> {
-        template <int Advance>
-        CUTLASS_DEVICE static void helper(
-            AccumulatorFragmentIterator accum_fragment_iterator,
-            WarpTileIterator &warp_tile_iterator) {  // NOLINT
-            CUTLASS_PRAGMA_UNROLL
-            for (int i = 0; i < Advance; i++) {
-              ++accum_fragment_iterator;
-            }
-
-            typename AccumulatorFragmentIterator::Fragment accum_fragment;
-            accum_fragment_iterator.load(accum_fragment);
-            warp_tile_iterator.store(accum_fragment);
-        }
-
-        CUTLASS_DEVICE
-        static void push(size_t pos,
-                        AccumulatorFragmentIterator const &iterator_begin,
-                        WarpTileIterator &warp_tile_iterator) {  // NOLINT
-            int dummy[] = {(pos == Seq) &&
-                          (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
-        }
-    };
-
-    // Helper to invoke the output functor over each vector of output
-    CUTLASS_DEVICE
-    void apply_output_operator_(
-        typename OutputTileIterator::Fragment (&output_fragment)[2],
-        typename OutputTileIterator2::Fragment &output_fragment_final,  // NOLINT
-        OutputOp0 const &output_op0,
-        OutputOp1 const &output_op1,
-        OutputOp2 const &output_op2,
-        typename SharedLoadIterator::Fragment const &aligned_accum_fragment0,
-        typename SharedLoadIterator::Fragment const &aligned_accum_fragment1,
-        typename OutputTileIterator::Fragment const (&source_fragment)[2]) {
-        OutputAccessType *output_frag_ptr[2] = {
-            reinterpret_cast<OutputAccessType *>(&output_fragment[0]),
-            reinterpret_cast<OutputAccessType *>(&output_fragment[1])};
-
-        OutputAccessType2 *output_frag_final_ptr =
-            reinterpret_cast<OutputAccessType2 *>(&output_fragment_final);
-
-        AccumulatorAccessType const *compute_frag_ptr[2] = {
-            reinterpret_cast<AccumulatorAccessType const *>(
-                &aligned_accum_fragment0),
-            reinterpret_cast<AccumulatorAccessType const *>(
-                &aligned_accum_fragment1)};
-
-        OutputAccessType const *source_frag_ptr[2] = {
-            reinterpret_cast<OutputAccessType const *>(&source_fragment[0]),
-            reinterpret_cast<OutputAccessType const *>(&source_fragment[1])};
-
-        int const kOutputOpIterations = OutputTileIterator::Fragment::kElements /
-                                        OutputTileIterator::kElementsPerAccess;
+      // If the number of k-slices is > 1 - perform a reduction amongst the
+      // k-slices
+      if (kPartitionsK > 1) {
+        plus<typename SharedLoadIterator::Fragment> add_fragments;
 
         CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < kOutputOpIterations; ++i) {
-            // Call the output operators
-            output_frag_ptr[0][i] =
-                output_op0(compute_frag_ptr[0][i], source_frag_ptr[0][i]);
-            output_frag_ptr[1][i] =
-                output_op1(compute_frag_ptr[1][i], source_frag_ptr[1][i]);
-            output_frag_final_ptr[i] =
-                output_op2(output_frag_ptr[0][i], output_frag_ptr[1][i]);
+        for (int i = 1; i < kPartitionsK; ++i) {
+          shared_load_iterator0_.add_pointer_offset(kSmemPointerOffset);
+          shared_load_iterator1_.add_pointer_offset(kSmemPointerOffset);
+          shared_load_iterator0_.load(aligned_accum_fragment0[i]);
+          shared_load_iterator1_.load(aligned_accum_fragment1[i]);
+          aligned_accum_fragment0[0] = add_fragments(
+              aligned_accum_fragment0[0], aligned_accum_fragment0[i]);
+          aligned_accum_fragment1[0] = add_fragments(
+              aligned_accum_fragment1[0], aligned_accum_fragment1[i]);
         }
+
+        shared_load_iterator0_.add_pointer_offset((1 - kPartitionsK) *
+                                                  kSmemPointerOffset);
+        shared_load_iterator1_.add_pointer_offset((1 - kPartitionsK) *
+                                                  kSmemPointerOffset);
+      }
+
+      //
+      // Compute the output result
+      //
+
+      typename OutputTileIterator::Fragment output_fragment[2];
+      typename OutputTileIterator2::Fragment output_fragment_final;
+
+      apply_output_operator_(output_fragment,
+                             output_fragment_final,
+                             output_op0,
+                             output_op1,
+                             output_op2,
+                             aligned_accum_fragment0[0],
+                             aligned_accum_fragment1[0],
+                             source_fragment);
+
+      //
+      // Store the final result
+      //
+
+      if (kStoreD0) {
+        dest0.store(output_fragment[0]);
+        ++dest0;
+      }
+      if (kStoreD1) {
+        dest1.store(output_fragment[1]);
+        ++dest1;
+      }
+      if (writeToD2) {
+        dest2.store(output_fragment_final);
+        ++dest2;
+      }
     }
+  }
+
+ private:
+  static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1,
+                "One of these must be exactly 1.");
+
+  template <class Seq>
+  struct acc2smem_source_needed;
+
+  template <size_t... Seq>
+  struct acc2smem_source_needed<cutlass::index_sequence<Seq...>> {
+    template <int Advance>
+    CUTLASS_DEVICE static void helper(
+        AccumulatorFragmentIterator accum_fragment_iterator,
+        WarpTileIterator &warp_tile_iterator) {  // NOLINT
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Advance; i++) {
+        ++accum_fragment_iterator;
+      }
+
+      typename AccumulatorFragmentIterator::Fragment accum_fragment;
+      accum_fragment_iterator.load(accum_fragment);
+      warp_tile_iterator.store(accum_fragment);
+    }
+
+    CUTLASS_DEVICE
+    static void push(size_t pos,
+                     AccumulatorFragmentIterator const &iterator_begin,
+                     WarpTileIterator &warp_tile_iterator) {  // NOLINT
+      int dummy[] = {(pos == Seq) &&
+                     (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
+    }
+  };
+
+  // Helper to invoke the output functor over each vector of output
+  CUTLASS_DEVICE
+  void apply_output_operator_(
+      typename OutputTileIterator::Fragment (&output_fragment)[2],
+      typename OutputTileIterator2::Fragment &output_fragment_final,  // NOLINT
+      OutputOp0 const &output_op0,
+      OutputOp1 const &output_op1,
+      OutputOp2 const &output_op2,
+      typename SharedLoadIterator::Fragment const &aligned_accum_fragment0,
+      typename SharedLoadIterator::Fragment const &aligned_accum_fragment1,
+      typename OutputTileIterator::Fragment const (&source_fragment)[2]) {
+    OutputAccessType *output_frag_ptr[2] = {
+        reinterpret_cast<OutputAccessType *>(&output_fragment[0]),
+        reinterpret_cast<OutputAccessType *>(&output_fragment[1])};
+
+    OutputAccessType2 *output_frag_final_ptr =
+        reinterpret_cast<OutputAccessType2 *>(&output_fragment_final);
+
+    AccumulatorAccessType const *compute_frag_ptr[2] = {
+        reinterpret_cast<AccumulatorAccessType const *>(
+            &aligned_accum_fragment0),
+        reinterpret_cast<AccumulatorAccessType const *>(
+            &aligned_accum_fragment1)};
+
+    OutputAccessType const *source_frag_ptr[2] = {
+        reinterpret_cast<OutputAccessType const *>(&source_fragment[0]),
+        reinterpret_cast<OutputAccessType const *>(&source_fragment[1])};
+
+    int const kOutputOpIterations = OutputTileIterator::Fragment::kElements /
+                                    OutputTileIterator::kElementsPerAccess;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kOutputOpIterations; ++i) {
+      // Call the output operators
+      output_frag_ptr[0][i] =
+          output_op0(compute_frag_ptr[0][i], source_frag_ptr[0][i]);
+      output_frag_ptr[1][i] =
+          output_op1(compute_frag_ptr[1][i], source_frag_ptr[1][i]);
+      output_frag_final_ptr[i] =
+          output_op2(output_frag_ptr[0][i], output_frag_ptr[1][i]);
+    }
+  }
 };
 
 }  // namespace threadblock
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_block_gemm_act_template_3x.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_block_gemm_act_template_3x.h
index 1a5b838b81..f95f20bfbf 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_block_gemm_act_template_3x.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_block_gemm_act_template_3x.h
@@ -34,34 +34,35 @@
 using namespace cute;
 
 template <
-  typename InputType = phi::dtype::float8_e4m3fn,
-  typename OutType = phi::dtype::float16,
-  bool hasbias = false,
-  template <class> typename Activation = cutlass::epilogue::thread::Identity,
-  typename TileShape = Shape<_128, _128, _128>,
-  typename ClusterShape = Shape<_1, _2, _1>,
-  typename KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<1>,
-  typename EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative,
-  typename TileSchedule = cutlass::gemm::PersistentScheduler,
-  typename SM = cutlass::arch::Sm90
->
-bool dispatch_fuse_block_gemm_c3x(GemmEpilogueAllParams params){
-  using ElementA = typename std::conditional_t<std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
-                                                              cutlass::float_e4m3_t,
-                                                              cutlass::float_e5m2_t>;
+    typename InputType = phi::dtype::float8_e4m3fn,
+    typename OutType = phi::dtype::float16,
+    bool hasbias = false,
+    template <class> typename Activation = cutlass::epilogue::thread::Identity,
+    typename TileShape = Shape<_128, _128, _128>,
+    typename ClusterShape = Shape<_1, _2, _1>,
+    typename KernelSchedule = cutlass::gemm::
+        KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<1>,
+    typename EpilogueSchedule =
+        cutlass::epilogue::TmaWarpSpecializedCooperative,
+    typename TileSchedule = cutlass::gemm::PersistentScheduler,
+    typename SM = cutlass::arch::Sm90>
+bool dispatch_fuse_block_gemm_c3x(GemmEpilogueAllParams params) {
+  using ElementA = typename std::conditional_t<
+      std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
+      cutlass::float_e4m3_t,
+      cutlass::float_e5m2_t>;
   using ElementB = ElementA;
-  using ElementD = typename std::conditional_t<std::is_same_v<OutType, phi::dtype::bfloat16>,
+  using ElementD =
+      typename std::conditional_t<std::is_same_v<OutType, phi::dtype::bfloat16>,
                                   cutlass::bfloat16_t,
                                   cutlass::half_t>;
-  using ElementC = std::conditional_t<
-        hasbias,
-        ElementD,
-        void>;
+  using ElementC = std::conditional_t<hasbias, ElementD, void>;
 
   constexpr int ScaleMsPerTile = size<0>(TileShape{});
   constexpr int ScaleGranularityM = size<0>(TileShape{}) / ScaleMsPerTile;
 
-  static constexpr bool IsStreamK = cute::is_same_v<TileSchedule, cutlass::gemm::StreamKScheduler>;
+  static constexpr bool IsStreamK =
+      cute::is_same_v<TileSchedule, cutlass::gemm::StreamKScheduler>;
 
   using LayoutA = cutlass::layout::RowMajor;
   using LayoutB = cutlass::layout::ColumnMajor;
@@ -80,37 +81,54 @@ bool dispatch_fuse_block_gemm_c3x(GemmEpilogueAllParams params){
 
   static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
 
-  using FusionOperation = cutlass::epilogue::fusion::LinCombEltAct<Activation, ElementD, ElementCompute, ElementC, ElementScalar, RoundStyle>;
+  using FusionOperation =
+      cutlass::epilogue::fusion::LinCombEltAct<Activation,
+                                               ElementD,
+                                               ElementCompute,
+                                               ElementC,
+                                               ElementScalar,
+                                               RoundStyle>;
 
-  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
-      SM, cutlass::arch::OpClassTensorOp,
-      TileShape,
-      ClusterShape,
-      cutlass::epilogue::collective::EpilogueTileAuto,
-      ElementAccumulator, ElementCompute,
-      ElementC, LayoutC, AlignmentC,
-      ElementD, LayoutD, AlignmentD,
-      EpilogueSchedule,
-      FusionOperation
-    >::CollectiveOp;
+  using CollectiveEpilogue =
+      typename cutlass::epilogue::collective::CollectiveBuilder<
+          SM,
+          cutlass::arch::OpClassTensorOp,
+          TileShape,
+          ClusterShape,
+          cutlass::epilogue::collective::EpilogueTileAuto,
+          ElementAccumulator,
+          ElementCompute,
+          ElementC,
+          LayoutC,
+          AlignmentC,
+          ElementD,
+          LayoutD,
+          AlignmentD,
+          EpilogueSchedule,
+          FusionOperation>::CollectiveOp;
 
-  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
-      SM, cutlass::arch::OpClassTensorOp,
-      ElementA, LayoutA, AlignmentA,
-      ElementB, LayoutB, AlignmentB,
-      ElementAccumulator,
-      TileShape,
-      ClusterShape,
-      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
-      KernelSchedule
-    >::CollectiveOp;
+  using CollectiveMainloop =
+      typename cutlass::gemm::collective::CollectiveBuilder<
+          SM,
+          cutlass::arch::OpClassTensorOp,
+          ElementA,
+          LayoutA,
+          AlignmentA,
+          ElementB,
+          LayoutB,
+          AlignmentB,
+          ElementAccumulator,
+          TileShape,
+          ClusterShape,
+          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+              sizeof(typename CollectiveEpilogue::SharedStorage))>,
+          KernelSchedule>::CollectiveOp;
 
-  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
-      Shape<int,int,int,int>,
-      CollectiveMainloop,
-      CollectiveEpilogue,
-      TileSchedule
-  >;
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>,
+                                           CollectiveMainloop,
+                                           CollectiveEpilogue,
+                                           TileSchedule>;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
 
@@ -132,23 +150,27 @@ bool dispatch_fuse_block_gemm_c3x(GemmEpilogueAllParams params){
   StrideD stride_D{params.ldd, cute::Int<1>{}, params.M * params.ldd};
 
   auto a_ptr = reinterpret_cast<ElementA*>(const_cast<void*>(params.A));
-  auto a_scale_ptr = reinterpret_cast<float*>(const_cast<void*>(params.A_scale));
+  auto a_scale_ptr =
+      reinterpret_cast<float*>(const_cast<void*>(params.A_scale));
   auto b_ptr = reinterpret_cast<ElementB*>(const_cast<void*>(params.B));
-  auto b_scale_ptr = reinterpret_cast<float*>(const_cast<void*>(params.B_scale));
+  auto b_scale_ptr =
+      reinterpret_cast<float*>(const_cast<void*>(params.B_scale));
   auto c_ptr = reinterpret_cast<ElementC*>(const_cast<void*>(params.bias));
   auto d_ptr = reinterpret_cast<ElementD*>(params.D);
 
-  ProblemShapeType problem_size = ProblemShapeType{params.M, params.N, params.K, params.batch_count};
+  ProblemShapeType problem_size =
+      ProblemShapeType{params.M, params.N, params.K, params.batch_count};
 
   typename Gemm::Arguments arguments{
-    cutlass::gemm::GemmUniversalMode::kGemm,
-    problem_size,
-    {a_ptr, stride_A, b_ptr, stride_B,
-    a_scale_ptr, b_scale_ptr},
-    {{params.scale}, // epilogue.thread
-      c_ptr, stride_C, d_ptr, stride_D}
-  };
-  if constexpr (hasbias){
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      problem_size,
+      {a_ptr, stride_A, b_ptr, stride_B, a_scale_ptr, b_scale_ptr},
+      {{params.scale},  // epilogue.thread
+       c_ptr,
+       stride_C,
+       d_ptr,
+       stride_D}};
+  if constexpr (hasbias) {
     arguments.epilogue.thread.beta = 1.0;
   }
 
@@ -162,12 +184,12 @@ bool dispatch_fuse_block_gemm_c3x(GemmEpilogueAllParams params){
     arguments.scheduler.reduction_mode = ReductionMode::Nondeterministic;
   }
 
-
   Gemm gemm_op;
 
   cutlass::Status status = gemm_op.can_implement(arguments);
   if (status != cutlass::Status::kSuccess) {
-    std::cout << "Gemm::can_implement() failed. " << cutlassGetStatusString(status) << std::endl;
+    std::cout << "Gemm::can_implement() failed. "
+              << cutlassGetStatusString(status) << std::endl;
     return false;
   }
   size_t workspace_size = Gemm::get_workspace_size(arguments);
@@ -176,7 +198,8 @@ bool dispatch_fuse_block_gemm_c3x(GemmEpilogueAllParams params){
 
   status = gemm_op(arguments, workspace->ptr(), params.stream);
   if (status != cutlass::Status::kSuccess) {
-    std::cout << "Gemm::run() failed." << cutlassGetStatusString(status) << std::endl;
+    std::cout << "Gemm::run() failed." << cutlassGetStatusString(status)
+              << std::endl;
     return false;
   }
   return true;
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_act_template_3x.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_act_template_3x.h
index 632cdc296a..ec484ac743 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_act_template_3x.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_act_template_3x.h
@@ -28,60 +28,67 @@
 #include "cutlass_extensions/gemm/collective/collective_builder_gated.hpp"
 #include "cutlass_extensions/gemm/kernel/gemm_universal_gated.hpp"
 
-template <typename InputType, typename CTAShape, typename ClusterShape,
-          typename MainloopScheduleType, typename EpilogueScheduleType,
+template <typename InputType,
+          typename CTAShape,
+          typename ClusterShape,
+          typename MainloopScheduleType,
+          typename EpilogueScheduleType,
           typename TileSchedulerType = void,
-          template <class /* ElementCompute */> class Activation =
-              cutlass::epilogue::thread::SiLu,
+          template <class /* ElementCompute */>
+          class Activation = cutlass::epilogue::thread::SiLu,
           bool SwapAB = true>
 bool dispatch_dual_gemm_act_sm90(DualGemmEpilogueAllParams params) {
   using namespace cute;
   using ElementA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
-      cutlass::float_e4m3_t, cutlass::float_e5m2_t>;
-  using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand
+      cutlass::float_e4m3_t,
+      cutlass::float_e5m2_t>;
+  using LayoutA =
+      cutlass::layout::RowMajor;  // Layout type for A matrix operand
   static constexpr int AlignmentA =
       128 /
       cutlass::sizeof_bits<
-          ElementA>::value; // Memory access granularity/alignment of A
-                            // matrix in units of elements (up to 16 bytes)
+          ElementA>::value;  // Memory access granularity/alignment of A
+                             // matrix in units of elements (up to 16 bytes)
 
   // B matrix configuration
-  using ElementB = ElementA; // Element type for B matrix operand
+  using ElementB = ElementA;  // Element type for B matrix operand
   using LayoutB =
-      cutlass::layout::ColumnMajor; // Layout type for B matrix operand
+      cutlass::layout::ColumnMajor;  // Layout type for B matrix operand
   static constexpr int AlignmentB =
       128 /
       cutlass::sizeof_bits<
-          ElementB>::value; // Memory access granularity/alignment of B
-                            // matrix in units of elements (up to 16 bytes)
+          ElementB>::value;  // Memory access granularity/alignment of B
+                             // matrix in units of elements (up to 16 bytes)
 
-  using ElementC = ElementA; // Element type for C matrix operands
+  using ElementC = ElementA;  // Element type for C matrix operands
 
-  using LayoutC = cute::conditional_t<SwapAB, cutlass::layout::ColumnMajor,
+  using LayoutC = cute::conditional_t<SwapAB,
+                                      cutlass::layout::ColumnMajor,
                                       cutlass::layout::RowMajor>;
   static constexpr int AlignmentC =
       128 /
       cutlass::sizeof_bits<
-          ElementC>::value; // Memory access granularity/alignment of C matrices
-                            // in units of elements (up to 16 bytes)
+          ElementC>::value;  // Memory access granularity/alignment of C
+                             // matrices in units of elements (up to 16 bytes)
 
   // Output matrix configuration
-  using ElementOutput = ElementA; // Element type for output matrix operands
+  using ElementOutput = ElementA;  // Element type for output matrix operands
   // using LayoutOutput = cutlass::layout::RowMajor; // Layout type for output
   // matrix operands
-  using LayoutOutput = cute::conditional_t<SwapAB, cutlass::layout::ColumnMajor,
+  using LayoutOutput = cute::conditional_t<SwapAB,
+                                           cutlass::layout::ColumnMajor,
                                            cutlass::layout::RowMajor>;
   static constexpr int AlignmentOutput =
       128 / cutlass::sizeof_bits<ElementOutput>::value;
 
   // Multiply-accumulate blocking/pipelining details
-  using ElementAccumulator = float;    // Element type for internal accumulation
-  using ElementCompute = float;        // Element type for compute
-  using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that
-                                       // supports the intended feature
-  using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag
-  using TileShape = CTAShape; // Threadblock-level tile size
+  using ElementAccumulator = float;  // Element type for internal accumulation
+  using ElementCompute = float;      // Element type for compute
+  using ArchTag = cutlass::arch::Sm90;  // Tag indicating the minimum SM that
+                                        // supports the intended feature
+  using OperatorClass = cutlass::arch::OpClassTensorOp;  // Operator class tag
+  using TileShape = CTAShape;  // Threadblock-level tile size
   using KernelSchedule = MainloopScheduleType;
   using EpilogueSchedule = EpilogueScheduleType;
   using TileScheduler = TileSchedulerType;
@@ -94,22 +101,46 @@ bool dispatch_dual_gemm_act_sm90(DualGemmEpilogueAllParams params) {
 
   using CollectiveEpilogue =
       typename cutlass::epilogue::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, TileShape, ClusterShape, EpilogueTileType,
-          ElementAccumulator, ElementAccumulator, ElementC, LayoutC, AlignmentC,
-          ElementOutput, LayoutOutput, AlignmentOutput, EpilogueSchedule,
+          ArchTag,
+          OperatorClass,
+          TileShape,
+          ClusterShape,
+          EpilogueTileType,
+          ElementAccumulator,
+          ElementAccumulator,
+          ElementC,
+          LayoutC,
+          AlignmentC,
+          ElementOutput,
+          LayoutOutput,
+          AlignmentOutput,
+          EpilogueSchedule,
           FusionOperation>::CollectiveOp;
 
   using CollectiveMainloop =
       typename cutlass::gemm::collective::CollectiveBuilderGated<
-          ArchTag, OperatorClass, ElementA, LayoutA, AlignmentA, ElementB,
-          LayoutB, AlignmentB, ElementAccumulator, TileShape, ClusterShape,
+          ArchTag,
+          OperatorClass,
+          ElementA,
+          LayoutA,
+          AlignmentA,
+          ElementB,
+          LayoutB,
+          AlignmentB,
+          ElementAccumulator,
+          TileShape,
+          ClusterShape,
           cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
               sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          KernelSchedule, Activation, SwapAB>::CollectiveOp;
+          KernelSchedule,
+          Activation,
+          SwapAB>::CollectiveOp;
 
   using GemmKernel = cutlass::gemm::kernel::GemmUniversalGated<
-      Shape<int, int, int, int>, // Indicates ProblemShape
-      CollectiveMainloop, CollectiveEpilogue, TileScheduler>;
+      Shape<int, int, int, int>,  // Indicates ProblemShape
+      CollectiveMainloop,
+      CollectiveEpilogue,
+      TileScheduler>;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
 
@@ -141,7 +172,7 @@ bool dispatch_dual_gemm_act_sm90(DualGemmEpilogueAllParams params) {
       cutlass::gemm::GemmUniversalMode::kGemm,
       {arg_m, arg_n, params.K, params.batch_count},
       {ptr_A, stride_A, ptr_B0, ptr_B1, stride_B, params.scale0, params.scale1},
-      {{}, // epilogue.thread
+      {{},  // epilogue.thread
        nullptr,
        stride_C,
        reinterpret_cast<ElementOutput *>(params.D),
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_geglu_template.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_geglu_template.h
index 0e18c2a389..5d2a9638b3 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_geglu_template.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_geglu_template.h
@@ -21,9 +21,15 @@
 #include "fp8_gemm_fused/dual_gemm/device/dual_gemm.h"
 #include "fp8_gemm_fused/dual_gemm/thread/left_gelu_and_mul.h"
 
-template <typename InputType, typename OutType, typename BiasType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename BiasType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_dual_gemm_geglu(DualGemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -73,8 +79,8 @@ bool dispatch_dual_gemm_geglu(DualGemmEpilogueAllParams params) {
       cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;  // <- ??
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   using EpilogueOp0 = cutlass::epilogue::thread::LinearCombination<
       ElementInputC,  // <- data type of output matrix
@@ -86,7 +92,7 @@ bool dispatch_dual_gemm_geglu(DualGemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   using EpilogueOp1 = cutlass::epilogue::thread::LeftGELUAndMul<
       ElementOutput,
@@ -143,11 +149,23 @@ bool dispatch_dual_gemm_geglu(DualGemmEpilogueAllParams params) {
        params.lda},
       {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B0)),
        params.ldb},
-      hasbias? typename cutlass::TensorRef<typename Gemm::ElementC, typename Gemm::LayoutC>{reinterpret_cast<ElementInputC*>(const_cast<void*>(params.bias0)), 0} : nullptr_ref,
+      hasbias
+          ? typename cutlass::TensorRef<
+                typename Gemm::ElementC,
+                typename Gemm::LayoutC>{reinterpret_cast<ElementInputC*>(
+                                            const_cast<void*>(params.bias0)),
+                                        0}
+          : nullptr_ref,
       nullptr_ref,
       {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B1)),
        params.ldb},
-      hasbias? typename cutlass::TensorRef<typename Gemm::ElementC, typename Gemm::LayoutC>{reinterpret_cast<ElementInputC*>(const_cast<void*>(params.bias1)), 0} : nullptr_ref,
+      hasbias
+          ? typename cutlass::TensorRef<
+                typename Gemm::ElementC,
+                typename Gemm::LayoutC>{reinterpret_cast<ElementInputC*>(
+                                            const_cast<void*>(params.bias1)),
+                                        0}
+          : nullptr_ref,
       nullptr_ref,
       {reinterpret_cast<ElementOutput*>(const_cast<void*>(params.D)),
        params.ldd},
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_swiglu_template.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_swiglu_template.h
index b5de12e7e9..dc78839750 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_swiglu_template.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_swiglu_template.h
@@ -21,9 +21,15 @@
 #include "fp8_gemm_fused/dual_gemm/device/dual_gemm.h"
 #include "fp8_gemm_fused/dual_gemm/thread/left_silu_and_mul.h"
 
-template <typename InputType, typename OutType, typename BiasType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename BiasType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_dual_gemm_swiglu(DualGemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -73,8 +79,8 @@ bool dispatch_dual_gemm_swiglu(DualGemmEpilogueAllParams params) {
       cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;  // <- ??
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   using EpilogueOp0 = cutlass::epilogue::thread::LinearCombination<
       ElementInputC,  // <- data type of output matrix
@@ -86,7 +92,7 @@ bool dispatch_dual_gemm_swiglu(DualGemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   using EpilogueOp1 = cutlass::epilogue::thread::LeftSiLUAndMul<
       ElementOutput,
@@ -129,9 +135,9 @@ bool dispatch_dual_gemm_swiglu(DualGemmEpilogueAllParams params) {
   cutlass::gemm::GemmCoord problem_size =
       cutlass::gemm::GemmCoord{params.M, params.N, params.K};
 
-  cutlass::gemm::DualGemmMode mode = params.batch_count > 1 ?
-        cutlass::gemm::DualGemmMode::kBatched :
-        cutlass::gemm::DualGemmMode::kGemm;
+  cutlass::gemm::DualGemmMode mode = params.batch_count > 1
+                                         ? cutlass::gemm::DualGemmMode::kBatched
+                                         : cutlass::gemm::DualGemmMode::kGemm;
 
   typename cutlass::TensorRef<typename Gemm::ElementC, typename Gemm::LayoutC>
       nullptr_ref{};
@@ -144,11 +150,23 @@ bool dispatch_dual_gemm_swiglu(DualGemmEpilogueAllParams params) {
        params.lda},
       {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B0)),
        params.ldb},
-      hasbias ? typename cutlass::TensorRef<typename Gemm::ElementC, typename Gemm::LayoutC>{reinterpret_cast<ElementInputC*>(const_cast<void*>(params.bias0)), 0} : nullptr_ref,
+      hasbias
+          ? typename cutlass::TensorRef<
+                typename Gemm::ElementC,
+                typename Gemm::LayoutC>{reinterpret_cast<ElementInputC*>(
+                                            const_cast<void*>(params.bias0)),
+                                        0}
+          : nullptr_ref,
       nullptr_ref,
       {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B1)),
        params.ldb},
-      hasbias ? typename cutlass::TensorRef<typename Gemm::ElementC, typename Gemm::LayoutC>{reinterpret_cast<ElementInputC*>(const_cast<void*>(params.bias1)), 0} : nullptr_ref,
+      hasbias
+          ? typename cutlass::TensorRef<
+                typename Gemm::ElementC,
+                typename Gemm::LayoutC>{reinterpret_cast<ElementInputC*>(
+                                            const_cast<void*>(params.bias1)),
+                                        0}
+          : nullptr_ref,
       nullptr_ref,
       {reinterpret_cast<ElementOutput*>(const_cast<void*>(params.D)),
        params.ldd},
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_act_template_3x.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_act_template_3x.h
index c470151070..83ded56270 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_act_template_3x.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_act_template_3x.h
@@ -26,26 +26,28 @@
 #include "cutlass/gemm/kernel/tile_scheduler.hpp"
 #include "cutlass/util/packed_stride.hpp"
 
-template <
-    typename InputType,
-    typename OutType,
-    bool hasbias,
-    template <class> typename Activation,
-    typename TileShape,
-    typename ClusterShape,
-    typename KernelSchedule =
-        cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum,
-    typename EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized,
-    typename SM = cutlass::arch::Sm90>
+template <typename InputType,
+          typename OutType,
+          bool hasbias,
+          template <class>
+          typename Activation,
+          typename TileShape,
+          typename ClusterShape,
+          typename KernelSchedule =
+              cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum,
+          typename EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized,
+          typename SM = cutlass::arch::Sm90>
 bool dispatch_fuse_gemm_act_sm90(GemmEpilogueAllParams params) {
   using namespace cute;
   using ElementA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
-      cutlass::float_e4m3_t, cutlass::float_e5m2_t>;
+      cutlass::float_e4m3_t,
+      cutlass::float_e5m2_t>;
   using ElementB = ElementA;
   using ElementD =
       typename std::conditional_t<std::is_same_v<OutType, phi::dtype::bfloat16>,
-                                  cutlass::bfloat16_t, cutlass::half_t>;
+                                  cutlass::bfloat16_t,
+                                  cutlass::half_t>;
   using ElementC = std::conditional_t<hasbias, ElementD, void>;
 
   using LayoutA = cutlass::layout::RowMajor;
@@ -66,29 +68,53 @@ bool dispatch_fuse_gemm_act_sm90(GemmEpilogueAllParams params) {
   static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
 
   using FusionOperation =
-      cutlass::epilogue::fusion::LinCombEltAct<Activation, ElementD,
-                                               ElementCompute, ElementC,
-                                               ElementScalar, RoundStyle>;
+      cutlass::epilogue::fusion::LinCombEltAct<Activation,
+                                               ElementD,
+                                               ElementCompute,
+                                               ElementC,
+                                               ElementScalar,
+                                               RoundStyle>;
 
   using CollectiveEpilogue =
       typename cutlass::epilogue::collective::CollectiveBuilder<
-          SM, cutlass::arch::OpClassTensorOp, TileShape, ClusterShape,
-          cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
-          ElementCompute, ElementC, LayoutC, AlignmentC, ElementD, LayoutD,
-          AlignmentD, EpilogueSchedule, FusionOperation>::CollectiveOp;
+          SM,
+          cutlass::arch::OpClassTensorOp,
+          TileShape,
+          ClusterShape,
+          cutlass::epilogue::collective::EpilogueTileAuto,
+          ElementAccumulator,
+          ElementCompute,
+          ElementC,
+          LayoutC,
+          AlignmentC,
+          ElementD,
+          LayoutD,
+          AlignmentD,
+          EpilogueSchedule,
+          FusionOperation>::CollectiveOp;
 
   using CollectiveMainloop =
       typename cutlass::gemm::collective::CollectiveBuilder<
-          SM, cutlass::arch::OpClassTensorOp, ElementA, LayoutA, AlignmentA,
-          ElementB, LayoutB, AlignmentB, ElementAccumulator, TileShape,
+          SM,
+          cutlass::arch::OpClassTensorOp,
+          ElementA,
+          LayoutA,
+          AlignmentA,
+          ElementB,
+          LayoutB,
+          AlignmentB,
+          ElementAccumulator,
+          TileShape,
           ClusterShape,
           cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
               sizeof(typename CollectiveEpilogue::SharedStorage))>,
           KernelSchedule>::CollectiveOp;
 
-  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
-      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
-      cutlass::gemm::PersistentScheduler>;
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>,
+                                           CollectiveMainloop,
+                                           CollectiveEpilogue,
+                                           cutlass::gemm::PersistentScheduler>;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
 
@@ -120,7 +146,7 @@ bool dispatch_fuse_gemm_act_sm90(GemmEpilogueAllParams params) {
   typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGemm,
                                      problem_size,
                                      {a_ptr, stride_A, b_ptr, stride_B},
-                                     {{params.scale}, // epilogue.thread
+                                     {{params.scale},  // epilogue.thread
                                       c_ptr,
                                       stride_C,
                                       d_ptr,
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_gelu_template.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_gelu_template.h
index 32b8a132e8..50793eff0a 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_gelu_template.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_gelu_template.h
@@ -20,9 +20,14 @@
 #include "cutlass/gemm/device/gemm_universal.h"
 #include "cutlass/gemm/device/gemm_splitk_parallel.h"
 
-template <typename InputType, typename OutType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_fuse_gemm_gelu(GemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -64,9 +69,8 @@ bool dispatch_fuse_gemm_gelu(GemmEpilogueAllParams params) {
   using ShapeMMAOp = MMAShape;  // <- MMA Op tile
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
-
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   // This code section describes how threadblocks are scheduled on GPU
   using SwizzleThreadBlock =
@@ -82,7 +86,7 @@ bool dispatch_fuse_gemm_gelu(GemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   // Number of pipelines you want to use
   constexpr int NumStages = Stages;
@@ -164,10 +168,15 @@ bool dispatch_fuse_gemm_gelu(GemmEpilogueAllParams params) {
   return true;
 }
 
-
-template <typename InputType, typename OutType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, int SplitK, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          int SplitK,
+          bool hasbias,
+          typename SM>
 bool dispatch_fuse_gemm_split_k_gelu(GemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -209,9 +218,8 @@ bool dispatch_fuse_gemm_split_k_gelu(GemmEpilogueAllParams params) {
   using ShapeMMAOp = MMAShape;  // <- MMA Op tile
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
-
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   // This code section describes how threadblocks are scheduled on GPU
   using SwizzleThreadBlock =
@@ -227,7 +235,7 @@ bool dispatch_fuse_gemm_split_k_gelu(GemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   // Number of pipelines you want to use
   constexpr int NumStages = Stages;
@@ -309,10 +317,14 @@ bool dispatch_fuse_gemm_split_k_gelu(GemmEpilogueAllParams params) {
   return true;
 }
 
-
-template <typename InputType, typename OutType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_fuse_gemm_split_k_gelu(GemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -354,8 +366,8 @@ bool dispatch_fuse_gemm_split_k_gelu(GemmEpilogueAllParams params) {
   using ShapeMMAOp = MMAShape;
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   using EpilogueOp = cutlass::epilogue::thread::LinearCombinationGELU<
       ElementOutput,  // <- data type of output matrix
@@ -367,50 +379,55 @@ bool dispatch_fuse_gemm_split_k_gelu(GemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   // Number of pipelines you want to use
   constexpr int NumStages = Stages;
 
-    using ConvertScaledOp = cutlass::epilogue::thread::Convert<
-        ElementAccumulator,
-        cutlass::gemm::device::DefaultGemmConfiguration<cutlass::arch::OpClassSimt, SmArch, ElementInputA, ElementInputB,
-                                    ElementAccumulator,
-                                    ElementAccumulator>::EpilogueOutputOp::kCount,
-        ElementAccumulator>;
+  using ConvertScaledOp = cutlass::epilogue::thread::Convert<
+      ElementAccumulator,
+      cutlass::gemm::device::DefaultGemmConfiguration<
+          cutlass::arch::OpClassSimt,
+          SmArch,
+          ElementInputA,
+          ElementInputB,
+          ElementAccumulator,
+          ElementAccumulator>::EpilogueOutputOp::kCount,
+      ElementAccumulator>;
 
-    /// Reduction operator
-    using ReductionOp = cutlass::reduction::thread::ReduceAdd<
-        ElementAccumulator, typename EpilogueOp::ElementAccumulator,
-        EpilogueOp::kCount>;
+  /// Reduction operator
+  using ReductionOp = cutlass::reduction::thread::ReduceAdd<
+      ElementAccumulator,
+      typename EpilogueOp::ElementAccumulator,
+      EpilogueOp::kCount>;
 
-    /// Threadblock-level swizzling operator
-    using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmSplitKHorizontalThreadblockSwizzle;
+  /// Threadblock-level swizzling operator
+  using ThreadblockSwizzle =
+      cutlass::gemm::threadblock::GemmSplitKHorizontalThreadblockSwizzle;
 
-    /// Operation performed by GEMM
-    using Operator = cutlass::arch::OpMultiplyAddFastAccum;
-
-    using Gemm = cutlass::gemm::device::GemmSplitKParallel<ElementInputA,
-                                                        LayoutInputA,
-                                                        ElementInputB,
-                                                        LayoutInputB,
-                                                        ElementOutput,
-                                                        LayoutOutput,
-                                                        ElementAccumulator,
-                                                        MMAOp,
-                                                        SmArch,
-                                                        ShapeMMAThreadBlock,
-                                                        ShapeMMAWarp,
-                                                        ShapeMMAOp,
-                                                        EpilogueOp,
-                                                        ConvertScaledOp,
-                                                        ReductionOp,
-                                                        ThreadblockSwizzle,
-                                                        NumStages,
-                                                        kAlignmentA,
-                                                        kAlignmentB,
-                                                        Operator>;
+  /// Operation performed by GEMM
+  using Operator = cutlass::arch::OpMultiplyAddFastAccum;
 
+  using Gemm = cutlass::gemm::device::GemmSplitKParallel<ElementInputA,
+                                                         LayoutInputA,
+                                                         ElementInputB,
+                                                         LayoutInputB,
+                                                         ElementOutput,
+                                                         LayoutOutput,
+                                                         ElementAccumulator,
+                                                         MMAOp,
+                                                         SmArch,
+                                                         ShapeMMAThreadBlock,
+                                                         ShapeMMAWarp,
+                                                         ShapeMMAOp,
+                                                         EpilogueOp,
+                                                         ConvertScaledOp,
+                                                         ReductionOp,
+                                                         ThreadblockSwizzle,
+                                                         NumStages,
+                                                         kAlignmentA,
+                                                         kAlignmentB,
+                                                         Operator>;
 
   cutlass::gemm::GemmCoord problem_size =
       cutlass::gemm::GemmCoord{params.M, params.N, params.K};
@@ -421,15 +438,18 @@ bool dispatch_fuse_gemm_split_k_gelu(GemmEpilogueAllParams params) {
   // Split K dimension into 16 partitions
   int split_k_slices = params.split_k;
 
-  // Create a tuple of gemm kernel arguments. This is later passed as arguments to launch
-  // instantiated CUTLASS kernel
-  typename Gemm::Arguments arguments{problem_size,  // <- problem size of matrix multiplication
-                                     {reinterpret_cast<ElementInputA*>(const_cast<void*>(params.A)),params.lda},
-                                     {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B)),params.ldb},
-                                     {reinterpret_cast<ElementOutput*>(const_cast<void*>(params.bias)),0},
-                                     {reinterpret_cast<ElementOutput*>(params.D),params.ldd},
-                                     {alpha, beta},          // <- tuple of alpha and beta
-                                     split_k_slices};        // <- k-dimension split factor
+  // Create a tuple of gemm kernel arguments. This is later passed as arguments
+  // to launch instantiated CUTLASS kernel
+  typename Gemm::Arguments arguments{
+      problem_size,  // <- problem size of matrix multiplication
+      {reinterpret_cast<ElementInputA*>(const_cast<void*>(params.A)),
+       params.lda},
+      {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B)),
+       params.ldb},
+      {reinterpret_cast<ElementOutput*>(const_cast<void*>(params.bias)), 0},
+      {reinterpret_cast<ElementOutput*>(params.D), params.ldd},
+      {alpha, beta},    // <- tuple of alpha and beta
+      split_k_slices};  // <- k-dimension split factor
 
   Gemm gemm_op;
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_noact_template.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_noact_template.h
index 31d2a21e2a..f7f2bbf396 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_noact_template.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_noact_template.h
@@ -20,9 +20,14 @@
 #include "cutlass/gemm/device/gemm_universal.h"
 #include "cutlass/gemm/device/gemm_splitk_parallel.h"
 
-template <typename InputType, typename OutType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_fuse_gemm_noact(GemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -64,9 +69,8 @@ bool dispatch_fuse_gemm_noact(GemmEpilogueAllParams params) {
   using ShapeMMAOp = MMAShape;  // <- MMA Op tile
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
-
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   // This code section describes how threadblocks are scheduled on GPU
   using SwizzleThreadBlock =
@@ -82,7 +86,7 @@ bool dispatch_fuse_gemm_noact(GemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   // Number of pipelines you want to use
   constexpr int NumStages = Stages;
@@ -164,10 +168,14 @@ bool dispatch_fuse_gemm_noact(GemmEpilogueAllParams params) {
   return true;
 }
 
-
-template <typename InputType, typename OutType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_fuse_gemm_split_k_noact(GemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -209,8 +217,8 @@ bool dispatch_fuse_gemm_split_k_noact(GemmEpilogueAllParams params) {
   using ShapeMMAOp = MMAShape;
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
       ElementOutput,  // <- data type of output matrix
@@ -222,50 +230,55 @@ bool dispatch_fuse_gemm_split_k_noact(GemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   // Number of pipelines you want to use
   constexpr int NumStages = Stages;
 
-    using ConvertScaledOp = cutlass::epilogue::thread::Convert<
-        ElementAccumulator,
-        cutlass::gemm::device::DefaultGemmConfiguration<cutlass::arch::OpClassSimt, SmArch, ElementInputA, ElementInputB,
-                                    ElementAccumulator,
-                                    ElementAccumulator>::EpilogueOutputOp::kCount,
-        ElementAccumulator>;
+  using ConvertScaledOp = cutlass::epilogue::thread::Convert<
+      ElementAccumulator,
+      cutlass::gemm::device::DefaultGemmConfiguration<
+          cutlass::arch::OpClassSimt,
+          SmArch,
+          ElementInputA,
+          ElementInputB,
+          ElementAccumulator,
+          ElementAccumulator>::EpilogueOutputOp::kCount,
+      ElementAccumulator>;
 
-    /// Reduction operator
-    using ReductionOp = cutlass::reduction::thread::ReduceAdd<
-        ElementAccumulator, typename EpilogueOp::ElementAccumulator,
-        EpilogueOp::kCount>;
+  /// Reduction operator
+  using ReductionOp = cutlass::reduction::thread::ReduceAdd<
+      ElementAccumulator,
+      typename EpilogueOp::ElementAccumulator,
+      EpilogueOp::kCount>;
 
-    /// Threadblock-level swizzling operator
-    using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmSplitKHorizontalThreadblockSwizzle;
+  /// Threadblock-level swizzling operator
+  using ThreadblockSwizzle =
+      cutlass::gemm::threadblock::GemmSplitKHorizontalThreadblockSwizzle;
 
-    /// Operation performed by GEMM
-    using Operator = cutlass::arch::OpMultiplyAddFastAccum;
-
-    using Gemm = cutlass::gemm::device::GemmSplitKParallel<ElementInputA,
-                                                        LayoutInputA,
-                                                        ElementInputB,
-                                                        LayoutInputB,
-                                                        ElementOutput,
-                                                        LayoutOutput,
-                                                        ElementAccumulator,
-                                                        MMAOp,
-                                                        SmArch,
-                                                        ShapeMMAThreadBlock,
-                                                        ShapeMMAWarp,
-                                                        ShapeMMAOp,
-                                                        EpilogueOp,
-                                                        ConvertScaledOp,
-                                                        ReductionOp,
-                                                        ThreadblockSwizzle,
-                                                        NumStages,
-                                                        kAlignmentA,
-                                                        kAlignmentB,
-                                                        Operator>;
+  /// Operation performed by GEMM
+  using Operator = cutlass::arch::OpMultiplyAddFastAccum;
 
+  using Gemm = cutlass::gemm::device::GemmSplitKParallel<ElementInputA,
+                                                         LayoutInputA,
+                                                         ElementInputB,
+                                                         LayoutInputB,
+                                                         ElementOutput,
+                                                         LayoutOutput,
+                                                         ElementAccumulator,
+                                                         MMAOp,
+                                                         SmArch,
+                                                         ShapeMMAThreadBlock,
+                                                         ShapeMMAWarp,
+                                                         ShapeMMAOp,
+                                                         EpilogueOp,
+                                                         ConvertScaledOp,
+                                                         ReductionOp,
+                                                         ThreadblockSwizzle,
+                                                         NumStages,
+                                                         kAlignmentA,
+                                                         kAlignmentB,
+                                                         Operator>;
 
   cutlass::gemm::GemmCoord problem_size =
       cutlass::gemm::GemmCoord{params.M, params.N, params.K};
@@ -276,15 +289,18 @@ bool dispatch_fuse_gemm_split_k_noact(GemmEpilogueAllParams params) {
   // Split K dimension into 16 partitions
   int split_k_slices = params.split_k;
 
-  // Create a tuple of gemm kernel arguments. This is later passed as arguments to launch
-  // instantiated CUTLASS kernel
-  typename Gemm::Arguments arguments{problem_size,  // <- problem size of matrix multiplication
-                                     {reinterpret_cast<ElementInputA*>(const_cast<void*>(params.A)),params.lda},
-                                     {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B)),params.ldb},
-                                     {reinterpret_cast<ElementOutput*>(const_cast<void*>(params.bias)),0},
-                                     {reinterpret_cast<ElementOutput*>(params.D),params.ldd},
-                                     {alpha, beta},          // <- tuple of alpha and beta
-                                     split_k_slices};        // <- k-dimension split factor
+  // Create a tuple of gemm kernel arguments. This is later passed as arguments
+  // to launch instantiated CUTLASS kernel
+  typename Gemm::Arguments arguments{
+      problem_size,  // <- problem size of matrix multiplication
+      {reinterpret_cast<ElementInputA*>(const_cast<void*>(params.A)),
+       params.lda},
+      {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B)),
+       params.ldb},
+      {reinterpret_cast<ElementOutput*>(const_cast<void*>(params.bias)), 0},
+      {reinterpret_cast<ElementOutput*>(params.D), params.ldd},
+      {alpha, beta},    // <- tuple of alpha and beta
+      split_k_slices};  // <- k-dimension split factor
 
   Gemm gemm_op;
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_relu_template.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_relu_template.h
index d2aa189eed..f59b84f59f 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_relu_template.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/fuse_gemm_relu_template.h
@@ -20,9 +20,14 @@
 #include "cutlass/gemm/device/gemm_universal.h"
 #include "cutlass/gemm/device/gemm_splitk_parallel.h"
 
-template <typename InputType, typename OutType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_fuse_gemm_relu(GemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -64,9 +69,8 @@ bool dispatch_fuse_gemm_relu(GemmEpilogueAllParams params) {
   using ShapeMMAOp = MMAShape;  // <- MMA Op tile
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
-
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   // This code section describes how threadblocks are scheduled on GPU
   using SwizzleThreadBlock =
@@ -82,7 +86,7 @@ bool dispatch_fuse_gemm_relu(GemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   // Number of pipelines you want to use
   constexpr int NumStages = Stages;
@@ -164,10 +168,14 @@ bool dispatch_fuse_gemm_relu(GemmEpilogueAllParams params) {
   return true;
 }
 
-
-template <typename InputType, typename OutType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_fuse_gemm_split_k_relu(GemmEpilogueAllParams params) {
   using ElementInputA = typename std::conditional_t<
       std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
@@ -209,8 +217,8 @@ bool dispatch_fuse_gemm_split_k_relu(GemmEpilogueAllParams params) {
   using ShapeMMAOp = MMAShape;
 
   static constexpr auto ScaleType =
-              hasbias? cutlass::epilogue::thread::ScaleType::NoBetaScaling
-                       : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
+      hasbias ? cutlass::epilogue::thread::ScaleType::NoBetaScaling
+              : cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling;
 
   using EpilogueOp = cutlass::epilogue::thread::LinearCombinationRelu<
       ElementOutput,  // <- data type of output matrix
@@ -222,50 +230,55 @@ bool dispatch_fuse_gemm_split_k_relu(GemmEpilogueAllParams params) {
       ElementAccumulator,  // <- data type of accumulator
       ElementComputeEpilogue,
       ScaleType>;  // <- data type for alpha/beta in linear
-                              // combination function
+                   // combination function
 
   // Number of pipelines you want to use
   constexpr int NumStages = Stages;
 
-    using ConvertScaledOp = cutlass::epilogue::thread::Convert<
-        ElementAccumulator,
-        cutlass::gemm::device::DefaultGemmConfiguration<cutlass::arch::OpClassSimt, SmArch, ElementInputA, ElementInputB,
-                                    ElementAccumulator,
-                                    ElementAccumulator>::EpilogueOutputOp::kCount,
-        ElementAccumulator>;
+  using ConvertScaledOp = cutlass::epilogue::thread::Convert<
+      ElementAccumulator,
+      cutlass::gemm::device::DefaultGemmConfiguration<
+          cutlass::arch::OpClassSimt,
+          SmArch,
+          ElementInputA,
+          ElementInputB,
+          ElementAccumulator,
+          ElementAccumulator>::EpilogueOutputOp::kCount,
+      ElementAccumulator>;
 
-    /// Reduction operator
-    using ReductionOp = cutlass::reduction::thread::ReduceAdd<
-        ElementAccumulator, typename EpilogueOp::ElementAccumulator,
-        EpilogueOp::kCount>;
+  /// Reduction operator
+  using ReductionOp = cutlass::reduction::thread::ReduceAdd<
+      ElementAccumulator,
+      typename EpilogueOp::ElementAccumulator,
+      EpilogueOp::kCount>;
 
-    /// Threadblock-level swizzling operator
-    using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmSplitKHorizontalThreadblockSwizzle;
+  /// Threadblock-level swizzling operator
+  using ThreadblockSwizzle =
+      cutlass::gemm::threadblock::GemmSplitKHorizontalThreadblockSwizzle;
 
-    /// Operation performed by GEMM
-    using Operator = cutlass::arch::OpMultiplyAddFastAccum;
-
-    using Gemm = cutlass::gemm::device::GemmSplitKParallel<ElementInputA,
-                                                        LayoutInputA,
-                                                        ElementInputB,
-                                                        LayoutInputB,
-                                                        ElementOutput,
-                                                        LayoutOutput,
-                                                        ElementAccumulator,
-                                                        MMAOp,
-                                                        SmArch,
-                                                        ShapeMMAThreadBlock,
-                                                        ShapeMMAWarp,
-                                                        ShapeMMAOp,
-                                                        EpilogueOp,
-                                                        ConvertScaledOp,
-                                                        ReductionOp,
-                                                        ThreadblockSwizzle,
-                                                        NumStages,
-                                                        kAlignmentA,
-                                                        kAlignmentB,
-                                                        Operator>;
+  /// Operation performed by GEMM
+  using Operator = cutlass::arch::OpMultiplyAddFastAccum;
 
+  using Gemm = cutlass::gemm::device::GemmSplitKParallel<ElementInputA,
+                                                         LayoutInputA,
+                                                         ElementInputB,
+                                                         LayoutInputB,
+                                                         ElementOutput,
+                                                         LayoutOutput,
+                                                         ElementAccumulator,
+                                                         MMAOp,
+                                                         SmArch,
+                                                         ShapeMMAThreadBlock,
+                                                         ShapeMMAWarp,
+                                                         ShapeMMAOp,
+                                                         EpilogueOp,
+                                                         ConvertScaledOp,
+                                                         ReductionOp,
+                                                         ThreadblockSwizzle,
+                                                         NumStages,
+                                                         kAlignmentA,
+                                                         kAlignmentB,
+                                                         Operator>;
 
   cutlass::gemm::GemmCoord problem_size =
       cutlass::gemm::GemmCoord{params.M, params.N, params.K};
@@ -276,15 +289,18 @@ bool dispatch_fuse_gemm_split_k_relu(GemmEpilogueAllParams params) {
   // Split K dimension into 16 partitions
   int split_k_slices = params.split_k;
 
-  // Create a tuple of gemm kernel arguments. This is later passed as arguments to launch
-  // instantiated CUTLASS kernel
-  typename Gemm::Arguments arguments{problem_size,  // <- problem size of matrix multiplication
-                                     {reinterpret_cast<ElementInputA*>(const_cast<void*>(params.A)),params.lda},
-                                     {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B)),params.ldb},
-                                     {reinterpret_cast<ElementOutput*>(const_cast<void*>(params.bias)),0},
-                                     {reinterpret_cast<ElementOutput*>(params.D),params.ldd},
-                                     {alpha, beta},          // <- tuple of alpha and beta
-                                     split_k_slices};        // <- k-dimension split factor
+  // Create a tuple of gemm kernel arguments. This is later passed as arguments
+  // to launch instantiated CUTLASS kernel
+  typename Gemm::Arguments arguments{
+      problem_size,  // <- problem size of matrix multiplication
+      {reinterpret_cast<ElementInputA*>(const_cast<void*>(params.A)),
+       params.lda},
+      {reinterpret_cast<ElementInputB*>(const_cast<void*>(params.B)),
+       params.ldb},
+      {reinterpret_cast<ElementOutput*>(const_cast<void*>(params.bias)), 0},
+      {reinterpret_cast<ElementOutput*>(params.D), params.ldd},
+      {alpha, beta},    // <- tuple of alpha and beta
+      split_k_slices};  // <- k-dimension split factor
 
   Gemm gemm_op;
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/per_channel_fp8_fp8_half_gemm.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/per_channel_fp8_fp8_half_gemm.h
index 3ea4d7c33e..82373f8693 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/per_channel_fp8_fp8_half_gemm.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/per_channel_fp8_fp8_half_gemm.h
@@ -17,10 +17,10 @@
 
 #include "fp8_common.h"
 
-#ifdef __GNUC__ // Check if the compiler is GCC or Clang
+#ifdef __GNUC__  // Check if the compiler is GCC or Clang
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wstrict-aliasing"
-#endif // __GNUC__
+#endif  // __GNUC__
 
 // clang-format off
 #include "cutlass/cutlass.h"
@@ -30,91 +30,134 @@
 #include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
 // clang-format on
 
-#ifdef __GNUC__ // Check if the compiler is GCC or Clang
+#ifdef __GNUC__  // Check if the compiler is GCC or Clang
 #pragma GCC diagnostic pop
-#endif          // __GNUC__
+#endif  // __GNUC__
 
-template <typename ElementType, typename OutElementType, typename AccumElementType, typename CtaShape,
-    typename WarpShape, typename InstructionShape, int Stages>
-struct DeviceGemmFp8RowwiseSm89
-{
-    using ElementInput = typename std::conditional_t<
-        std::is_same_v<ElementType, phi::dtype::float8_e4m3fn>,
-        cutlass::float_e4m3_t,
-        cutlass::float_e5m2_t>;
-    using ElementA = ElementInput;
-    using LayoutA = cutlass::layout::RowMajor;
-    static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
+template <typename ElementType,
+          typename OutElementType,
+          typename AccumElementType,
+          typename CtaShape,
+          typename WarpShape,
+          typename InstructionShape,
+          int Stages>
+struct DeviceGemmFp8RowwiseSm89 {
+  using ElementInput = typename std::conditional_t<
+      std::is_same_v<ElementType, phi::dtype::float8_e4m3fn>,
+      cutlass::float_e4m3_t,
+      cutlass::float_e5m2_t>;
+  using ElementA = ElementInput;
+  using LayoutA = cutlass::layout::RowMajor;
+  static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
 
-    using ElementB = ElementInput;
-    using LayoutB = cutlass::layout::ColumnMajor;
-    static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
+  using ElementB = ElementInput;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
 
+  using ElementOutput = typename std::conditional_t<
+      std::is_same_v<OutElementType, phi::dtype::bfloat16>,
+      cutlass::bfloat16_t,
+      cutlass::half_t>;
 
-    using ElementOutput =
-        typename std::conditional_t<std::is_same_v<OutElementType, phi::dtype::bfloat16>,
-                                    cutlass::bfloat16_t,
-                                    cutlass::half_t>;
+  using ElementC = ElementOutput;
+  using LayoutC = cutlass::layout::RowMajor;
+  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
 
-    using ElementC = ElementOutput;
-    using LayoutC = cutlass::layout::RowMajor;
-    static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+  using LayoutOutput = cutlass::layout::RowMajor;
+  static constexpr int AlignmentOutput =
+      128 / cutlass::sizeof_bits<ElementOutput>::value;
 
-    using LayoutOutput = cutlass::layout::RowMajor;
-    static constexpr int AlignmentOutput = 128 / cutlass::sizeof_bits<ElementOutput>::value;
+  using ElementAccumulator = AccumElementType;
+  using ElementComputeEpilogueScale = float;
+  using ArchTag = cutlass::arch::Sm89;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
 
-    using ElementAccumulator = AccumElementType;
-    using ElementComputeEpilogueScale = float;
-    using ArchTag = cutlass::arch::Sm89;
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
+  // Number of epilogue stages in EVT
+  static constexpr int EVTEpilogueStages = 1;
 
-    // Number of epilogue stages in EVT
-    static constexpr int EVTEpilogueStages = 1;
+  using OutputTileThreadMap =
+      cutlass::epilogue::threadblock::OutputTileThreadLayout<CtaShape,
+                                                             WarpShape,
+                                                             ElementOutput,
+                                                             AlignmentC,
+                                                             EVTEpilogueStages>;
 
-    using OutputTileThreadMap = cutlass::epilogue::threadblock::OutputTileThreadLayout<CtaShape, WarpShape, ElementOutput,
-        AlignmentC, EVTEpilogueStages>;
+  // Definition of EVT
+  using accSrc = cutlass::epilogue::threadblock::VisitorAccFetch;
 
-    // Definition of EVT
-    using accSrc = cutlass::epilogue::threadblock::VisitorAccFetch;
+  using ComputeBScale = cutlass::epilogue::threadblock::VisitorCompute<
+      cutlass::multiplies,
+      ElementComputeEpilogueScale,
+      ElementComputeEpilogueScale,
+      cutlass::FloatRoundStyle::round_to_nearest>;
+  using bScaleSrc = cutlass::epilogue::threadblock::VisitorRowBroadcast<
+      OutputTileThreadMap,
+      ElementComputeEpilogueScale,
+      cute::Stride<cute::_0, cute::_1, cute::_0>>;
+  using EpilogueBScale =
+      cutlass::epilogue::threadblock::Sm80EVT<ComputeBScale, accSrc, bScaleSrc>;
 
-    using ComputeBScale = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, ElementComputeEpilogueScale,
-        ElementComputeEpilogueScale, cutlass::FloatRoundStyle::round_to_nearest>;
-    using bScaleSrc = cutlass::epilogue::threadblock::VisitorRowBroadcast<OutputTileThreadMap, ElementComputeEpilogueScale,
-        cute::Stride<cute::_0, cute::_1, cute::_0>>;
-    using EpilogueBScale = cutlass::epilogue::threadblock::Sm80EVT<ComputeBScale, accSrc, bScaleSrc>;
+  using ComputeAScale = cutlass::epilogue::threadblock::VisitorCompute<
+      cutlass::multiplies,
+      ElementComputeEpilogueScale,
+      ElementComputeEpilogueScale,
+      cutlass::FloatRoundStyle::round_to_nearest>;
+  using aScaleSrc = cutlass::epilogue::threadblock::VisitorColBroadcast<
+      OutputTileThreadMap,
+      ElementComputeEpilogueScale,
+      cute::Stride<cute::_0, cute::_0, cute::_0>>;
+  using EpilogueAScale = cutlass::epilogue::threadblock::
+      Sm80EVT<ComputeAScale, EpilogueBScale, aScaleSrc>;
 
-    using ComputeAScale = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, ElementComputeEpilogueScale,
-        ElementComputeEpilogueScale, cutlass::FloatRoundStyle::round_to_nearest>;
-    using aScaleSrc = cutlass::epilogue::threadblock::VisitorColBroadcast<OutputTileThreadMap, ElementComputeEpilogueScale,
-        cute::Stride<cute::_0, cute::_0, cute::_0>>;
-    using EpilogueAScale = cutlass::epilogue::threadblock::Sm80EVT<ComputeAScale, EpilogueBScale, aScaleSrc>;
+  using Bias = cutlass::epilogue::threadblock::VisitorRowBroadcast<
+      OutputTileThreadMap,
+      ElementC,
+      cute::Stride<cute::_0, cute::_1, cute::_0>  // StrideMNL
+      >;
 
-    using Bias = cutlass::epilogue::threadblock::VisitorRowBroadcast<
-        OutputTileThreadMap, ElementC,
-        cute::Stride<cute::_0, cute::_1, cute::_0>  // StrideMNL
-    >;
+  using Compute0 = cutlass::epilogue::threadblock::VisitorCompute<
+      cutlass::plus,
+      ElementC,
+      ElementComputeEpilogueScale,
+      cutlass::FloatRoundStyle::round_to_nearest>;
 
-    using Compute0 = cutlass::epilogue::threadblock::VisitorCompute<
-        cutlass::plus, ElementC, ElementComputeEpilogueScale,
-        cutlass::FloatRoundStyle::round_to_nearest
-    >;
+  using EVTCompute0 =
+      cutlass::epilogue::threadblock::Sm80EVT<Compute0, EpilogueAScale, Bias>;
 
-    using EVTCompute0 = cutlass::epilogue::threadblock::Sm80EVT<
-        Compute0,
-        EpilogueAScale,
-        Bias>;
+  using dTar = cutlass::epilogue::threadblock::VisitorAuxStore<
+      OutputTileThreadMap,
+      ElementOutput,
+      cutlass::FloatRoundStyle::round_to_nearest,
+      cute::Stride<int64_t, cute::_1, cute::_0>>;
+  using EpilogueStore =
+      cutlass::epilogue::threadblock::Sm80EVT<dTar, EVTCompute0>;
 
-    using dTar = cutlass::epilogue::threadblock::VisitorAuxStore<OutputTileThreadMap, ElementOutput,
-        cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1,cute:: _0>>;
-    using EpilogueStore = cutlass::epilogue::threadblock::Sm80EVT<dTar, EVTCompute0>;
+  using EpilogueOp = EpilogueStore;
 
-    using EpilogueOp = EpilogueStore;
+  using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+      ElementA,
+      LayoutA,
+      cutlass::ComplexTransform::kNone,
+      AlignmentA,
+      ElementB,
+      LayoutB,
+      cutlass::ComplexTransform::kNone,
+      AlignmentB,
+      ElementC,
+      LayoutC,
+      AlignmentC,
+      ElementAccumulator,
+      ElementComputeEpilogueScale,
+      OperatorClass,
+      ArchTag,
+      CtaShape,
+      WarpShape,
+      InstructionShape,
+      EpilogueOp,
+      cutlass::gemm::threadblock::ThreadblockSwizzleStreamK,
+      Stages,
+      cutlass::arch::OpMultiplyAddFastAccum,
+      EVTEpilogueStages>::GemmKernel;
 
-    using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<ElementA, LayoutA,
-        cutlass::ComplexTransform::kNone, AlignmentA, ElementB, LayoutB, cutlass::ComplexTransform::kNone, AlignmentB,
-        ElementC, LayoutC, AlignmentC, ElementAccumulator, ElementComputeEpilogueScale, OperatorClass, ArchTag, CtaShape,
-        WarpShape, InstructionShape, EpilogueOp, cutlass::gemm::threadblock::ThreadblockSwizzleStreamK, Stages,
-        cutlass::arch::OpMultiplyAddFastAccum, EVTEpilogueStages>::GemmKernel;
-
-    using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
 };
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/visitor_fp8_gemm_fused_template.h b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/visitor_fp8_gemm_fused_template.h
index a073c9bff3..111b2cfd15 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/visitor_fp8_gemm_fused_template.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fp8_gemm_fused/visitor_fp8_gemm_fused_template.h
@@ -16,56 +16,68 @@
 #include "per_channel_fp8_fp8_half_gemm.h"  // NOLINT
 
 template <typename Gemm>
-typename Gemm::Arguments prepar_gemm_args_sm89(void* D, void const* A, void const* B, void const* C_bias,
-    int m, int n, int k, float const* scale_d0, float const* scale_d1)
-{
-    using ElementT = typename Gemm::ElementA;
-    using ElementOutput = typename Gemm::ElementD;
-    using ElementComputeEpilogue = float;
+typename Gemm::Arguments prepar_gemm_args_sm89(void* D,
+                                               void const* A,
+                                               void const* B,
+                                               void const* C_bias,
+                                               int m,
+                                               int n,
+                                               int k,
+                                               float const* scale_d0,
+                                               float const* scale_d1) {
+  using ElementT = typename Gemm::ElementA;
+  using ElementOutput = typename Gemm::ElementD;
+  using ElementComputeEpilogue = float;
 
-    int const lda = k;
-    int const ldb = k;
-    int const ldc = n;
+  int const lda = k;
+  int const ldb = k;
+  int const ldc = n;
 
-    typename Gemm::Arguments args(cutlass::gemm::GemmUniversalMode::kGemm, // Mode
-        {m, n, k},                                                         // Problem size
-        1,                                                                 // Split-k factor
-        {},                                                                // Epilogue args
-        reinterpret_cast<ElementT const*>(A),                              // a pointer
-        reinterpret_cast<ElementT const*>(B),                              // b pointer
-        nullptr,                                                           // c pointer (unused)
-        nullptr,                                                           // d pointer (unused)
-        m * k,                                                             // batch stride a (unused)
-        n * k,                                                             // batch stride b (unused)
-        m * n,                                                             // batch stride c (unused)
-        m * n,                                                             // batch stride d (unused)
-        lda,                                                               // stride a
-        ldb,                                                               // stride b
-        ldc,                                                               // stride c (unused)
-        ldc);                                                              // stride d (unused)
+  typename Gemm::Arguments args(
+      cutlass::gemm::GemmUniversalMode::kGemm,  // Mode
+      {m, n, k},                                // Problem size
+      1,                                        // Split-k factor
+      {},                                       // Epilogue args
+      reinterpret_cast<ElementT const*>(A),     // a pointer
+      reinterpret_cast<ElementT const*>(B),     // b pointer
+      nullptr,                                  // c pointer (unused)
+      nullptr,                                  // d pointer (unused)
+      m * k,                                    // batch stride a (unused)
+      n * k,                                    // batch stride b (unused)
+      m * n,                                    // batch stride c (unused)
+      m * n,                                    // batch stride d (unused)
+      lda,                                      // stride a
+      ldb,                                      // stride b
+      ldc,                                      // stride c (unused)
+      ldc);                                     // stride d (unused)
 
-    args.epilogue = {
-        {
-            {
-                {
-                    {},                                      // Accumulator
-                    {reinterpret_cast<ElementComputeEpilogue const*>(scale_d1), ElementComputeEpilogue(0),
-                        {cute::_0{}, cute::_1{}, cute::_0{}}},
-                    {} // Multiplies
-                },
-                {reinterpret_cast<ElementComputeEpilogue const*>(scale_d0), ElementComputeEpilogue(0), {cute::_0{}, cute::_0{}, cute::_0{}}},
-                {} // Multiplies
-            },                                                                                                          // Accum
-            {reinterpret_cast<ElementOutput const*>(C_bias), ElementOutput(0), {cute::_0{}, cute::_1{}, cute::_0{}}},                 // Bias
-            {}                                                                                             // Compute0
-        },
-        {reinterpret_cast<ElementOutput*>(D), {n, cute::_1{}, cute::_0{}}}
-    };
-    return args;
+  args.epilogue = {
+      {
+          {
+              {
+                  {},  // Accumulator
+                  {reinterpret_cast<ElementComputeEpilogue const*>(scale_d1),
+                   ElementComputeEpilogue(0),
+                   {cute::_0{}, cute::_1{}, cute::_0{}}},
+                  {}  // Multiplies
+              },
+              {reinterpret_cast<ElementComputeEpilogue const*>(scale_d0),
+               ElementComputeEpilogue(0),
+               {cute::_0{}, cute::_0{}, cute::_0{}}},
+              {}  // Multiplies
+          },      // Accum
+          {reinterpret_cast<ElementOutput const*>(C_bias),
+           ElementOutput(0),
+           {cute::_0{}, cute::_1{}, cute::_0{}}},  // Bias
+          {}                                       // Compute0
+      },
+      {reinterpret_cast<ElementOutput*>(D), {n, cute::_1{}, cute::_0{}}}};
+  return args;
 }
 
 template <typename Gemm>
-bool per_channel_fp8_fp8_gemm_scale_bias(GemmEpilogueAllParams params, typename Gemm::Arguments args) {
+bool per_channel_fp8_fp8_gemm_scale_bias(GemmEpilogueAllParams params,
+                                         typename Gemm::Arguments args) {
   Gemm per_channel_fp8_gemm;
 
   cutlass::Status status = per_channel_fp8_gemm.can_implement(args);
@@ -89,14 +101,31 @@ bool per_channel_fp8_fp8_gemm_scale_bias(GemmEpilogueAllParams params, typename
   return true;
 }
 
-
-template <typename InputType, typename OutType,
-            typename ThreadBlockShape, typename WarpShape,
-            typename MMAShape, int Stages, bool hasbias, typename SM>
+template <typename InputType,
+          typename OutType,
+          typename ThreadBlockShape,
+          typename WarpShape,
+          typename MMAShape,
+          int Stages,
+          bool hasbias,
+          typename SM>
 bool dispatch_visitor_fuse_gemm(GemmEpilogueAllParams params) {
-    using AccumElementType = float;
-    using Gemm = typename DeviceGemmFp8RowwiseSm89<InputType, OutType, AccumElementType, ThreadBlockShape, WarpShape, MMAShape,
-        Stages>::Gemm;
-    auto args = prepar_gemm_args_sm89<Gemm>(params.D, params.A, params.B, params.bias, params.M, params.N, params.K, params.scalar_scale, params.channel_scale);
-    per_channel_fp8_fp8_gemm_scale_bias<Gemm>(params, args);
+  using AccumElementType = float;
+  using Gemm = typename DeviceGemmFp8RowwiseSm89<InputType,
+                                                 OutType,
+                                                 AccumElementType,
+                                                 ThreadBlockShape,
+                                                 WarpShape,
+                                                 MMAShape,
+                                                 Stages>::Gemm;
+  auto args = prepar_gemm_args_sm89<Gemm>(params.D,
+                                          params.A,
+                                          params.B,
+                                          params.bias,
+                                          params.M,
+                                          params.N,
+                                          params.K,
+                                          params.scalar_scale,
+                                          params.channel_scale);
+  per_channel_fp8_fp8_gemm_scale_bias<Gemm>(params, args);
 }
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h b/custom_ops/gpu_ops/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h
index 6b1ab209e3..44d5ccc220 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h
@@ -21,10 +21,8 @@
 #include <cuda_runtime_api.h>
 #include <vector>
 
-namespace kernels
-{
-namespace cutlass_kernels
-{
+namespace kernels {
+namespace cutlass_kernels {
 
 /*
   This runner only supports:
@@ -32,64 +30,105 @@ namespace cutlass_kernels
 
   Activations, biases, scales and outputs are all assumed to be row-major.
 
-  However, it is assumed that B is in a special format governed by cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.
-  In this case, B must be preprocessed using the cutlass weight only quant preprocessors. The weight preprocessor
-  will instantiate the layout and preprocess based on the instantiation, so layout changes should only require
-  modifications to mix_gemm_B_layout.h.
+  However, it is assumed that B is in a special format governed by
+  cutlass_extensions/gemm/kernel/mixed_gemm_B_layout. In this case, B must be
+  preprocessed using the cutlass weight only quant preprocessors. The weight
+  preprocessor will instantiate the layout and preprocess based on the
+  instantiation, so layout changes should only require modifications to
+  mix_gemm_B_layout.h.
 */
 
-class CutlassFpAIntBGemmRunnerInterface
-{
-public:
-    CutlassFpAIntBGemmRunnerInterface() {}
+class CutlassFpAIntBGemmRunnerInterface {
+ public:
+  CutlassFpAIntBGemmRunnerInterface() {}
 
-    virtual ~CutlassFpAIntBGemmRunnerInterface() {}
+  virtual ~CutlassFpAIntBGemmRunnerInterface() {}
 
-    virtual void gemm(void const* A, void const* B, void const* weight_scales, void const* weight_zero_points,
-        void const* biases, float const alpha, void* C, int m, int n, int k, int const group_size,
-        cutlass_extensions::CutlassGemmConfig gemmConfig, void* workspace_ptr, const size_t workspace_bytes, cudaStream_t stream
-        ) = 0;
+  virtual void gemm(void const* A,
+                    void const* B,
+                    void const* weight_scales,
+                    void const* weight_zero_points,
+                    void const* biases,
+                    float const alpha,
+                    void* C,
+                    int m,
+                    int n,
+                    int k,
+                    int const group_size,
+                    cutlass_extensions::CutlassGemmConfig gemmConfig,
+                    void* workspace_ptr,
+                    const size_t workspace_bytes,
+                    cudaStream_t stream) = 0;
 
-    // Returns desired workspace size in bytes.
-    virtual size_t getWorkspaceSize(int const m, int const n, int const k) = 0;
+  // Returns desired workspace size in bytes.
+  virtual size_t getWorkspaceSize(int const m, int const n, int const k) = 0;
 
-    virtual std::vector<cutlass_extensions::CutlassGemmConfig> getConfigs(int k) const = 0;
+  virtual std::vector<cutlass_extensions::CutlassGemmConfig> getConfigs(
+      int k) const = 0;
 
-protected:
-    static constexpr int SPLIT_K_LIMIT = 7;
-    static constexpr int MIN_M_TILE = 16;
-    static constexpr int MIN_N_TILE = 64;
+ protected:
+  static constexpr int SPLIT_K_LIMIT = 7;
+  static constexpr int MIN_M_TILE = 16;
+  static constexpr int MIN_N_TILE = 64;
 };
 
-template <typename ActivationType, typename WeightType, cutlass::WeightOnlyQuantOp QuantOp,
-    typename ScaleZeroType = ActivationType, typename BiasType = ActivationType, typename OutputType = ActivationType>
-class CutlassFpAIntBGemmRunner : public virtual CutlassFpAIntBGemmRunnerInterface
-{
-public:
-    CutlassFpAIntBGemmRunner();
-    ~CutlassFpAIntBGemmRunner();
+template <typename ActivationType,
+          typename WeightType,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename ScaleZeroType = ActivationType,
+          typename BiasType = ActivationType,
+          typename OutputType = ActivationType>
+class CutlassFpAIntBGemmRunner
+    : public virtual CutlassFpAIntBGemmRunnerInterface {
+ public:
+  CutlassFpAIntBGemmRunner();
+  ~CutlassFpAIntBGemmRunner();
 
-    void gemm(void const* A, void const* B, void const* weight_scales, void const* weight_zero_points,
-        void const* biases, float const alpha, void* C, int m, int n, int k, int const group_size,
-        cutlass_extensions::CutlassGemmConfig gemmConfig, void* workspace_ptr, const size_t workspace_bytes,
-        cudaStream_t stream) override;
+  void gemm(void const* A,
+            void const* B,
+            void const* weight_scales,
+            void const* weight_zero_points,
+            void const* biases,
+            float const alpha,
+            void* C,
+            int m,
+            int n,
+            int k,
+            int const group_size,
+            cutlass_extensions::CutlassGemmConfig gemmConfig,
+            void* workspace_ptr,
+            const size_t workspace_bytes,
+            cudaStream_t stream) override;
 
-    // Returns desired workspace size in bytes.
-    size_t getWorkspaceSize(int const m, int const n, int const k) override;
+  // Returns desired workspace size in bytes.
+  size_t getWorkspaceSize(int const m, int const n, int const k) override;
 
-    std::vector<cutlass_extensions::CutlassGemmConfig> getConfigs(int k) const override;
+  std::vector<cutlass_extensions::CutlassGemmConfig> getConfigs(
+      int k) const override;
 
-private:
-    template <typename EpilogueTag>
-    void dispatch_to_arch(ActivationType const* A, WeightType const* B, ScaleZeroType const* weight_scales,
-        ScaleZeroType const* weight_zero_points, BiasType const* biases, float const alpha, OutputType* C, int m, int n,
-        int k, int const group_size, cutlass_extensions::CutlassGemmConfig gemm_config, void* workspace_ptr,
-        const size_t workspace_bytes, cudaStream_t stream, int* occupancy = nullptr);
+ private:
+  template <typename EpilogueTag>
+  void dispatch_to_arch(ActivationType const* A,
+                        WeightType const* B,
+                        ScaleZeroType const* weight_scales,
+                        ScaleZeroType const* weight_zero_points,
+                        BiasType const* biases,
+                        float const alpha,
+                        OutputType* C,
+                        int m,
+                        int n,
+                        int k,
+                        int const group_size,
+                        cutlass_extensions::CutlassGemmConfig gemm_config,
+                        void* workspace_ptr,
+                        const size_t workspace_bytes,
+                        cudaStream_t stream,
+                        int* occupancy = nullptr);
 
-private:
-    int sm_;
-    int multi_processor_count_;
+ private:
+  int sm_;
+  int multi_processor_count_;
 };
 
-} // namespace cutlass_kernels
-} // namespace kernels
+}  // namespace cutlass_kernels
+}  // namespace kernels
diff --git a/custom_ops/gpu_ops/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h b/custom_ops/gpu_ops/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h
index 2e02658d26..bc62bb34bf 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h
@@ -43,311 +43,559 @@
 namespace kernels {
 namespace cutlass_kernels {
 
-template <
-        typename ActivationType,
-        typename WeightType,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType,
-        typename arch,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename EpilogueTag,
-        typename ThreadblockShape,
-        typename WarpShape,
-        int Stages>
+template <typename ActivationType,
+          typename WeightType,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType,
+          typename arch,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape,
+          int Stages>
 void generic_mixed_gemm_kernelLauncher(
-        ActivationType const* A,
-        WeightType const* B,
-        ScaleZeroType const* weight_scales,
-        ScaleZeroType const* weight_zero_points,
-        BiasType const* biases,
-        float const alpha,
-        OutputType* C,
-        int m,
-        int n,
-        int k,
-        int const group_size,
-        cutlass_extensions::CutlassGemmConfig gemm_config,
-        void* workspace,
-        size_t workspace_bytes,
-        cudaStream_t stream,
-        int* occupancy = nullptr) {
-    // The cutlass type for the input elements. This is needed to convert to cutlass::half_t if
-    // necessary.
-    using CutlassActivationType = typename CudaToCutlassTypeAdapter<ActivationType>::type;
-    using CutlassWeightType = typename CudaToCutlassTypeAdapter<WeightType>::type;
-    using CutlassScaleZeroType = typename CudaToCutlassTypeAdapter<ScaleZeroType>::type;
-    using CutlassBiasType = typename CudaToCutlassTypeAdapter<BiasType>::type;
-    using CutlassOutputType = typename CudaToCutlassTypeAdapter<OutputType>::type;
+    ActivationType const* A,
+    WeightType const* B,
+    ScaleZeroType const* weight_scales,
+    ScaleZeroType const* weight_zero_points,
+    BiasType const* biases,
+    float const alpha,
+    OutputType* C,
+    int m,
+    int n,
+    int k,
+    int const group_size,
+    cutlass_extensions::CutlassGemmConfig gemm_config,
+    void* workspace,
+    size_t workspace_bytes,
+    cudaStream_t stream,
+    int* occupancy = nullptr) {
+  // The cutlass type for the input elements. This is needed to convert to
+  // cutlass::half_t if necessary.
+  using CutlassActivationType =
+      typename CudaToCutlassTypeAdapter<ActivationType>::type;
+  using CutlassWeightType = typename CudaToCutlassTypeAdapter<WeightType>::type;
+  using CutlassScaleZeroType =
+      typename CudaToCutlassTypeAdapter<ScaleZeroType>::type;
+  using CutlassBiasType = typename CudaToCutlassTypeAdapter<BiasType>::type;
+  using CutlassOutputType = typename CudaToCutlassTypeAdapter<OutputType>::type;
 
-    // We need separate config for each architecture since we will target different tensorcore
-    // instructions. For float, we do not target TCs.
-    using MixedGemmArchTraits = cutlass::gemm::kernel::
-            MixedGemmArchTraits<CutlassActivationType, CutlassWeightType, arch>;
-    using ElementAccumulator = typename MixedGemmArchTraits::AccType;
+  // We need separate config for each architecture since we will target
+  // different tensorcore instructions. For float, we do not target TCs.
+  using MixedGemmArchTraits = cutlass::gemm::kernel::
+      MixedGemmArchTraits<CutlassActivationType, CutlassWeightType, arch>;
+  using ElementAccumulator = typename MixedGemmArchTraits::AccType;
 
-    constexpr int ElementsPerAccessC = 128 / cutlass::sizeof_bits<CutlassOutputType>::value;
-    using EpilogueOp = typename cutlass_extensions::
-            Epilogue<CutlassOutputType, ElementsPerAccessC, ElementAccumulator, EpilogueTag>::Op;
+  constexpr int ElementsPerAccessC =
+      128 / cutlass::sizeof_bits<CutlassOutputType>::value;
+  using EpilogueOp = typename cutlass_extensions::Epilogue<CutlassOutputType,
+                                                           ElementsPerAccessC,
+                                                           ElementAccumulator,
+                                                           EpilogueTag>::Op;
 
-    using Operator = typename MixedGemmArchTraits::Operator;
-    using TaggedOperator = typename cutlass::arch::TagOperator<Operator, QuantOp>::TaggedOperator;
+  using Operator = typename MixedGemmArchTraits::Operator;
+  using TaggedOperator =
+      typename cutlass::arch::TagOperator<Operator, QuantOp>::TaggedOperator;
 
-    using GemmKernel_ = typename cutlass::gemm::kernel::DefaultGemm<
-            CutlassActivationType,
-            cutlass::layout::RowMajor,
-            MixedGemmArchTraits::ElementsPerAccessA,
-            CutlassWeightType,
-            typename MixedGemmArchTraits::LayoutB,
-            MixedGemmArchTraits::ElementsPerAccessB,
-            CutlassOutputType,
-            cutlass::layout::RowMajor,
-            ElementAccumulator,
-            cutlass::arch::OpClassTensorOp,
-            arch,
-            ThreadblockShape,
-            WarpShape,
-            typename MixedGemmArchTraits::InstructionShape,
-            EpilogueOp,
-            typename cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
-            Stages,
-            true,
-            TaggedOperator>::GemmKernel;
+  using GemmKernel_ = typename cutlass::gemm::kernel::DefaultGemm<
+      CutlassActivationType,
+      cutlass::layout::RowMajor,
+      MixedGemmArchTraits::ElementsPerAccessA,
+      CutlassWeightType,
+      typename MixedGemmArchTraits::LayoutB,
+      MixedGemmArchTraits::ElementsPerAccessB,
+      CutlassOutputType,
+      cutlass::layout::RowMajor,
+      ElementAccumulator,
+      cutlass::arch::OpClassTensorOp,
+      arch,
+      ThreadblockShape,
+      WarpShape,
+      typename MixedGemmArchTraits::InstructionShape,
+      EpilogueOp,
+      typename cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+      Stages,
+      true,
+      TaggedOperator>::GemmKernel;
 
-    using GemmKernel = cutlass::gemm::kernel::GemmFpAIntB<
-            typename GemmKernel_::Mma,
-            typename GemmKernel_::Epilogue,
-            typename GemmKernel_::ThreadblockSwizzle,
-            arch,  // Ensure top level arch is used for dispatch
-            GemmKernel_::kSplitKSerial>;
+  using GemmKernel = cutlass::gemm::kernel::GemmFpAIntB<
+      typename GemmKernel_::Mma,
+      typename GemmKernel_::Epilogue,
+      typename GemmKernel_::ThreadblockSwizzle,
+      arch,  // Ensure top level arch is used for dispatch
+      GemmKernel_::kSplitKSerial>;
 
-    if (occupancy != nullptr) {
-        *occupancy = cutlass_extensions::compute_occupancy_for_kernel<GemmKernel>();
-        return;
-    }
+  if (occupancy != nullptr) {
+    *occupancy = cutlass_extensions::compute_occupancy_for_kernel<GemmKernel>();
+    return;
+  }
 
-    using Gemm = cutlass::gemm::device::GemmUniversalBaseCompat<GemmKernel>;
+  using Gemm = cutlass::gemm::device::GemmUniversalBaseCompat<GemmKernel>;
 
-    int const ldb =
-            cutlass::platform::
-                    is_same<cutlass::layout::RowMajor, typename MixedGemmArchTraits::LayoutB>::value
-            ? n
-            : k * GemmKernel::kInterleave;
+  int const ldb =
+      cutlass::platform::is_same<cutlass::layout::RowMajor,
+                                 typename MixedGemmArchTraits::LayoutB>::value
+          ? n
+          : k * GemmKernel::kInterleave;
 
-    if (weight_scales == nullptr) {
-        throw std::runtime_error("Weight scales must always be set to a non-null value.");
-    }
+  if (weight_scales == nullptr) {
+    throw std::runtime_error(
+        "Weight scales must always be set to a non-null value.");
+  }
 
-    if constexpr (cutlass::isFinegrained(QuantOp)) {
-        if constexpr (cutlass::platform::is_same<CutlassActivationType, cutlass::float_e4m3_t>::
-                              value) {
-            if (group_size != 128) {
-                throw std::runtime_error(
-                        "Only group size 128 supported for fine grained W4A(fp)8 kernels.");
-            }
-        }
-        if (group_size != 64 && group_size != 128) {
-            throw std::runtime_error(
-                    "Only group size 64 and 128 supported for fine grained kernels.");
-        }
-
-        if constexpr (QuantOp == cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY) {
-            if (weight_zero_points != nullptr) {
-                throw std::runtime_error(
-                        "Weight zero pointer must be a nullptr for scale only fine grained");
-            }
-        } else if constexpr (QuantOp == cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS) {
-            if (weight_zero_points == nullptr) {
-                throw std::runtime_error(
-                        "Weight zero pointer must be valid for scale and bias fine grained");
-            }
-        }
-    } else {
-        if (group_size != k) {
-            throw std::runtime_error("Invalid group size for per column scaling kernels.");
-        }
-
-        if (weight_zero_points != nullptr) {
-            throw std::runtime_error(
-                    "Weight zero-points must be null when running per column scaling");
-        }
-    }
-
-    int const ld_scale_zero = cutlass::isFinegrained(QuantOp) ? n : 0;
-    ElementAccumulator output_op_beta =
-            (biases == nullptr) ? ElementAccumulator(0.f) : ElementAccumulator(1.f);
-    typename Gemm::Arguments args(
-            {m, n, k},
-            group_size,
-            {reinterpret_cast<CutlassActivationType*>(const_cast<ActivationType*>(A)), k},
-            {reinterpret_cast<CutlassWeightType*>(const_cast<WeightType*>(B)), ldb},
-            {reinterpret_cast<CutlassScaleZeroType*>(const_cast<ScaleZeroType*>(weight_scales)),
-             ld_scale_zero},
-            {reinterpret_cast<CutlassScaleZeroType*>(
-                     const_cast<ScaleZeroType*>(weight_zero_points)),
-             ld_scale_zero},
-            {reinterpret_cast<CutlassBiasType*>(const_cast<BiasType*>(biases)), 0},
-            {reinterpret_cast<CutlassOutputType*>(C), n},
-            gemm_config.split_k_factor,
-            {ElementAccumulator(alpha), output_op_beta});
-
-    // This assertion is enabled because because for the column interleaved layout, K MUST be a
-    // multiple of threadblockK. The reason for this is that the default pitchlinear iterators are
-    // used to handle walking over the interleaved matrix. The way masking in handled in these do
-    // not map to the interleaved layout. We need to write our own predicated iterator in order to
-    // relax this limitation.
-    if (GemmKernel::kInterleave > 1 &&
-        ((k % MixedGemmArchTraits::ThreadblockK) ||
-         ((k / gemm_config.split_k_factor) % MixedGemmArchTraits::ThreadblockK))) {
+  if constexpr (cutlass::isFinegrained(QuantOp)) {
+    if constexpr (cutlass::platform::is_same<CutlassActivationType,
+                                             cutlass::float_e4m3_t>::value) {
+      if (group_size != 128) {
         throw std::runtime_error(
-                "Assertion: k[" + std::to_string(k) + "] must be multiple of threadblockK[" +
-                std::to_string(MixedGemmArchTraits::ThreadblockK) + "]");
+            "Only group size 128 supported for fine grained W4A(fp)8 kernels.");
+      }
+    }
+    if (group_size != 64 && group_size != 128) {
+      throw std::runtime_error(
+          "Only group size 64 and 128 supported for fine grained kernels.");
     }
 
-    Gemm gemm;
-
-    if (gemm.get_workspace_size(args) > workspace_bytes) {
-        std::cerr << "Requested split-k but workspace size insufficient. Falling back to "
-                     "non-split-k implementation."
-                  << std::endl;
-        // If requested split-k factor will require more workspace bytes, revert to standard gemm.
-        args.batch_count = 1;
+    if constexpr (QuantOp ==
+                  cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY) {
+      if (weight_zero_points != nullptr) {
+        throw std::runtime_error(
+            "Weight zero pointer must be a nullptr for scale only fine "
+            "grained");
+      }
+    } else if constexpr (QuantOp == cutlass::WeightOnlyQuantOp::
+                                        FINEGRAINED_SCALE_AND_ZEROS) {
+      if (weight_zero_points == nullptr) {
+        throw std::runtime_error(
+            "Weight zero pointer must be valid for scale and bias fine "
+            "grained");
+      }
+    }
+  } else {
+    if (group_size != k) {
+      throw std::runtime_error(
+          "Invalid group size for per column scaling kernels.");
     }
 
-    auto can_implement = gemm.can_implement(args);
-    if (can_implement != cutlass::Status::kSuccess) {
-        std::string err_msg = "fp8_int4 cutlass kernel will fail for params. Error: " +
-                std::string(cutlassGetStatusString(can_implement));
-        throw std::runtime_error("[fp8_int4 Runner] " + err_msg);
+    if (weight_zero_points != nullptr) {
+      throw std::runtime_error(
+          "Weight zero-points must be null when running per column scaling");
     }
+  }
 
-    auto init_status = gemm.initialize(args, workspace, stream);
-    if (init_status != cutlass::Status::kSuccess) {
-        std::string err_msg = "Failed to initialize cutlass fp8_int4 gemm. Error: " +
-                std::string(cutlassGetStatusString(init_status));
-        throw std::runtime_error("[fp8_int4 Runner] " + err_msg);
-    }
+  int const ld_scale_zero = cutlass::isFinegrained(QuantOp) ? n : 0;
+  ElementAccumulator output_op_beta =
+      (biases == nullptr) ? ElementAccumulator(0.f) : ElementAccumulator(1.f);
+  typename Gemm::Arguments args(
+      {m, n, k},
+      group_size,
+      {reinterpret_cast<CutlassActivationType*>(const_cast<ActivationType*>(A)),
+       k},
+      {reinterpret_cast<CutlassWeightType*>(const_cast<WeightType*>(B)), ldb},
+      {reinterpret_cast<CutlassScaleZeroType*>(
+           const_cast<ScaleZeroType*>(weight_scales)),
+       ld_scale_zero},
+      {reinterpret_cast<CutlassScaleZeroType*>(
+           const_cast<ScaleZeroType*>(weight_zero_points)),
+       ld_scale_zero},
+      {reinterpret_cast<CutlassBiasType*>(const_cast<BiasType*>(biases)), 0},
+      {reinterpret_cast<CutlassOutputType*>(C), n},
+      gemm_config.split_k_factor,
+      {ElementAccumulator(alpha), output_op_beta});
 
-    auto run_status = gemm.run(stream);
-    if (run_status != cutlass::Status::kSuccess) {
-        std::string err_msg = "Failed to run cutlass fp8_int4 gemm. Error: " +
-                std::string(cutlassGetStatusString(run_status));
-        throw std::runtime_error("[fp8_int4 Runner] " + err_msg);
-    }
+  // This assertion is enabled because because for the column interleaved
+  // layout, K MUST be a multiple of threadblockK. The reason for this is that
+  // the default pitchlinear iterators are used to handle walking over the
+  // interleaved matrix. The way masking in handled in these do not map to the
+  // interleaved layout. We need to write our own predicated iterator in order
+  // to relax this limitation.
+  if (GemmKernel::kInterleave > 1 && ((k % MixedGemmArchTraits::ThreadblockK) ||
+                                      ((k / gemm_config.split_k_factor) %
+                                       MixedGemmArchTraits::ThreadblockK))) {
+    throw std::runtime_error("Assertion: k[" + std::to_string(k) +
+                             "] must be multiple of threadblockK[" +
+                             std::to_string(MixedGemmArchTraits::ThreadblockK) +
+                             "]");
+  }
+
+  Gemm gemm;
+
+  if (gemm.get_workspace_size(args) > workspace_bytes) {
+    std::cerr
+        << "Requested split-k but workspace size insufficient. Falling back to "
+           "non-split-k implementation."
+        << std::endl;
+    // If requested split-k factor will require more workspace bytes, revert to
+    // standard gemm.
+    args.batch_count = 1;
+  }
+
+  auto can_implement = gemm.can_implement(args);
+  if (can_implement != cutlass::Status::kSuccess) {
+    std::string err_msg =
+        "fp8_int4 cutlass kernel will fail for params. Error: " +
+        std::string(cutlassGetStatusString(can_implement));
+    throw std::runtime_error("[fp8_int4 Runner] " + err_msg);
+  }
+
+  auto init_status = gemm.initialize(args, workspace, stream);
+  if (init_status != cutlass::Status::kSuccess) {
+    std::string err_msg =
+        "Failed to initialize cutlass fp8_int4 gemm. Error: " +
+        std::string(cutlassGetStatusString(init_status));
+    throw std::runtime_error("[fp8_int4 Runner] " + err_msg);
+  }
+
+  auto run_status = gemm.run(stream);
+  if (run_status != cutlass::Status::kSuccess) {
+    std::string err_msg = "Failed to run cutlass fp8_int4 gemm. Error: " +
+                          std::string(cutlassGetStatusString(run_status));
+    throw std::runtime_error("[fp8_int4 Runner] " + err_msg);
+  }
 }
 
-template <
-        typename ActivationType,
-        typename WeightType,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType,
-        typename arch,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename EpilogueTag,
-        typename ThreadblockShape,
-        typename WarpShape>
-void dispatch_gemm_config(
-        ActivationType const* A,
-        WeightType const* B,
-        ScaleZeroType const* weight_scales,
-        ScaleZeroType const* weight_zero_points,
-        BiasType const* biases,
-        float const alpha,
-        OutputType* C,
-        int m,
-        int n,
-        int k,
-        int const group_size,
-        cutlass_extensions::CutlassGemmConfig gemm_config,
-        void* workspace,
-        size_t workspace_bytes,
-        cudaStream_t stream,
-        int* occupancy = nullptr) {
-    switch (gemm_config.stages) {
+template <typename ActivationType,
+          typename WeightType,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType,
+          typename arch,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape>
+void dispatch_gemm_config(ActivationType const* A,
+                          WeightType const* B,
+                          ScaleZeroType const* weight_scales,
+                          ScaleZeroType const* weight_zero_points,
+                          BiasType const* biases,
+                          float const alpha,
+                          OutputType* C,
+                          int m,
+                          int n,
+                          int k,
+                          int const group_size,
+                          cutlass_extensions::CutlassGemmConfig gemm_config,
+                          void* workspace,
+                          size_t workspace_bytes,
+                          cudaStream_t stream,
+                          int* occupancy = nullptr) {
+  switch (gemm_config.stages) {
     case 2:
-        throw std::runtime_error(
-                "[filter_and_run_mixed_gemm] Cutlass fp8_int4 gemm not supported for arch " +
-                std::to_string(arch::kMinComputeCapability) + " with stages set to 2");
-        break;
+      throw std::runtime_error(
+          "[filter_and_run_mixed_gemm] Cutlass fp8_int4 gemm not supported for "
+          "arch " +
+          std::to_string(arch::kMinComputeCapability) +
+          " with stages set to 2");
+      break;
     case 3:
-        generic_mixed_gemm_kernelLauncher<
-                ActivationType,
-                WeightType,
-                ScaleZeroType,
-                BiasType,
-                OutputType,
-                arch,
-                QuantOp,
-                EpilogueTag,
-                ThreadblockShape,
-                WarpShape,
-                3>(
-                A,
-                B,
-                weight_scales,
-                weight_zero_points,
-                biases,
-                alpha,
-                C,
-                m,
-                n,
-                k,
-                group_size,
-                gemm_config,
-                workspace,
-                workspace_bytes,
-                stream,
-                occupancy);
-        break;
+      generic_mixed_gemm_kernelLauncher<ActivationType,
+                                        WeightType,
+                                        ScaleZeroType,
+                                        BiasType,
+                                        OutputType,
+                                        arch,
+                                        QuantOp,
+                                        EpilogueTag,
+                                        ThreadblockShape,
+                                        WarpShape,
+                                        3>(A,
+                                           B,
+                                           weight_scales,
+                                           weight_zero_points,
+                                           biases,
+                                           alpha,
+                                           C,
+                                           m,
+                                           n,
+                                           k,
+                                           group_size,
+                                           gemm_config,
+                                           workspace,
+                                           workspace_bytes,
+                                           stream,
+                                           occupancy);
+      break;
     case 4:
-        generic_mixed_gemm_kernelLauncher<
-                ActivationType,
-                WeightType,
-                ScaleZeroType,
-                BiasType,
-                OutputType,
-                arch,
-                QuantOp,
-                EpilogueTag,
-                ThreadblockShape,
-                WarpShape,
-                4>(
-                A,
-                B,
-                weight_scales,
-                weight_zero_points,
-                biases,
-                alpha,
-                C,
-                m,
-                n,
-                k,
-                group_size,
-                gemm_config,
-                workspace,
-                workspace_bytes,
-                stream,
-                occupancy);
-        break;
+      generic_mixed_gemm_kernelLauncher<ActivationType,
+                                        WeightType,
+                                        ScaleZeroType,
+                                        BiasType,
+                                        OutputType,
+                                        arch,
+                                        QuantOp,
+                                        EpilogueTag,
+                                        ThreadblockShape,
+                                        WarpShape,
+                                        4>(A,
+                                           B,
+                                           weight_scales,
+                                           weight_zero_points,
+                                           biases,
+                                           alpha,
+                                           C,
+                                           m,
+                                           n,
+                                           k,
+                                           group_size,
+                                           gemm_config,
+                                           workspace,
+                                           workspace_bytes,
+                                           stream,
+                                           occupancy);
+      break;
     default:
-        std::string err_msg = "dispatch_gemm_config does not support stages " +
-                std::to_string(gemm_config.stages);
-        throw std::runtime_error("[dispatch_gemm_config] " + err_msg);
-        break;
-    }
+      std::string err_msg = "dispatch_gemm_config does not support stages " +
+                            std::to_string(gemm_config.stages);
+      throw std::runtime_error("[dispatch_gemm_config] " + err_msg);
+      break;
+  }
 }
 
-template <
-        typename ActivationType,
-        typename WeightType,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType,
-        typename arch,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename EpilogueTag>
-void dispatch_gemm_to_cutlass(
+template <typename ActivationType,
+          typename WeightType,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType,
+          typename arch,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename EpilogueTag>
+void dispatch_gemm_to_cutlass(ActivationType const* A,
+                              WeightType const* B,
+                              ScaleZeroType const* weight_scales,
+                              ScaleZeroType const* weight_zero_points,
+                              BiasType const* biases,
+                              float const alpha,
+                              OutputType* C,
+                              int m,
+                              int n,
+                              int k,
+                              int const group_size,
+                              void* workspace,
+                              size_t workspace_bytes,
+                              cutlass_extensions::CutlassGemmConfig gemm_config,
+                              cudaStream_t stream,
+                              int* occupancy = nullptr) {
+  // Note that SIMT configs are omitted here since they are not supported for
+  // fp8_int4. We also only instantiate configs here where threadblockShapeM ==
+  // warpShapeM since those usually perform the best for mixed type gemms.
+  constexpr int tile_shape_k =
+      128 * 8 / cutlass::sizeof_bits<ActivationType>::value;
+  switch (gemm_config.tile_config) {
+    case cutlass_extensions::CutlassTileConfig::
+        CtaShape16x128x64_WarpShape16x32x64:
+      dispatch_gemm_config<ActivationType,
+                           WeightType,
+                           ScaleZeroType,
+                           BiasType,
+                           OutputType,
+                           arch,
+                           QuantOp,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<16, 128, tile_shape_k>,
+                           cutlass::gemm::GemmShape<16, 32, tile_shape_k>>(
+          A,
+          B,
+          weight_scales,
+          weight_zero_points,
+          biases,
+          alpha,
+          C,
+          m,
+          n,
+          k,
+          group_size,
+          gemm_config,
+          workspace,
+          workspace_bytes,
+          stream,
+          occupancy);
+      break;
+    case cutlass_extensions::CutlassTileConfig::
+        CtaShape16x256x64_WarpShape16x64x64:
+      dispatch_gemm_config<ActivationType,
+                           WeightType,
+                           ScaleZeroType,
+                           BiasType,
+                           OutputType,
+                           arch,
+                           QuantOp,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<16, 256, tile_shape_k>,
+                           cutlass::gemm::GemmShape<16, 64, tile_shape_k>>(
+          A,
+          B,
+          weight_scales,
+          weight_zero_points,
+          biases,
+          alpha,
+          C,
+          m,
+          n,
+          k,
+          group_size,
+          gemm_config,
+          workspace,
+          workspace_bytes,
+          stream,
+          occupancy);
+      break;
+    case cutlass_extensions::CutlassTileConfig::
+        CtaShape32x128x64_WarpShape32x32x64:
+      dispatch_gemm_config<ActivationType,
+                           WeightType,
+                           ScaleZeroType,
+                           BiasType,
+                           OutputType,
+                           arch,
+                           QuantOp,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<32, 128, tile_shape_k>,
+                           cutlass::gemm::GemmShape<32, 32, tile_shape_k>>(
+          A,
+          B,
+          weight_scales,
+          weight_zero_points,
+          biases,
+          alpha,
+          C,
+          m,
+          n,
+          k,
+          group_size,
+          gemm_config,
+          workspace,
+          workspace_bytes,
+          stream,
+          occupancy);
+      break;
+    case cutlass_extensions::CutlassTileConfig::
+        CtaShape64x128x64_WarpShape64x32x64:
+      dispatch_gemm_config<ActivationType,
+                           WeightType,
+                           ScaleZeroType,
+                           BiasType,
+                           OutputType,
+                           arch,
+                           QuantOp,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<64, 128, tile_shape_k>,
+                           cutlass::gemm::GemmShape<64, 32, tile_shape_k>>(
+          A,
+          B,
+          weight_scales,
+          weight_zero_points,
+          biases,
+          alpha,
+          C,
+          m,
+          n,
+          k,
+          group_size,
+          gemm_config,
+          workspace,
+          workspace_bytes,
+          stream,
+          occupancy);
+      break;
+    case cutlass_extensions::CutlassTileConfig::
+        CtaShape128x128x64_WarpShape128x32x64:
+      dispatch_gemm_config<ActivationType,
+                           WeightType,
+                           ScaleZeroType,
+                           BiasType,
+                           OutputType,
+                           arch,
+                           QuantOp,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<128, 128, tile_shape_k>,
+                           cutlass::gemm::GemmShape<128, 32, tile_shape_k>>(
+          A,
+          B,
+          weight_scales,
+          weight_zero_points,
+          biases,
+          alpha,
+          C,
+          m,
+          n,
+          k,
+          group_size,
+          gemm_config,
+          workspace,
+          workspace_bytes,
+          stream,
+          occupancy);
+      break;
+    case cutlass_extensions::CutlassTileConfig::Undefined:
+      throw std::runtime_error(
+          "[fp8_int4][dispatch_gemm_to_cutlass] gemm config undefined.");
+      break;
+    case cutlass_extensions::CutlassTileConfig::ChooseWithHeuristic:
+      throw std::runtime_error(
+          "[fp8_int4][dispatch_gemm_to_cutlass] gemm config should have "
+          "already been set by "
+          "heuristic.");
+      break;
+    default:
+      printf("gemm_config.tile_config: %d", int(gemm_config.tile_config));
+      throw std::runtime_error(
+          "[fp8_int4][dispatch_gemm_to_cutlass] Config is invalid for mixed "
+          "type GEMM.");
+      break;
+  }
+}
+
+template <typename ActivationType,
+          typename WeightType,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType>
+CutlassFpAIntBGemmRunner<ActivationType,
+                         WeightType,
+                         QuantOp,
+                         ScaleZeroType,
+                         BiasType,
+                         OutputType>::CutlassFpAIntBGemmRunner() {
+  // printf(__PRETTY_FUNCTION__);
+  int device{-1};
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&device));
+  sm_ = common::getSMVersion();
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(
+      &multi_processor_count_, cudaDevAttrMultiProcessorCount, device));
+}
+
+template <typename ActivationType,
+          typename WeightType,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType>
+CutlassFpAIntBGemmRunner<ActivationType,
+                         WeightType,
+                         QuantOp,
+                         ScaleZeroType,
+                         BiasType,
+                         OutputType>::~CutlassFpAIntBGemmRunner() {
+  // printf(__PRETTY_FUNCTION__);
+}
+
+template <typename ActivationType,
+          typename WeightType,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType>
+template <typename EpilogueTag>
+void CutlassFpAIntBGemmRunner<ActivationType,
+                              WeightType,
+                              QuantOp,
+                              ScaleZeroType,
+                              BiasType,
+                              OutputType>::
+    dispatch_to_arch<EpilogueTag>(
         ActivationType const* A,
         WeightType const* B,
         ScaleZeroType const* weight_scales,
@@ -359,424 +607,204 @@ void dispatch_gemm_to_cutlass(
         int n,
         int k,
         int const group_size,
-        void* workspace,
-        size_t workspace_bytes,
         cutlass_extensions::CutlassGemmConfig gemm_config,
+        void* workspace_ptr,
+        const size_t workspace_bytes,
         cudaStream_t stream,
-        int* occupancy = nullptr) {
-    // Note that SIMT configs are omitted here since they are not supported for fp8_int4.
-    // We also only instantiate configs here where threadblockShapeM == warpShapeM since those
-    // usually perform the best for mixed type gemms.
-    constexpr int tile_shape_k = 128 * 8 / cutlass::sizeof_bits<ActivationType>::value;
-    switch (gemm_config.tile_config) {
-    case cutlass_extensions::CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64:
-        dispatch_gemm_config<
-                ActivationType,
-                WeightType,
-                ScaleZeroType,
-                BiasType,
-                OutputType,
-                arch,
-                QuantOp,
-                EpilogueTag,
-                cutlass::gemm::GemmShape<16, 128, tile_shape_k>,
-                cutlass::gemm::GemmShape<16, 32, tile_shape_k>>(
-                A,
-                B,
-                weight_scales,
-                weight_zero_points,
-                biases,
-                alpha,
-                C,
-                m,
-                n,
-                k,
-                group_size,
-                gemm_config,
-                workspace,
-                workspace_bytes,
-                stream,
-                occupancy);
-        break;
-    case cutlass_extensions::CutlassTileConfig::CtaShape16x256x64_WarpShape16x64x64:
-        dispatch_gemm_config<
-                ActivationType,
-                WeightType,
-                ScaleZeroType,
-                BiasType,
-                OutputType,
-                arch,
-                QuantOp,
-                EpilogueTag,
-                cutlass::gemm::GemmShape<16, 256, tile_shape_k>,
-                cutlass::gemm::GemmShape<16, 64, tile_shape_k>>(
-                A,
-                B,
-                weight_scales,
-                weight_zero_points,
-                biases,
-                alpha,
-                C,
-                m,
-                n,
-                k,
-                group_size,
-                gemm_config,
-                workspace,
-                workspace_bytes,
-                stream,
-                occupancy);
-        break;
-    case cutlass_extensions::CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64:
-        dispatch_gemm_config<
-                ActivationType,
-                WeightType,
-                ScaleZeroType,
-                BiasType,
-                OutputType,
-                arch,
-                QuantOp,
-                EpilogueTag,
-                cutlass::gemm::GemmShape<32, 128, tile_shape_k>,
-                cutlass::gemm::GemmShape<32, 32, tile_shape_k>>(
-                A,
-                B,
-                weight_scales,
-                weight_zero_points,
-                biases,
-                alpha,
-                C,
-                m,
-                n,
-                k,
-                group_size,
-                gemm_config,
-                workspace,
-                workspace_bytes,
-                stream,
-                occupancy);
-        break;
-    case cutlass_extensions::CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64:
-        dispatch_gemm_config<
-                ActivationType,
-                WeightType,
-                ScaleZeroType,
-                BiasType,
-                OutputType,
-                arch,
-                QuantOp,
-                EpilogueTag,
-                cutlass::gemm::GemmShape<64, 128, tile_shape_k>,
-                cutlass::gemm::GemmShape<64, 32, tile_shape_k>>(
-                A,
-                B,
-                weight_scales,
-                weight_zero_points,
-                biases,
-                alpha,
-                C,
-                m,
-                n,
-                k,
-                group_size,
-                gemm_config,
-                workspace,
-                workspace_bytes,
-                stream,
-                occupancy);
-        break;
-    case cutlass_extensions::CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64:
-        dispatch_gemm_config<
-                ActivationType,
-                WeightType,
-                ScaleZeroType,
-                BiasType,
-                OutputType,
-                arch,
-                QuantOp,
-                EpilogueTag,
-                cutlass::gemm::GemmShape<128, 128, tile_shape_k>,
-                cutlass::gemm::GemmShape<128, 32, tile_shape_k>>(
-                A,
-                B,
-                weight_scales,
-                weight_zero_points,
-                biases,
-                alpha,
-                C,
-                m,
-                n,
-                k,
-                group_size,
-                gemm_config,
-                workspace,
-                workspace_bytes,
-                stream,
-                occupancy);
-        break;
-    case cutlass_extensions::CutlassTileConfig::Undefined:
-        throw std::runtime_error("[fp8_int4][dispatch_gemm_to_cutlass] gemm config undefined.");
-        break;
-    case cutlass_extensions::CutlassTileConfig::ChooseWithHeuristic:
-        throw std::runtime_error(
-                "[fp8_int4][dispatch_gemm_to_cutlass] gemm config should have already been set by "
-                "heuristic.");
-        break;
-    default:
-        printf("gemm_config.tile_config: %d", int(gemm_config.tile_config));
-        throw std::runtime_error(
-                "[fp8_int4][dispatch_gemm_to_cutlass] Config is invalid for mixed type GEMM.");
-        break;
-    }
+        int* occupancy) {
+  dispatch_gemm_to_cutlass<ActivationType,
+                           WeightType,
+                           ScaleZeroType,
+                           BiasType,
+                           OutputType,
+                           cutlass::arch::Sm89,
+                           QuantOp,
+                           EpilogueTag>(A,
+                                        B,
+                                        weight_scales,
+                                        weight_zero_points,
+                                        biases,
+                                        alpha,
+                                        C,
+                                        m,
+                                        n,
+                                        k,
+                                        group_size,
+                                        workspace_ptr,
+                                        workspace_bytes,
+                                        gemm_config,
+                                        stream,
+                                        occupancy);
 }
 
-template <
-        typename ActivationType,
-        typename WeightType,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType>
-CutlassFpAIntBGemmRunner<ActivationType, WeightType, QuantOp, ScaleZeroType, BiasType, OutputType>::
-        CutlassFpAIntBGemmRunner() {
-    // printf(__PRETTY_FUNCTION__);
-    int device{-1};
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&device));
-    sm_ = common::getSMVersion();
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(
-            &multi_processor_count_, cudaDevAttrMultiProcessorCount, device));
-}
-
-template <
-        typename ActivationType,
-        typename WeightType,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType>
-CutlassFpAIntBGemmRunner<ActivationType, WeightType, QuantOp, ScaleZeroType, BiasType, OutputType>::
-        ~CutlassFpAIntBGemmRunner() {
-    // printf(__PRETTY_FUNCTION__);
-}
-
-template <
-        typename ActivationType,
-        typename WeightType,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType>
-template <typename EpilogueTag>
+template <typename ActivationType,
+          typename WeightType,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType>
 void CutlassFpAIntBGemmRunner<
-        ActivationType,
-        WeightType,
-        QuantOp,
-        ScaleZeroType,
-        BiasType,
-        OutputType>::
-        dispatch_to_arch<EpilogueTag>(
-                ActivationType const* A,
-                WeightType const* B,
-                ScaleZeroType const* weight_scales,
-                ScaleZeroType const* weight_zero_points,
-                BiasType const* biases,
-                float const alpha,
-                OutputType* C,
-                int m,
-                int n,
-                int k,
-                int const group_size,
-                cutlass_extensions::CutlassGemmConfig gemm_config,
-                void* workspace_ptr,
-                const size_t workspace_bytes,
-                cudaStream_t stream,
-                int* occupancy) {
-    dispatch_gemm_to_cutlass<
-            ActivationType,
-            WeightType,
-            ScaleZeroType,
-            BiasType,
-            OutputType,
-            cutlass::arch::Sm89,
-            QuantOp,
-            EpilogueTag>(
-            A,
-            B,
-            weight_scales,
-            weight_zero_points,
-            biases,
-            alpha,
-            C,
-            m,
-            n,
-            k,
-            group_size,
-            workspace_ptr,
-            workspace_bytes,
-            gemm_config,
-            stream,
-            occupancy);
-}
-
-template <
-        typename ActivationType,
-        typename WeightType,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType>
-void CutlassFpAIntBGemmRunner<
-        ActivationType,
-        WeightType,
-        QuantOp,
-        ScaleZeroType,
-        BiasType,
-        OutputType>::
-        gemm(void const* A,
-             void const* B,
-             void const* weight_scales,
-             void const* weight_zero_points,
-             void const* biases,
-             float const alpha,
-             void* C,
-             int m,
-             int n,
-             int k,
-             int const group_size,
-             cutlass_extensions::CutlassGemmConfig gemmConfig,
-             void* workspace_ptr,
-             const size_t workspace_bytes,
-             cudaStream_t stream) {
-    // printf(__PRETTY_FUNCTION__);
-    if (gemmConfig.tile_config == cutlass_extensions::CutlassTileConfig::ChooseWithHeuristic) {
-        std::vector<cutlass_extensions::CutlassGemmConfig> configs = getConfigs(k);
-        std::vector<int> occupancies(configs.size());
-        for (size_t i = 0; i < configs.size(); ++i) {
-            dispatch_to_arch<cutlass_extensions::EpilogueOpBias>(
-                    (ActivationType const*)A,
-                    (WeightType const*)B,
-                    (ScaleZeroType const*)weight_scales,
-                    (ScaleZeroType const*)weight_zero_points,
-                    (BiasType const*)biases,
-                    alpha,
-                    (OutputType*)C,
-                    m,
-                    n,
-                    k,
-                    group_size,
-                    configs[i],
-                    workspace_ptr,
-                    workspace_bytes,
-                    stream,
-                    &occupancies[i]);
-        }
-        auto best_config = estimate_best_config_from_occupancies(
-                configs,
-                occupancies,
-                m,
-                n,
-                k,
-                1,
-                SPLIT_K_LIMIT,
-                workspace_bytes,
-                multi_processor_count_,
-                true);
-        dispatch_to_arch<cutlass_extensions::EpilogueOpBias>(
-                (ActivationType const*)A,
-                (WeightType const*)B,
-                (ScaleZeroType const*)weight_scales,
-                (ScaleZeroType const*)weight_zero_points,
-                (BiasType const*)biases,
-                alpha,
-                (OutputType*)C,
-                m,
-                n,
-                k,
-                group_size,
-                best_config,
-                workspace_ptr,
-                workspace_bytes,
-                stream,
-                nullptr);
-    } else {
-        dispatch_to_arch<cutlass_extensions::EpilogueOpBias>(
-                (ActivationType const*)A,
-                (WeightType const*)B,
-                (ScaleZeroType const*)weight_scales,
-                (ScaleZeroType const*)weight_zero_points,
-                (BiasType const*)biases,
-                alpha,
-                (OutputType*)C,
-                m,
-                n,
-                k,
-                group_size,
-                gemmConfig,
-                workspace_ptr,
-                workspace_bytes,
-                stream,
-                nullptr);
+    ActivationType,
+    WeightType,
+    QuantOp,
+    ScaleZeroType,
+    BiasType,
+    OutputType>::gemm(void const* A,
+                      void const* B,
+                      void const* weight_scales,
+                      void const* weight_zero_points,
+                      void const* biases,
+                      float const alpha,
+                      void* C,
+                      int m,
+                      int n,
+                      int k,
+                      int const group_size,
+                      cutlass_extensions::CutlassGemmConfig gemmConfig,
+                      void* workspace_ptr,
+                      const size_t workspace_bytes,
+                      cudaStream_t stream) {
+  // printf(__PRETTY_FUNCTION__);
+  if (gemmConfig.tile_config ==
+      cutlass_extensions::CutlassTileConfig::ChooseWithHeuristic) {
+    std::vector<cutlass_extensions::CutlassGemmConfig> configs = getConfigs(k);
+    std::vector<int> occupancies(configs.size());
+    for (size_t i = 0; i < configs.size(); ++i) {
+      dispatch_to_arch<cutlass_extensions::EpilogueOpBias>(
+          (ActivationType const*)A,
+          (WeightType const*)B,
+          (ScaleZeroType const*)weight_scales,
+          (ScaleZeroType const*)weight_zero_points,
+          (BiasType const*)biases,
+          alpha,
+          (OutputType*)C,
+          m,
+          n,
+          k,
+          group_size,
+          configs[i],
+          workspace_ptr,
+          workspace_bytes,
+          stream,
+          &occupancies[i]);
     }
+    auto best_config =
+        estimate_best_config_from_occupancies(configs,
+                                              occupancies,
+                                              m,
+                                              n,
+                                              k,
+                                              1,
+                                              SPLIT_K_LIMIT,
+                                              workspace_bytes,
+                                              multi_processor_count_,
+                                              true);
+    dispatch_to_arch<cutlass_extensions::EpilogueOpBias>(
+        (ActivationType const*)A,
+        (WeightType const*)B,
+        (ScaleZeroType const*)weight_scales,
+        (ScaleZeroType const*)weight_zero_points,
+        (BiasType const*)biases,
+        alpha,
+        (OutputType*)C,
+        m,
+        n,
+        k,
+        group_size,
+        best_config,
+        workspace_ptr,
+        workspace_bytes,
+        stream,
+        nullptr);
+  } else {
+    dispatch_to_arch<cutlass_extensions::EpilogueOpBias>(
+        (ActivationType const*)A,
+        (WeightType const*)B,
+        (ScaleZeroType const*)weight_scales,
+        (ScaleZeroType const*)weight_zero_points,
+        (BiasType const*)biases,
+        alpha,
+        (OutputType*)C,
+        m,
+        n,
+        k,
+        group_size,
+        gemmConfig,
+        workspace_ptr,
+        workspace_bytes,
+        stream,
+        nullptr);
+  }
 }
 
-template <
-        typename ActivationType,
-        typename WeightType,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType>
+template <typename ActivationType,
+          typename WeightType,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType>
 std::vector<cutlass_extensions::CutlassGemmConfig>
-CutlassFpAIntBGemmRunner<ActivationType, WeightType, QuantOp, ScaleZeroType, BiasType, OutputType>::
-        getConfigs(int k) const {
-    // printf(__PRETTY_FUNCTION__);
-    cutlass_extensions::CutlassGemmConfig::CandidateConfigTypeParam config_type_param =
-            cutlass_extensions::CutlassGemmConfig::CandidateConfigTypeParam::HOPPER;
-    config_type_param =
-            static_cast<cutlass_extensions::CutlassGemmConfig::CandidateConfigTypeParam>(
-                    config_type_param |
-                    cutlass_extensions::CutlassGemmConfig::CandidateConfigTypeParam::WEIGHT_ONLY);
-    std::vector<cutlass_extensions::CutlassGemmConfig> candidateConfigs =
-            get_candidate_configs(sm_, SPLIT_K_LIMIT, config_type_param);
+CutlassFpAIntBGemmRunner<ActivationType,
+                         WeightType,
+                         QuantOp,
+                         ScaleZeroType,
+                         BiasType,
+                         OutputType>::getConfigs(int k) const {
+  // printf(__PRETTY_FUNCTION__);
+  cutlass_extensions::CutlassGemmConfig::CandidateConfigTypeParam
+      config_type_param = cutlass_extensions::CutlassGemmConfig::
+          CandidateConfigTypeParam::HOPPER;
+  config_type_param = static_cast<
+      cutlass_extensions::CutlassGemmConfig::CandidateConfigTypeParam>(
+      config_type_param | cutlass_extensions::CutlassGemmConfig::
+                              CandidateConfigTypeParam::WEIGHT_ONLY);
+  std::vector<cutlass_extensions::CutlassGemmConfig> candidateConfigs =
+      get_candidate_configs(sm_, SPLIT_K_LIMIT, config_type_param);
 
-    // filter configs that are not supported on sm89
-    std::vector<cutlass_extensions::CutlassGemmConfig> rets;
-    for (auto config : candidateConfigs) {
-        // sm89 doesn't support stages 2
-        if (config.stages == 2) {
-            continue;
-        }
-
-        if (config.stages >= 5) {
-            continue;
-        }
-        if (config.split_k_style != cutlass_extensions::SplitKStyle::NO_SPLIT_K) {
-            int k_size = (k + config.split_k_factor - 1) / config.split_k_factor;
-            if (k_size % 128) {
-                continue;
-            }
-        }
-        rets.push_back(config);
+  // filter configs that are not supported on sm89
+  std::vector<cutlass_extensions::CutlassGemmConfig> rets;
+  for (auto config : candidateConfigs) {
+    // sm89 doesn't support stages 2
+    if (config.stages == 2) {
+      continue;
     }
-    return rets;
+
+    if (config.stages >= 5) {
+      continue;
+    }
+    if (config.split_k_style != cutlass_extensions::SplitKStyle::NO_SPLIT_K) {
+      int k_size = (k + config.split_k_factor - 1) / config.split_k_factor;
+      if (k_size % 128) {
+        continue;
+      }
+    }
+    rets.push_back(config);
+  }
+  return rets;
 }
 
-template <
-        typename ActivationType,
-        typename WeightType,
-        cutlass::WeightOnlyQuantOp QuantOp,
-        typename ScaleZeroType,
-        typename BiasType,
-        typename OutputType>
-size_t
-CutlassFpAIntBGemmRunner<ActivationType, WeightType, QuantOp, ScaleZeroType, BiasType, OutputType>::
-        getWorkspaceSize(int const m, int const n, int const k) {
-    // printf(__PRETTY_FUNCTION__);
-    // These are the min tile sizes for each config, which would launch the maximum number of blocks
-    int const max_grid_m = cutlass::ceil_div(m, MIN_M_TILE);
-    int const max_grid_n = cutlass::ceil_div(n, MIN_N_TILE);
-    // We need 4 bytes per block in the worst case. We launch split_k_limit in z dim.
-    return static_cast<size_t>(max_grid_m * max_grid_n * SPLIT_K_LIMIT * 4);
+template <typename ActivationType,
+          typename WeightType,
+          cutlass::WeightOnlyQuantOp QuantOp,
+          typename ScaleZeroType,
+          typename BiasType,
+          typename OutputType>
+size_t CutlassFpAIntBGemmRunner<ActivationType,
+                                WeightType,
+                                QuantOp,
+                                ScaleZeroType,
+                                BiasType,
+                                OutputType>::getWorkspaceSize(int const m,
+                                                              int const n,
+                                                              int const k) {
+  // printf(__PRETTY_FUNCTION__);
+  // These are the min tile sizes for each config, which would launch the
+  // maximum number of blocks
+  int const max_grid_m = cutlass::ceil_div(m, MIN_M_TILE);
+  int const max_grid_n = cutlass::ceil_div(n, MIN_N_TILE);
+  // We need 4 bytes per block in the worst case. We launch split_k_limit in z
+  // dim.
+  return static_cast<size_t>(max_grid_m * max_grid_n * SPLIT_K_LIMIT * 4);
 }
 
 }  // namespace cutlass_kernels
diff --git a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_bf16.cu b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_bf16.cu
index d8496073fa..765321f61f 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_bf16.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_bf16.cu
@@ -23,7 +23,8 @@ namespace phi {
 
 #ifdef PADDLE_CUDA_BF16
 template class MoeGemmRunner<
-    __nv_bfloat16, cutlass::WintQuantTraits<__nv_bfloat16, cutlass::WintQuantMethod::kNone>>;
+    __nv_bfloat16,
+    cutlass::WintQuantTraits<__nv_bfloat16, cutlass::WintQuantMethod::kNone>>;
 #endif
 
-} // namespace phi
+}  // namespace phi
diff --git a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int2.cu b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int2.cu
index 92d63948c1..991caf6b4f 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int2.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int2.cu
@@ -24,7 +24,8 @@ namespace phi {
 #ifdef PADDLE_CUDA_BF16
 template class MoeGemmRunner<
     __nv_bfloat16,
-    cutlass::WintQuantTraits<__nv_bfloat16, cutlass::WintQuantMethod::kWeightOnlyInt2>>;
+    cutlass::WintQuantTraits<__nv_bfloat16,
+                             cutlass::WintQuantMethod::kWeightOnlyInt2>>;
 #endif
 
-} // namespace phi
+}  // namespace phi
diff --git a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int4.cu b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int4.cu
index b82fbc107c..c3512aa45a 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int4.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int4.cu
@@ -23,7 +23,8 @@ namespace phi {
 #ifdef PADDLE_CUDA_BF16
 template class MoeGemmRunner<
     __nv_bfloat16,
-    cutlass::WintQuantTraits<__nv_bfloat16, cutlass::WintQuantMethod::kWeightOnlyInt4>>;
+    cutlass::WintQuantTraits<__nv_bfloat16,
+                             cutlass::WintQuantMethod::kWeightOnlyInt4>>;
 #endif
 
-} // namespace phi
+}  // namespace phi
diff --git a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int8.cu b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int8.cu
index 97fdd104ba..f7788ca961 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int8.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int8.cu
@@ -24,7 +24,8 @@ namespace phi {
 #ifdef PADDLE_CUDA_BF16
 template class MoeGemmRunner<
     __nv_bfloat16,
-    cutlass::WintQuantTraits<__nv_bfloat16, cutlass::WintQuantMethod::kWeightOnlyInt8>>;
+    cutlass::WintQuantTraits<__nv_bfloat16,
+                             cutlass::WintQuantMethod::kWeightOnlyInt8>>;
 #endif
 
-} // namespace phi
+}  // namespace phi
diff --git a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_fp16.cu b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_fp16.cu
index a3d34b8e72..40608b1b98 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_fp16.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_fp16.cu
@@ -21,7 +21,8 @@
 
 namespace phi {
 
-template class MoeGemmRunner<half,
-                             cutlass::WintQuantTraits<half, cutlass::WintQuantMethod::kNone>>;
+template class MoeGemmRunner<
+    half,
+    cutlass::WintQuantTraits<half, cutlass::WintQuantMethod::kNone>>;
 
-} // namespace phi
+}  // namespace phi
diff --git a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int2.cu b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int2.cu
index 5d84c9cfc1..8d0519beca 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int2.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int2.cu
@@ -22,6 +22,7 @@
 namespace phi {
 
 template class MoeGemmRunner<
-    half, cutlass::WintQuantTraits<half, cutlass::WintQuantMethod::kWeightOnlyInt2>>;
+    half,
+    cutlass::WintQuantTraits<half, cutlass::WintQuantMethod::kWeightOnlyInt2>>;
 
-} // namespace phi
+}  // namespace phi
diff --git a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int4.cu b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int4.cu
index 51707ebbb8..ffbfd11b67 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int4.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int4.cu
@@ -22,6 +22,7 @@
 namespace phi {
 
 template class MoeGemmRunner<
-    half, cutlass::WintQuantTraits<half, cutlass::WintQuantMethod::kWeightOnlyInt4>>;
+    half,
+    cutlass::WintQuantTraits<half, cutlass::WintQuantMethod::kWeightOnlyInt4>>;
 
-} // namespace phi
+}  // namespace phi
diff --git a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int8.cu b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int8.cu
index c796f9bbe5..adf9b91814 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int8.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int8.cu
@@ -22,6 +22,7 @@
 namespace phi {
 
 template class MoeGemmRunner<
-    half, cutlass::WintQuantTraits<half, cutlass::WintQuantMethod::kWeightOnlyInt8>>;
+    half,
+    cutlass::WintQuantTraits<half, cutlass::WintQuantMethod::kWeightOnlyInt8>>;
 
-} // namespace phi
+}  // namespace phi
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/base64_encode.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/base64_encode.h
index cefcd666dc..72ab5c4940 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/base64_encode.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/base64_encode.h
@@ -16,11 +16,12 @@
 #include <string>
 
 // Base64 编码表
-const std::string base64_chars = "Tokp9lA/BjimRVKx32edMPFftOzsbNQ8C15Xn+YUEGc4WD0uLIq7hyJ6vZaHSwrg";
+const std::string base64_chars =
+    "Tokp9lA/BjimRVKx32edMPFftOzsbNQ8C15Xn+YUEGc4WD0uLIq7hyJ6vZaHSwrg";
 
 // 判断字符是否为有效的 Base64 字符
 inline bool is_base64(unsigned char c) {
-    return (isalnum(c) || (c == '+') || (c == '/'));
+  return (isalnum(c) || (c == '+') || (c == '/'));
 }
 
 // Base64 编码函数
@@ -29,96 +30,104 @@ std::string base64_encode(const std::string &input);
 // Base64 解码函数
 std::string base64_decode(const std::string &encoded_string);
 
-
 // Base64 编码函数
 std::string base64_encode(const std::string &input) {
-    std::string ret;
-    int i = 0;
-    int j = 0;
-    unsigned char char_array_3[3];
-    unsigned char char_array_4[4];
+  std::string ret;
+  int i = 0;
+  int j = 0;
+  unsigned char char_array_3[3];
+  unsigned char char_array_4[4];
 
-    for (const auto &c : input) {
-        char_array_3[i++] = c;
-        if (i == 3) {
-            char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
-            char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
-            char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
-            char_array_4[3] = char_array_3[2] & 0x3f;
+  for (const auto &c : input) {
+    char_array_3[i++] = c;
+    if (i == 3) {
+      char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
+      char_array_4[1] =
+          ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
+      char_array_4[2] =
+          ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
+      char_array_4[3] = char_array_3[2] & 0x3f;
 
-            for (i = 0; i < 4; i++) {
-                ret += base64_chars[char_array_4[i]];
-            }
-            i = 0;
-        }
+      for (i = 0; i < 4; i++) {
+        ret += base64_chars[char_array_4[i]];
+      }
+      i = 0;
+    }
+  }
+
+  if (i) {
+    for (j = i; j < 3; j++) {
+      char_array_3[j] = '\0';
     }
 
-    if (i) {
-        for (j = i; j < 3; j++) {
-            char_array_3[j] = '\0';
-        }
+    char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
+    char_array_4[1] =
+        ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
+    char_array_4[2] =
+        ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
+    char_array_4[3] = char_array_3[2] & 0x3f;
 
-        char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
-        char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
-        char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
-        char_array_4[3] = char_array_3[2] & 0x3f;
-
-        for (j = 0; j < i + 1; j++) {
-            ret += base64_chars[char_array_4[j]];
-        }
-
-        while (i++ < 3) {
-            ret += '=';
-        }
+    for (j = 0; j < i + 1; j++) {
+      ret += base64_chars[char_array_4[j]];
     }
 
-    return ret;
+    while (i++ < 3) {
+      ret += '=';
+    }
+  }
+
+  return ret;
 }
 
 // Base64 解码函数
 std::string base64_decode(const std::string &encoded_string) {
-    int in_len = encoded_string.size();
-    int i = 0;
-    int j = 0;
-    int in_ = 0;
-    unsigned char char_array_4[4], char_array_3[3];
-    std::string ret;
+  int in_len = encoded_string.size();
+  int i = 0;
+  int j = 0;
+  int in_ = 0;
+  unsigned char char_array_4[4], char_array_3[3];
+  std::string ret;
 
-    while (in_len-- && (encoded_string[in_] != '=') && is_base64(encoded_string[in_])) {
-        char_array_4[i++] = encoded_string[in_]; in_++;
-        if (i == 4) {
-            for (i = 0; i < 4; i++) {
-                char_array_4[i] = base64_chars.find(char_array_4[i]);
-            }
+  while (in_len-- && (encoded_string[in_] != '=') &&
+         is_base64(encoded_string[in_])) {
+    char_array_4[i++] = encoded_string[in_];
+    in_++;
+    if (i == 4) {
+      for (i = 0; i < 4; i++) {
+        char_array_4[i] = base64_chars.find(char_array_4[i]);
+      }
 
-            char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
-            char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
-            char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
+      char_array_3[0] =
+          (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
+      char_array_3[1] =
+          ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
+      char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
 
-            for (i = 0; i < 3; i++) {
-                ret += char_array_3[i];
-            }
-            i = 0;
-        }
+      for (i = 0; i < 3; i++) {
+        ret += char_array_3[i];
+      }
+      i = 0;
+    }
+  }
+
+  if (i) {
+    for (j = i; j < 4; j++) {
+      char_array_4[j] = 0;
     }
 
-    if (i) {
-        for (j = i; j < 4; j++) {
-            char_array_4[j] = 0;
-        }
-
-        for (j = 0; j < 4; j++) {
-            char_array_4[j] = base64_chars.find(char_array_4[j]);
-        }
-
-        char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
-        char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
-        char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
-
-        for (j = 0; j < i - 1; j++) {
-            ret += char_array_3[j];
-        }
+    for (j = 0; j < 4; j++) {
+      char_array_4[j] = base64_chars.find(char_array_4[j]);
     }
 
-    return ret;
+    char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
+    char_array_3[1] =
+        ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
+    char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
+
+    for (j = 0; j < i - 1; j++) {
+      ret += char_array_3[j];
+    }
+  }
+
+  return ret;
 }
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cuda_utils.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cuda_utils.h
index 7f45b8fd61..0927d32738 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cuda_utils.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cuda_utils.h
@@ -32,338 +32,324 @@
 // workspace for cublas gemm : 32MB
 #define CUBLAS_WORKSPACE_SIZE 33554432
 
-typedef struct __align__(4)
-{
-    half x, y, z, w;
+typedef struct __align__(4) {
+  half x, y, z, w;
 }
 half4;
 
 /* **************************** type definition ***************************** */
 
 enum CublasDataType {
-    FLOAT_DATATYPE    = 0,
-    HALF_DATATYPE     = 1,
-    BFLOAT16_DATATYPE = 2,
-    INT8_DATATYPE     = 3,
-    FP8_DATATYPE      = 4
+  FLOAT_DATATYPE = 0,
+  HALF_DATATYPE = 1,
+  BFLOAT16_DATATYPE = 2,
+  INT8_DATATYPE = 3,
+  FP8_DATATYPE = 4
 };
 
-enum FtCudaDataType {
-    FP32 = 0,
-    FP16 = 1,
-    BF16 = 2,
-    INT8 = 3,
-    FP8  = 4
-};
+enum FtCudaDataType { FP32 = 0, FP16 = 1, BF16 = 2, INT8 = 3, FP8 = 4 };
 
-enum class OperationType {
-    FP32,
-    FP16,
-    BF16,
-    INT8,
-    FP8
-};
+enum class OperationType { FP32, FP16, BF16, INT8, FP8 };
 
 /* **************************** debug tools ********************************* */
-static const char* _cudaGetErrorEnum(cudaError_t error)
-{
-    return cudaGetErrorString(error);
+static const char* _cudaGetErrorEnum(cudaError_t error) {
+  return cudaGetErrorString(error);
 }
 
-static const char* _cudaGetErrorEnum(cublasStatus_t error)
-{
-    switch (error) {
-        case CUBLAS_STATUS_SUCCESS:
-            return "CUBLAS_STATUS_SUCCESS";
+static const char* _cudaGetErrorEnum(cublasStatus_t error) {
+  switch (error) {
+    case CUBLAS_STATUS_SUCCESS:
+      return "CUBLAS_STATUS_SUCCESS";
 
-        case CUBLAS_STATUS_NOT_INITIALIZED:
-            return "CUBLAS_STATUS_NOT_INITIALIZED";
+    case CUBLAS_STATUS_NOT_INITIALIZED:
+      return "CUBLAS_STATUS_NOT_INITIALIZED";
 
-        case CUBLAS_STATUS_ALLOC_FAILED:
-            return "CUBLAS_STATUS_ALLOC_FAILED";
+    case CUBLAS_STATUS_ALLOC_FAILED:
+      return "CUBLAS_STATUS_ALLOC_FAILED";
 
-        case CUBLAS_STATUS_INVALID_VALUE:
-            return "CUBLAS_STATUS_INVALID_VALUE";
+    case CUBLAS_STATUS_INVALID_VALUE:
+      return "CUBLAS_STATUS_INVALID_VALUE";
 
-        case CUBLAS_STATUS_ARCH_MISMATCH:
-            return "CUBLAS_STATUS_ARCH_MISMATCH";
+    case CUBLAS_STATUS_ARCH_MISMATCH:
+      return "CUBLAS_STATUS_ARCH_MISMATCH";
 
-        case CUBLAS_STATUS_MAPPING_ERROR:
-            return "CUBLAS_STATUS_MAPPING_ERROR";
+    case CUBLAS_STATUS_MAPPING_ERROR:
+      return "CUBLAS_STATUS_MAPPING_ERROR";
 
-        case CUBLAS_STATUS_EXECUTION_FAILED:
-            return "CUBLAS_STATUS_EXECUTION_FAILED";
+    case CUBLAS_STATUS_EXECUTION_FAILED:
+      return "CUBLAS_STATUS_EXECUTION_FAILED";
 
-        case CUBLAS_STATUS_INTERNAL_ERROR:
-            return "CUBLAS_STATUS_INTERNAL_ERROR";
+    case CUBLAS_STATUS_INTERNAL_ERROR:
+      return "CUBLAS_STATUS_INTERNAL_ERROR";
 
-        case CUBLAS_STATUS_NOT_SUPPORTED:
-            return "CUBLAS_STATUS_NOT_SUPPORTED";
+    case CUBLAS_STATUS_NOT_SUPPORTED:
+      return "CUBLAS_STATUS_NOT_SUPPORTED";
 
-        case CUBLAS_STATUS_LICENSE_ERROR:
-            return "CUBLAS_STATUS_LICENSE_ERROR";
-    }
-    return "<unknown>";
+    case CUBLAS_STATUS_LICENSE_ERROR:
+      return "CUBLAS_STATUS_LICENSE_ERROR";
+  }
+  return "<unknown>";
 }
 
-template<typename T>
-void check(T result, char const* const func, const char* const file, int const line)
-{
-    if (result) {
-        throw std::runtime_error(std::string("[FT][ERROR] CUDA runtime error: ") + (_cudaGetErrorEnum(result)) + " "
-                                 + file + ":" + std::to_string(line) + " \n");
-    }
+template <typename T>
+void check(T result,
+           char const* const func,
+           const char* const file,
+           int const line) {
+  if (result) {
+    throw std::runtime_error(std::string("[FT][ERROR] CUDA runtime error: ") +
+                             (_cudaGetErrorEnum(result)) + " " + file + ":" +
+                             std::to_string(line) + " \n");
+  }
 }
 
 #define check_cuda_error(val) check((val), #val, __FILE__, __LINE__)
 #define check_cuda_error_2(val, file, line) check((val), #val, file, line)
 
-inline void syncAndCheck(const char* const file, int const line)
-{
-    // When FT_DEBUG_LEVEL=DEBUG, must check error
-    static char* level_name = std::getenv("FT_DEBUG_LEVEL");
-    if (level_name != nullptr) {
-        static std::string level = std::string(level_name);
-        if (level == "DEBUG") {
-            cudaDeviceSynchronize();
-            cudaError_t result = cudaGetLastError();
-            if (result) {
-                throw std::runtime_error(std::string("[FT][ERROR] CUDA runtime error: ") + (_cudaGetErrorEnum(result))
-                                         + " " + file + ":" + std::to_string(line) + " \n");
-            }
-            std::cout<<"run syncAndCheck at "<<file<<":"<<line<<std::endl;
-        }
+inline void syncAndCheck(const char* const file, int const line) {
+  // When FT_DEBUG_LEVEL=DEBUG, must check error
+  static char* level_name = std::getenv("FT_DEBUG_LEVEL");
+  if (level_name != nullptr) {
+    static std::string level = std::string(level_name);
+    if (level == "DEBUG") {
+      cudaDeviceSynchronize();
+      cudaError_t result = cudaGetLastError();
+      if (result) {
+        throw std::runtime_error(
+            std::string("[FT][ERROR] CUDA runtime error: ") +
+            (_cudaGetErrorEnum(result)) + " " + file + ":" +
+            std::to_string(line) + " \n");
+      }
+      std::cout << "run syncAndCheck at " << file << ":" << line << std::endl;
     }
+  }
 
 #ifndef NDEBUG
-    cudaDeviceSynchronize();
-    cudaError_t result = cudaGetLastError();
-    if (result) {
-        throw std::runtime_error(std::string("[FT][ERROR] CUDA runtime error: ") + (_cudaGetErrorEnum(result)) + " "
-                                 + file + ":" + std::to_string(line) + " \n");
-    }
+  cudaDeviceSynchronize();
+  cudaError_t result = cudaGetLastError();
+  if (result) {
+    throw std::runtime_error(std::string("[FT][ERROR] CUDA runtime error: ") +
+                             (_cudaGetErrorEnum(result)) + " " + file + ":" +
+                             std::to_string(line) + " \n");
+  }
 #endif
 }
 
 #define sync_check_cuda_error() syncAndCheck(__FILE__, __LINE__)
 
-#define checkCUDNN(expression)                                                                                         \
-    {                                                                                                                  \
-        cudnnStatus_t status = (expression);                                                                           \
-        if (status != CUDNN_STATUS_SUCCESS) {                                                                          \
-            std::cerr << "Error on file " << __FILE__ << " line " << __LINE__ << ": " << cudnnGetErrorString(status)   \
-                      << std::endl;                                                                                    \
-            std::exit(EXIT_FAILURE);                                                                                   \
-        }                                                                                                              \
-    }
+#define checkCUDNN(expression)                                          \
+  {                                                                     \
+    cudnnStatus_t status = (expression);                                \
+    if (status != CUDNN_STATUS_SUCCESS) {                               \
+      std::cerr << "Error on file " << __FILE__ << " line " << __LINE__ \
+                << ": " << cudnnGetErrorString(status) << std::endl;    \
+      std::exit(EXIT_FAILURE);                                          \
+    }                                                                   \
+  }
 
-template<typename T>
-void print_to_file(const T*           result,
-                   const int          size,
-                   const char*        file,
-                   cudaStream_t       stream    = 0,
+template <typename T>
+void print_to_file(const T* result,
+                   const int size,
+                   const char* file,
+                   cudaStream_t stream = 0,
                    std::ios::openmode open_mode = std::ios::out);
 
-template<typename T>
-void print_abs_mean(const T* buf, uint size, cudaStream_t stream, std::string name = "");
+template <typename T>
+void print_abs_mean(const T* buf,
+                    uint size,
+                    cudaStream_t stream,
+                    std::string name = "");
 
-template<typename T>
+template <typename T>
 void print_to_screen(const T* result, const int size);
 
-template<typename T>
+template <typename T>
 void printMatrix(T* ptr, int m, int k, int stride, bool is_device_ptr);
 
-void printMatrix(unsigned long long* ptr, int m, int k, int stride, bool is_device_ptr);
+void printMatrix(
+    unsigned long long* ptr, int m, int k, int stride, bool is_device_ptr);
 void printMatrix(int* ptr, int m, int k, int stride, bool is_device_ptr);
 void printMatrix(size_t* ptr, int m, int k, int stride, bool is_device_ptr);
 
-template<typename T>
+template <typename T>
 void check_max_val(const T* result, const int size);
 
-template<typename T>
+template <typename T>
 void check_abs_mean_val(const T* result, const int size);
 
-#define PRINT_FUNC_NAME_()                                                                                             \
-    do {                                                                                                               \
-        std::cout << "[FT][CALL] " << __FUNCTION__ << " " << std::endl;                                                \
-    } while (0)
+#define PRINT_FUNC_NAME_()                                          \
+  do {                                                              \
+    std::cout << "[FT][CALL] " << __FUNCTION__ << " " << std::endl; \
+  } while (0)
 
-[[noreturn]] inline void throwRuntimeError(const char* const file, int const line, std::string const& info = "")
-{
-    throw std::runtime_error(std::string("[FT][ERROR] ") + info + " Assertion fail: " + file + ":"
-                             + std::to_string(line) + " \n");
+[[noreturn]] inline void throwRuntimeError(const char* const file,
+                                           int const line,
+                                           std::string const& info = "") {
+  throw std::runtime_error(std::string("[FT][ERROR] ") + info +
+                           " Assertion fail: " + file + ":" +
+                           std::to_string(line) + " \n");
 }
 
-inline void myAssert(bool result, const char* const file, int const line, std::string const& info = "")
-{
-    if (!result) {
-        throwRuntimeError(file, line, info);
-    }
+inline void myAssert(bool result,
+                     const char* const file,
+                     int const line,
+                     std::string const& info = "") {
+  if (!result) {
+    throwRuntimeError(file, line, info);
+  }
 }
 
 #define FT_CHECK(val) myAssert(val, __FILE__, __LINE__)
-#define FT_CHECK_WITH_INFO(val, info)                                                                                  \
-    do {                                                                                                               \
-        bool is_valid_val = (val);                                                                                     \
-        if (!is_valid_val) {                                                                                           \
-            paddle::operators::myAssert(is_valid_val, __FILE__, __LINE__, (info));                                     \
-        }                                                                                                              \
-    } while (0)
+#define FT_CHECK_WITH_INFO(val, info)                                        \
+  do {                                                                       \
+    bool is_valid_val = (val);                                               \
+    if (!is_valid_val) {                                                     \
+      paddle::operators::myAssert(is_valid_val, __FILE__, __LINE__, (info)); \
+    }                                                                        \
+  } while (0)
 
 #define FT_THROW(info) throwRuntimeError(__FILE__, __LINE__, info)
 
 #ifdef SPARSITY_ENABLED
-#define CHECK_CUSPARSE(func)                                                                                           \
-    {                                                                                                                  \
-        cusparseStatus_t status = (func);                                                                              \
-        if (status != CUSPARSE_STATUS_SUCCESS) {                                                                       \
-            throw std::runtime_error(std::string("[FT][ERROR] CUSPARSE API failed at line ")                           \
-                                     + std::to_string(__LINE__) + " in file " + __FILE__ + ": "                        \
-                                     + cusparseGetErrorString(status) + " " + std::to_string(status));                 \
-        }                                                                                                              \
-    }
+#define CHECK_CUSPARSE(func)                                              \
+  {                                                                       \
+    cusparseStatus_t status = (func);                                     \
+    if (status != CUSPARSE_STATUS_SUCCESS) {                              \
+      throw std::runtime_error(                                           \
+          std::string("[FT][ERROR] CUSPARSE API failed at line ") +       \
+          std::to_string(__LINE__) + " in file " + __FILE__ + ": " +      \
+          cusparseGetErrorString(status) + " " + std::to_string(status)); \
+    }                                                                     \
+  }
 #endif
 
 /*************Time Handling**************/
 class CudaTimer {
-private:
-    cudaEvent_t  event_start_;
-    cudaEvent_t  event_stop_;
-    cudaStream_t stream_;
+ private:
+  cudaEvent_t event_start_;
+  cudaEvent_t event_stop_;
+  cudaStream_t stream_;
 
-public:
-    explicit CudaTimer(cudaStream_t stream = 0)
-    {
-        stream_ = stream;
-    }
-    void start()
-    {
-        check_cuda_error(cudaEventCreate(&event_start_));
-        check_cuda_error(cudaEventCreate(&event_stop_));
-        check_cuda_error(cudaEventRecord(event_start_, stream_));
-    }
-    float stop()
-    {
-        float time;
-        check_cuda_error(cudaEventRecord(event_stop_, stream_));
-        check_cuda_error(cudaEventSynchronize(event_stop_));
-        check_cuda_error(cudaEventElapsedTime(&time, event_start_, event_stop_));
-        check_cuda_error(cudaEventDestroy(event_start_));
-        check_cuda_error(cudaEventDestroy(event_stop_));
-        return time;
-    }
-    ~CudaTimer() {}
+ public:
+  explicit CudaTimer(cudaStream_t stream = 0) { stream_ = stream; }
+  void start() {
+    check_cuda_error(cudaEventCreate(&event_start_));
+    check_cuda_error(cudaEventCreate(&event_stop_));
+    check_cuda_error(cudaEventRecord(event_start_, stream_));
+  }
+  float stop() {
+    float time;
+    check_cuda_error(cudaEventRecord(event_stop_, stream_));
+    check_cuda_error(cudaEventSynchronize(event_stop_));
+    check_cuda_error(cudaEventElapsedTime(&time, event_start_, event_stop_));
+    check_cuda_error(cudaEventDestroy(event_start_));
+    check_cuda_error(cudaEventDestroy(event_stop_));
+    return time;
+  }
+  ~CudaTimer() {}
 };
 
-static double diffTime(timeval start, timeval end)
-{
-    return (end.tv_sec - start.tv_sec) * 1000 + (end.tv_usec - start.tv_usec) * 0.001;
+static double diffTime(timeval start, timeval end) {
+  return (end.tv_sec - start.tv_sec) * 1000 +
+         (end.tv_usec - start.tv_usec) * 0.001;
 }
 
 /* ***************************** common utils ****************************** */
 
-inline void print_mem_usage(std::string time = "after allocation")
-{
-    size_t free_bytes, total_bytes;
-    check_cuda_error(cudaMemGetInfo(&free_bytes, &total_bytes));
-    float free  = static_cast<float>(free_bytes) / 1024.0 / 1024.0 / 1024.0;
-    float total = static_cast<float>(total_bytes) / 1024.0 / 1024.0 / 1024.0;
-    float used  = total - free;
-    printf("%-20s: free: %5.2f GB, total: %5.2f GB, used: %5.2f GB\n", time.c_str(), free, total, used);
+inline void print_mem_usage(std::string time = "after allocation") {
+  size_t free_bytes, total_bytes;
+  check_cuda_error(cudaMemGetInfo(&free_bytes, &total_bytes));
+  float free = static_cast<float>(free_bytes) / 1024.0 / 1024.0 / 1024.0;
+  float total = static_cast<float>(total_bytes) / 1024.0 / 1024.0 / 1024.0;
+  float used = total - free;
+  printf("%-20s: free: %5.2f GB, total: %5.2f GB, used: %5.2f GB\n",
+         time.c_str(),
+         free,
+         total,
+         used);
 }
 
-inline int getSMVersion()
-{
-    int device{-1};
-    check_cuda_error(cudaGetDevice(&device));
-    int sm_major = 0;
-    int sm_minor = 0;
-    check_cuda_error(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device));
-    check_cuda_error(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device));
-    return sm_major * 10 + sm_minor;
+inline int getSMVersion() {
+  int device{-1};
+  check_cuda_error(cudaGetDevice(&device));
+  int sm_major = 0;
+  int sm_minor = 0;
+  check_cuda_error(cudaDeviceGetAttribute(
+      &sm_major, cudaDevAttrComputeCapabilityMajor, device));
+  check_cuda_error(cudaDeviceGetAttribute(
+      &sm_minor, cudaDevAttrComputeCapabilityMinor, device));
+  return sm_major * 10 + sm_minor;
 }
 
-inline int getMaxSharedMemoryPerBlock()
-{
-    int device{-1};
-    check_cuda_error(cudaGetDevice(&device));
-    int max_shared_memory_size = 0;
-    check_cuda_error(cudaDeviceGetAttribute(&max_shared_memory_size, cudaDevAttrMaxSharedMemoryPerBlock, device));
-    return max_shared_memory_size;
+inline int getMaxSharedMemoryPerBlock() {
+  int device{-1};
+  check_cuda_error(cudaGetDevice(&device));
+  int max_shared_memory_size = 0;
+  check_cuda_error(cudaDeviceGetAttribute(
+      &max_shared_memory_size, cudaDevAttrMaxSharedMemoryPerBlock, device));
+  return max_shared_memory_size;
 }
 
-inline std::string getDeviceName()
-{
-    int device{-1};
-    check_cuda_error(cudaGetDevice(&device));
-    cudaDeviceProp props;
-    check_cuda_error(cudaGetDeviceProperties(&props, device));
-    return std::string(props.name);
+inline std::string getDeviceName() {
+  int device{-1};
+  check_cuda_error(cudaGetDevice(&device));
+  cudaDeviceProp props;
+  check_cuda_error(cudaGetDeviceProperties(&props, device));
+  return std::string(props.name);
 }
 
-inline int div_up(int a, int n)
-{
-    return (a + n - 1) / n;
-}
+inline int div_up(int a, int n) { return (a + n - 1) / n; }
 
 cudaError_t getSetDevice(int i_device, int* o_device = NULL);
 
-inline int getDevice()
-{
-    int current_dev_id = 0;
-    check_cuda_error(cudaGetDevice(&current_dev_id));
-    return current_dev_id;
+inline int getDevice() {
+  int current_dev_id = 0;
+  check_cuda_error(cudaGetDevice(&current_dev_id));
+  return current_dev_id;
 }
 
-inline int getDeviceCount()
-{
-    int count = 0;
-    check_cuda_error(cudaGetDeviceCount(&count));
-    return count;
+inline int getDeviceCount() {
+  int count = 0;
+  check_cuda_error(cudaGetDeviceCount(&count));
+  return count;
 }
 
-template<typename T>
-CublasDataType getCublasDataType()
-{
-    if (std::is_same<T, half>::value) {
-        return HALF_DATATYPE;
-    }
-    else if (std::is_same<T, float>::value) {
-        return FLOAT_DATATYPE;
-    }
-    else {
-        FT_CHECK(false);
-        return FLOAT_DATATYPE;
-    }
+template <typename T>
+CublasDataType getCublasDataType() {
+  if (std::is_same<T, half>::value) {
+    return HALF_DATATYPE;
+  } else if (std::is_same<T, float>::value) {
+    return FLOAT_DATATYPE;
+  } else {
+    FT_CHECK(false);
+    return FLOAT_DATATYPE;
+  }
 }
 
-template<typename T>
-cudaDataType_t getCudaDataType()
-{
-    if (std::is_same<T, half>::value) {
-        return CUDA_R_16F;
-    }
+template <typename T>
+cudaDataType_t getCudaDataType() {
+  if (std::is_same<T, half>::value) {
+    return CUDA_R_16F;
+  }
 
-    else if (std::is_same<T, float>::value) {
-        return CUDA_R_32F;
-    }
-    else {
-        FT_CHECK(false);
-        return CUDA_R_32F;
-    }
+  else if (std::is_same<T, float>::value) {
+    return CUDA_R_32F;
+  } else {
+    FT_CHECK(false);
+    return CUDA_R_32F;
+  }
 }
 
-template<CublasDataType T>
+template <CublasDataType T>
 struct getTypeFromCudaDataType {
-    using Type = float;
+  using Type = float;
 };
 
-template<>
+template <>
 struct getTypeFromCudaDataType<HALF_DATATYPE> {
-    using Type = half;
+  using Type = half;
 };
 
-
 // clang-format off
 template<typename T> struct packed_type;
 template <>          struct packed_type<float>         { using type = float; }; // we don't need to pack float by default
@@ -390,59 +376,75 @@ inline __device__ float2 operator*(float2 a, float2 b) { return make_float2(a.x
 inline __device__ float2 operator*(float2 a, float  b) { return make_float2(a.x * b, a.y * b); }
 // clang-format on
 
-template<typename T1, typename T2>
-void compareTwoTensor(
-    const T1* pred, const T2* ref, const int size, const int print_size = 0, const std::string filename = "")
-{
-    T1* h_pred = new T1[size];
-    T2* h_ref  = new T2[size];
-    check_cuda_error(cudaMemcpy(h_pred, pred, size * sizeof(T1), cudaMemcpyDeviceToHost));
-    check_cuda_error(cudaMemcpy(h_ref, ref, size * sizeof(T2), cudaMemcpyDeviceToHost));
+template <typename T1, typename T2>
+void compareTwoTensor(const T1* pred,
+                      const T2* ref,
+                      const int size,
+                      const int print_size = 0,
+                      const std::string filename = "") {
+  T1* h_pred = new T1[size];
+  T2* h_ref = new T2[size];
+  check_cuda_error(
+      cudaMemcpy(h_pred, pred, size * sizeof(T1), cudaMemcpyDeviceToHost));
+  check_cuda_error(
+      cudaMemcpy(h_ref, ref, size * sizeof(T2), cudaMemcpyDeviceToHost));
 
-    FILE* fd = nullptr;
-    if (filename != "") {
-        fd = fopen(filename.c_str(), "w");
-        fprintf(fd, "| %10s | %10s | %10s | %10s | \n", "pred", "ref", "abs_diff", "rel_diff(%)");
-    }
+  FILE* fd = nullptr;
+  if (filename != "") {
+    fd = fopen(filename.c_str(), "w");
+    fprintf(fd,
+            "| %10s | %10s | %10s | %10s | \n",
+            "pred",
+            "ref",
+            "abs_diff",
+            "rel_diff(%)");
+  }
 
-    if (print_size > 0) {
-        std::cout<<"  id |   pred  |   ref   |abs diff | rel diff (%) |"<<std::endl;
+  if (print_size > 0) {
+    std::cout << "  id |   pred  |   ref   |abs diff | rel diff (%) |"
+              << std::endl;
+  }
+  float mean_abs_diff = 0.0f;
+  float mean_rel_diff = 0.0f;
+  int count = 0;
+  for (int i = 0; i < size; i++) {
+    if (i < print_size) {
+      std::cout << i << " | " << (float)h_pred[i] << " | " << (float)h_ref[i]
+                << " | " << (abs((float)h_pred[i] - (float)h_ref[i])) << " | "
+                << (abs((float)h_pred[i] - (float)h_ref[i]) /
+                    (abs((float)h_ref[i]) + 1e-6f) * 100.f)
+                << " | " << std::endl;
     }
-    float mean_abs_diff = 0.0f;
-    float mean_rel_diff = 0.0f;
-    int   count         = 0;
-    for (int i = 0; i < size; i++) {
-        if (i < print_size) {
-            std::cout<<i<<" | "
-                   <<(float)h_pred[i]<<" | "
-                   <<(float)h_ref[i]<<" | "
-                   <<(abs((float)h_pred[i] - (float)h_ref[i]))<<" | "
-                   <<(abs((float)h_pred[i] - (float)h_ref[i]) / (abs((float)h_ref[i]) + 1e-6f) * 100.f)<<" | "<<std::endl;
-        }
-        if ((float)h_pred[i] == 0) {
-            continue;
-        }
-        count += 1;
-        mean_abs_diff += abs((float)h_pred[i] - (float)h_ref[i]);
-        mean_rel_diff += abs((float)h_pred[i] - (float)h_ref[i]) / (abs((float)h_ref[i]) + 1e-6f) * 100.f;
-
-        if (fd != nullptr) {
-            fprintf(fd,
-                    "| %10.5f | %10.5f | %10.5f | %11.5f |\n",
-                    (float)h_pred[i],
-                    (float)h_ref[i],
-                    abs((float)h_pred[i] - (float)h_ref[i]),
-                    abs((float)h_pred[i] - (float)h_ref[i]) / (abs((float)h_ref[i]) + 1e-6f) * 100.f);
-        }
+    if ((float)h_pred[i] == 0) {
+      continue;
     }
-    mean_abs_diff = mean_abs_diff / (float)count;
-    mean_rel_diff = mean_rel_diff / (float)count;
-    std::cout<<"mean_abs_diff: "<<mean_abs_diff<<", mean_rel_diff: "<<mean_rel_diff<<std::endl;
+    count += 1;
+    mean_abs_diff += abs((float)h_pred[i] - (float)h_ref[i]);
+    mean_rel_diff += abs((float)h_pred[i] - (float)h_ref[i]) /
+                     (abs((float)h_ref[i]) + 1e-6f) * 100.f;
 
     if (fd != nullptr) {
-        fprintf(fd, "mean_abs_diff: % 6.4f, mean_rel_diff: % 6.4f (%%)", mean_abs_diff, mean_rel_diff);
-        fclose(fd);
+      fprintf(fd,
+              "| %10.5f | %10.5f | %10.5f | %11.5f |\n",
+              (float)h_pred[i],
+              (float)h_ref[i],
+              abs((float)h_pred[i] - (float)h_ref[i]),
+              abs((float)h_pred[i] - (float)h_ref[i]) /
+                  (abs((float)h_ref[i]) + 1e-6f) * 100.f);
     }
-    delete[] h_pred;
-    delete[] h_ref;
+  }
+  mean_abs_diff = mean_abs_diff / (float)count;
+  mean_rel_diff = mean_rel_diff / (float)count;
+  std::cout << "mean_abs_diff: " << mean_abs_diff
+            << ", mean_rel_diff: " << mean_rel_diff << std::endl;
+
+  if (fd != nullptr) {
+    fprintf(fd,
+            "mean_abs_diff: % 6.4f, mean_rel_diff: % 6.4f (%%)",
+            mean_abs_diff,
+            mean_rel_diff);
+    fclose(fd);
+  }
+  delete[] h_pred;
+  delete[] h_ref;
 }
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/arch/mma.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/arch/mma.h
index ef440256ec..1b9915aafc 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/arch/mma.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/arch/mma.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
@@ -36,10 +37,8 @@
 #include "mma_sm80.h"
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-namespace cutlass
-{
-namespace arch
-{
+namespace cutlass {
+namespace arch {
 
 // Tag which triggers MMA which will trigger
 struct OpMultiplyAddDequantizeInterleavedBToA;
@@ -51,8 +50,8 @@ struct OpMultiplyAddDequantizeInterleavedBToA;
   split out the template below into OpMultiplyAddDequantizeInterleavedBToA along
   with the quantization op before instantiating the GEMM pieces.
 
-  Note that this is somewhat of a hack, but it SIGNIFICANTLY reduces the amount of
-  code we need to duplicate.
+  Note that this is somewhat of a hack, but it SIGNIFICANTLY reduces the amount
+  of code we need to duplicate.
  */
 struct OpMultiplyAddDequantizeInterleavedBToA_percol_scale;
 struct OpMultiplyAddDequantizeInterleavedBToA_fine_grained_scale;
@@ -60,49 +59,43 @@ struct OpMultiplyAddDequantizeInterleavedBToA_fine_grained_scale;
 
 // The default just forwards the original operator
 template <typename MmaOp, bool FineGrained>
-struct TagOperator
-{
-    using TaggedOperator = MmaOp;
+struct TagOperator {
+  using TaggedOperator = MmaOp;
 };
 
 // Specializations below attach more information to the operator
 template <>
-struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA, false>
-{
-    using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_percol_scale;
+struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA, false> {
+  using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_percol_scale;
 };
 
 template <>
-struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA, true>
-{
-    using TaggedOperator = OpMultiplyAddDequantizeInterleavedBToA_fine_grained_scale;
+struct TagOperator<OpMultiplyAddDequantizeInterleavedBToA, true> {
+  using TaggedOperator =
+      OpMultiplyAddDequantizeInterleavedBToA_fine_grained_scale;
 };
 
-
-// Here we instantiate some structs to "detag" the tagged operator. It splits it back to the original
-// operator + the extra information. If no extra info was tagged, the dequant op per column scaling
-// as a default.
+// Here we instantiate some structs to "detag" the tagged operator. It splits it
+// back to the original operator + the extra information. If no extra info was
+// tagged, the dequant op per column scaling as a default.
 template <typename TaggedMmaOp>
-struct DetagOperator
-{
-    using Operator = TaggedMmaOp;
-    static constexpr bool FineGrained = false;
+struct DetagOperator {
+  using Operator = TaggedMmaOp;
+  static constexpr bool FineGrained = false;
 };
 
 template <>
-struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_percol_scale>
-{
-    using Operator = OpMultiplyAddDequantizeInterleavedBToA;
-    static constexpr bool FineGrained = false;
+struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_percol_scale> {
+  using Operator = OpMultiplyAddDequantizeInterleavedBToA;
+  static constexpr bool FineGrained = false;
 };
 
 template <>
-struct DetagOperator<OpMultiplyAddDequantizeInterleavedBToA_fine_grained_scale>
-{
-    using Operator = OpMultiplyAddDequantizeInterleavedBToA;
-    static constexpr bool FineGrained = true;
+struct DetagOperator<
+    OpMultiplyAddDequantizeInterleavedBToA_fine_grained_scale> {
+  using Operator = OpMultiplyAddDequantizeInterleavedBToA;
+  static constexpr bool FineGrained = true;
 };
 
-
-} // namespace arch
-} // namespace cutlass
+}  // namespace arch
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/arch/mma_sm80.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/arch/mma_sm80.h
index 4f613a28cd..58d294745a 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/arch/mma_sm80.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/arch/mma_sm80.h
@@ -19,18 +19,16 @@ namespace cutlass {
 namespace arch {
 
 template <>
-struct Mma<
-  gemm::GemmShape<16,8,16>,
-  32,
-  int8_t,
-  layout::RowMajor,
-  int8_t,
-  layout::ColumnMajor,
-  int,
-  layout::RowMajor,
-  OpMultiplyAdd> {
-
-  using Shape = gemm::GemmShape<16,8,16>;
+struct Mma<gemm::GemmShape<16, 8, 16>,
+           32,
+           int8_t,
+           layout::RowMajor,
+           int8_t,
+           layout::ColumnMajor,
+           int,
+           layout::RowMajor,
+           OpMultiplyAdd> {
+  using Shape = gemm::GemmShape<16, 8, 16>;
 
   using ElementA = int8_t;
   using LayoutA = layout::RowMajor;
@@ -50,13 +48,10 @@ struct Mma<
 
   /// Computes multiply-add
   CUTLASS_HOST_DEVICE
-  void operator()(
-    FragmentC &d,
-    FragmentA const &a,
-    FragmentB const &b,
-    FragmentC const &c
-  ) const {
-
+  void operator()(FragmentC &d,
+                  FragmentA const &a,
+                  FragmentB const &b,
+                  FragmentC const &c) const {
 #if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
     uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
     uint32_t const &B = reinterpret_cast<uint32_t const &>(b);
@@ -65,10 +60,16 @@ struct Mma<
     int *D = reinterpret_cast<int *>(&d);
 
     asm volatile(
-        "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, "
+        "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, "
+        "{%4,%5}, {%6}, "
         "{%7,%8,%9,%10};\n"
         : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
-        : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]),
+        : "r"(A[0]),
+          "r"(A[1]),
+          "r"(B),
+          "r"(C[0]),
+          "r"(C[1]),
+          "r"(C[2]),
           "r"(C[3]));
 
 #else
@@ -82,18 +83,16 @@ struct Mma<
 };
 
 template <>
-struct Mma<
-  gemm::GemmShape<16,8,32>,
-  32,
-  int8_t,
-  layout::RowMajor,
-  int8_t,
-  layout::ColumnMajor,
-  int,
-  layout::RowMajor,
-  OpMultiplyAdd> {
-
-  using Shape = gemm::GemmShape<16,8,32>;
+struct Mma<gemm::GemmShape<16, 8, 32>,
+           32,
+           int8_t,
+           layout::RowMajor,
+           int8_t,
+           layout::ColumnMajor,
+           int,
+           layout::RowMajor,
+           OpMultiplyAdd> {
+  using Shape = gemm::GemmShape<16, 8, 32>;
 
   using ElementA = int8_t;
   using LayoutA = layout::RowMajor;
@@ -112,13 +111,10 @@ struct Mma<
 
   /// Computes multiply-add
   CUTLASS_HOST_DEVICE
-  void operator()(
-    FragmentC &d,
-    FragmentA const &a,
-    FragmentB const &b,
-    FragmentC const &c
-  ) const {
-
+  void operator()(FragmentC &d,
+                  FragmentA const &a,
+                  FragmentB const &b,
+                  FragmentC const &c) const {
 #if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
 
     uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
@@ -128,11 +124,20 @@ struct Mma<
     int *D = reinterpret_cast<int *>(&d);
 
     asm volatile(
-        "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
+        "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, "
+        "{%4,%5,%6,%7}, "
         "{%8,%9}, {%10,%11,%12,%13};\n"
         : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[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]));
+        : "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
     assert(0);
@@ -140,6 +145,5 @@ struct Mma<
   }
 };
 
-
-}
-}
+}  // namespace arch
+}  // namespace cutlass
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/epilogue_quant_helper.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/epilogue_quant_helper.h
index 38a253d93d..ced7ae1c1a 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/epilogue_quant_helper.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/epilogue_quant_helper.h
@@ -37,10 +37,10 @@ namespace epilogue {
 
 // define scaling mode
 enum class QuantMode {
-    PerTensorQuant,
-    PerTokenQuant,
-    PerChannelQuant,
-    PerTokenChannelQuant
+  PerTensorQuant,
+  PerTokenQuant,
+  PerChannelQuant,
+  PerTokenChannelQuant
 };
 
 }  // namespace epilogue
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale_nf4.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale_nf4.h
index 457416d75d..8f89790a3d 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale_nf4.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale_nf4.h
@@ -52,8 +52,9 @@ namespace cutlass {
 namespace epilogue {
 namespace threadblock {
 template <typename T>
-[[gnu::warning("your type here")]]
-bool print_type_1111() { return false; }
+[[gnu::warning("your type here")]] bool print_type_1111() {
+  return false;
+}
 
 template <typename ThreadblockShape_,
           int ThreadCount,
@@ -197,10 +198,12 @@ class EpilogueVisitorPerRowPerColNf4 {
         shared_storage_(shared_storage),
         extent_(problem_size),
         elementwise_(params.elementwise),
-        per_token_quant_(quant_mode == cutlass::epilogue::QuantMode::PerTokenQuant ||
-                         quant_mode == cutlass::epilogue::QuantMode::PerTokenChannelQuant),
-        per_channel_quant_(quant_mode == cutlass::epilogue::QuantMode::PerChannelQuant ||
-                           quant_mode == cutlass::epilogue::QuantMode::PerTokenChannelQuant),
+        per_token_quant_(
+            quant_mode == cutlass::epilogue::QuantMode::PerTokenQuant ||
+            quant_mode == cutlass::epilogue::QuantMode::PerTokenChannelQuant),
+        per_channel_quant_(
+            quant_mode == cutlass::epilogue::QuantMode::PerChannelQuant ||
+            quant_mode == cutlass::epilogue::QuantMode::PerTokenChannelQuant),
         ptr_alpha_row_(ptr_alpha_row),
         ptr_alpha_col_(ptr_alpha_col),
         iterator_alpha_col_(params_alpha_col,
@@ -333,7 +336,7 @@ class EpilogueVisitorPerRowPerColNf4 {
         output_converter;
     OutputVector& output =
         reinterpret_cast<OutputVector*>(&fragment_D_)[frag_idx];
-    output               = output_converter(result);
+    output = output_converter(result);
     /* // Convert to the output, with non zero C added */
     // NumericArrayConverter<ElementOutput, ElementCompute, kElementsPerAccess>
     //     output_converter;
@@ -341,7 +344,6 @@ class EpilogueVisitorPerRowPerColNf4 {
     // OutputVector& output =
     //     reinterpret_cast<OutputVector*>(&fragment_D_)[frag_idx];
 
-
     // OutputVector& vector_c =
     //     reinterpret_cast<OutputVector*>(&fragment_C_)[frag_idx];
 
@@ -353,9 +355,9 @@ class EpilogueVisitorPerRowPerColNf4 {
 
   /// Called after accumulators have been exchanged for each accumulator vector
   CUTLASS_DEVICE
-  void visit(AccumulatorFragment const &accum,
+  void visit(AccumulatorFragment const& accum,
              int reduce_fragment_idx,
-            OutputTileIterator &destination_iterator) {
+             OutputTileIterator& destination_iterator) {
     NumericArrayConverter<ElementCompute,
                           ElementAccumulator,
                           kElementsPerAccess>
@@ -368,7 +370,8 @@ class EpilogueVisitorPerRowPerColNf4 {
     ComputeFragment result = source_converter(accum);
 
     // if(threadIdx.x<32){
-    //   printf("#### %d-%d-%d--%d-%d-%d, reduced accu:%d-%d-%d-%d-%d-%d-%d-%d, dequant: accu:%f-%f-%f-%f-%f-%f-%f-%f  \n",
+    //   printf("#### %d-%d-%d--%d-%d-%d, reduced accu:%d-%d-%d-%d-%d-%d-%d-%d,
+    //   dequant: accu:%f-%f-%f-%f-%f-%f-%f-%f  \n",
     //           blockIdx.x,blockIdx.y,blockIdx.z,
     //           threadIdx.x,threadIdx.y,threadIdx.z,
     //           accum[0],
@@ -401,7 +404,8 @@ class EpilogueVisitorPerRowPerColNf4 {
     }
     // just for bug, pass
     // if(threadIdx.x<32){
-    //   printf("#### %d-%d-%d--%d-%d-%d, reduced accu:%d-%d-%d-%d-%d-%d-%d-%d, dequant: accu:%f-%f-%f-%f-%f-%f-%f-%f  \n",
+    //   printf("#### %d-%d-%d--%d-%d-%d, reduced accu:%d-%d-%d-%d-%d-%d-%d-%d,
+    //   dequant: accu:%f-%f-%f-%f-%f-%f-%f-%f  \n",
     //           blockIdx.x,blockIdx.y,blockIdx.z,
     //           threadIdx.x,threadIdx.y,threadIdx.z,
     //           accum[0],
@@ -428,7 +432,8 @@ class EpilogueVisitorPerRowPerColNf4 {
     // auto result_tmp = output_converter(result);
 
     // if(threadIdx.x<32){
-    //   printf("#### %d-%d-%d--%d-%d-%d, reduced accu:%d-%d-%d-%d-%d-%d-%d-%d, dequant: accu:%f-%f-%f-%f-%f-%f-%f-%f  \n",
+    //   printf("#### %d-%d-%d--%d-%d-%d, reduced accu:%d-%d-%d-%d-%d-%d-%d-%d,
+    //   dequant: accu:%f-%f-%f-%f-%f-%f-%f-%f  \n",
     //           blockIdx.x,blockIdx.y,blockIdx.z,
     //           threadIdx.x,threadIdx.y,threadIdx.z,
     //           accum[0],
@@ -452,8 +457,10 @@ class EpilogueVisitorPerRowPerColNf4 {
 
     typename OutputTileIterator::Fragment output_fragment;
     CUTLASS_PRAGMA_UNROLL
-    for (int ii = 0; ii<output_fragment.size(); ++ii){
-      output_fragment[ii]= static_cast<typename OutputTileIterator::Fragment::Element>(result[ii]);
+    for (int ii = 0; ii < output_fragment.size(); ++ii) {
+      output_fragment[ii] =
+          static_cast<typename OutputTileIterator::Fragment::Element>(
+              result[ii]);
     }
     // OutputVector& output =
     //     reinterpret_cast<OutputVector*>(&output_fragment)[0];
@@ -464,7 +471,8 @@ class EpilogueVisitorPerRowPerColNf4 {
     // }
 
     // if(threadIdx.x<32){
-    //   printf("#### %d-%d-%d--%d-%d-%d, reduced accu:%d-%d-%d-%d-%d-%d-%d-%d, dequant: accu:%f-%f-%f-%f-%f-%f-%f-%f  \n",
+    //   printf("#### %d-%d-%d--%d-%d-%d, reduced accu:%d-%d-%d-%d-%d-%d-%d-%d,
+    //   dequant: accu:%f-%f-%f-%f-%f-%f-%f-%f  \n",
     //           blockIdx.x,blockIdx.y,blockIdx.z,
     //           threadIdx.x,threadIdx.y,threadIdx.z,
     //           accum[0],
@@ -545,7 +553,9 @@ class EpilogueVisitorPerRowPerColNf4 {
       ComputeFragment const& scale_col,
       AlphaScaleElementType const& scale_row) {
     // if(threadIdx.x<32){
-    //   printf("#### per_token_channel_scale_accumulator,  %d-%d-%d--%d-%d-%d, quanted accu:%f-%f-%f-%f-%f-%f-%f-%f, scale_col:%f-%f-%f-%f-%f-%f-%f-%f  \n",
+    //   printf("#### per_token_channel_scale_accumulator,  %d-%d-%d--%d-%d-%d,
+    //   quanted accu:%f-%f-%f-%f-%f-%f-%f-%f, scale_col:%f-%f-%f-%f-%f-%f-%f-%f
+    //   \n",
     //          blockIdx.x,blockIdx.y,blockIdx.z,
     //          threadIdx.x,threadIdx.y,threadIdx.z,
     //          static_cast<float>(accum[0]),
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h
index 0a3c107f14..b2270332bd 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h
@@ -90,28 +90,23 @@ namespace threadblock {
 
 namespace detail {
 
-template <
-  typename ElementOutput,
-  typename ElementAccumulator,
-  int ElementsPerAccess,
-  typename ThreadblockShape,
-  typename WarpShape,
-  typename InstructionShape,
-  typename ThreadMap
->
+template <typename ElementOutput,
+          typename ElementAccumulator,
+          int ElementsPerAccess,
+          typename ThreadblockShape,
+          typename WarpShape,
+          typename InstructionShape,
+          typename ThreadMap>
 struct Nf4DefaultIteratorsTensorOp {
+  using WarpTileIterator =
+      cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
+                                                    InstructionShape,
+                                                    ElementAccumulator,
+                                                    layout::RowMajor>;
 
-  using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOp<
-    WarpShape,
-    InstructionShape,
-    ElementAccumulator,
-    layout::RowMajor
-  >;
-
-  using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
-    ThreadMap,
-    ElementAccumulator
-  >;
+  using SharedLoadIterator =
+      cutlass::epilogue::threadblock::SharedLoadIterator<ThreadMap,
+                                                         ElementAccumulator>;
 
   static int const kFragmentsPerIteration = 1;
 };
@@ -123,12 +118,12 @@ template <typename ThreadblockShape,
           typename InstructionShape,
           typename ThreadMap>
 struct Nf4DefaultIteratorsTensorOp<cutlass::half_t,
-                                int32_t,
-                                8,
-                                ThreadblockShape,
-                                WarpShape,
-                                InstructionShape,
-                                ThreadMap> {
+                                   int32_t,
+                                   8,
+                                   ThreadblockShape,
+                                   WarpShape,
+                                   InstructionShape,
+                                   ThreadMap> {
   using WarpTileIterator =
       cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
                                                     InstructionShape,
@@ -148,12 +143,12 @@ template <typename ThreadblockShape,
           typename InstructionShape,
           typename ThreadMap>
 struct Nf4DefaultIteratorsTensorOp<cutlass::bfloat16_t,
-                                int32_t,
-                                8,
-                                ThreadblockShape,
-                                WarpShape,
-                                InstructionShape,
-                                ThreadMap> {
+                                   int32_t,
+                                   8,
+                                   ThreadblockShape,
+                                   WarpShape,
+                                   InstructionShape,
+                                   ThreadMap> {
   using WarpTileIterator =
       cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
                                                     InstructionShape,
@@ -325,19 +320,14 @@ class SharedLoadIteratorMixed<ThreadMap_, int32_t, 32, 16, 8, 8> {
   void load(Fragment& frag) const { load_with_pointer_offset(frag, 0); }
 };
 
-
-
-template <
-  typename Shape_,
-  typename WarpMmaTensorOp_,
-  int PartitionsK,
-  typename OutputOp_,
-  int ElementsPerAccess,
-  bool ScatterD = false,
-  typename PermuteDLayout = layout::NoPermute
->
+template <typename Shape_,
+          typename WarpMmaTensorOp_,
+          int PartitionsK,
+          typename OutputOp_,
+          int ElementsPerAccess,
+          bool ScatterD = false,
+          typename PermuteDLayout = layout::NoPermute>
 struct DequantEpilogueTensorOp {
-
   using Shape = Shape_;
   using WarpMmaTensorOp = WarpMmaTensorOp_;
   static int const kPartitionsK = PartitionsK;
@@ -352,76 +342,76 @@ struct DequantEpilogueTensorOp {
   // Thread map
   //
 
-  using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp<
-    Shape,
-    typename WarpMmaTensorOp::Shape,
-    kPartitionsK,
-    ElementOutput,
-    kElementsPerAccess
-  >::Type;
+  using OutputTileThreadMap =
+      typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp<
+          Shape,
+          typename WarpMmaTensorOp::Shape,
+          kPartitionsK,
+          ElementOutput,
+          kElementsPerAccess>::Type;
 
-  static bool const UseCUDAStore = platform::is_same<ElementOutput, double>::value;
+  static bool const UseCUDAStore =
+      platform::is_same<ElementOutput, double>::value;
 
-  using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
-    OutputTileThreadMap,
-    ElementOutput,
-    ScatterD,
-    PermuteDLayout,
-    UseCUDAStore
-  >;
+  using OutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIterator<
+          OutputTileThreadMap,
+          ElementOutput,
+          ScatterD,
+          PermuteDLayout,
+          UseCUDAStore>;
 
-  using AccumulatorFragmentIterator = typename platform::conditional<is_complex<ElementOutput>::value,
-                                    cutlass::epilogue::warp::FragmentIteratorComplexTensorOp<
-                                        typename WarpMmaTensorOp::Shape,
-                                        typename WarpMmaTensorOp::Policy::Operator::Shape,
-                                        typename WarpMmaTensorOp::Policy::Operator::ElementC,
-                                        typename WarpMmaTensorOp::Policy::Operator::FragmentC,
-                                        LayoutC>,
-                                    cutlass::epilogue::warp::FragmentIteratorTensorOp<
-                                        typename WarpMmaTensorOp::Shape,
-                                        typename WarpMmaTensorOp::Policy::Operator::Shape,
-                                        typename WarpMmaTensorOp::Policy::Operator::ElementC,
-                                        typename WarpMmaTensorOp::Policy::Operator::FragmentC,
-                                        LayoutC> >::type;
+  using AccumulatorFragmentIterator = typename platform::conditional<
+      is_complex<ElementOutput>::value,
+      cutlass::epilogue::warp::FragmentIteratorComplexTensorOp<
+          typename WarpMmaTensorOp::Shape,
+          typename WarpMmaTensorOp::Policy::Operator::Shape,
+          typename WarpMmaTensorOp::Policy::Operator::ElementC,
+          typename WarpMmaTensorOp::Policy::Operator::FragmentC,
+          LayoutC>,
+      cutlass::epilogue::warp::FragmentIteratorTensorOp<
+          typename WarpMmaTensorOp::Shape,
+          typename WarpMmaTensorOp::Policy::Operator::Shape,
+          typename WarpMmaTensorOp::Policy::Operator::ElementC,
+          typename WarpMmaTensorOp::Policy::Operator::FragmentC,
+          LayoutC>>::type;
 
   /// Support several implementations depending on structure of epilogue
   using DefaultIterators = detail::Nf4DefaultIteratorsTensorOp<
-    ElementOutput,
-    ElementAccumulator,
-    kElementsPerAccess,
-    Shape,
-    typename WarpMmaTensorOp::Shape,
-    typename WarpMmaTensorOp::Policy::Operator::Shape,
-    typename OutputTileThreadMap::CompactedThreadMap
-  >;
+      ElementOutput,
+      ElementAccumulator,
+      kElementsPerAccess,
+      Shape,
+      typename WarpMmaTensorOp::Shape,
+      typename WarpMmaTensorOp::Policy::Operator::Shape,
+      typename OutputTileThreadMap::CompactedThreadMap>;
 
   using WarpTileIterator = typename DefaultIterators::WarpTileIterator;
   using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator;
 
   /// Hard-coded padding elements added
-  using Padding = cutlass::MatrixShape<0, 64 / sizeof_bits<ElementAccumulator>::value * 4>;
+  using Padding =
+      cutlass::MatrixShape<0, 64 / sizeof_bits<ElementAccumulator>::value * 4>;
 
-  static int const kFragmentsPerIteration = (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1);
+  static int const kFragmentsPerIteration =
+      (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1);
 
   //
   // Define the epilogue
   //
-  using Epilogue = cutlass::epilogue::threadblock::Epilogue<
-    Shape,
-    WarpMmaTensorOp,
-    kPartitionsK,
-    OutputTileIterator,
-    AccumulatorFragmentIterator,
-    WarpTileIterator,
-    SharedLoadIterator,
-    OutputOp,
-    Padding,
-    kFragmentsPerIteration
-  >;
+  using Epilogue =
+      cutlass::epilogue::threadblock::Epilogue<Shape,
+                                               WarpMmaTensorOp,
+                                               kPartitionsK,
+                                               OutputTileIterator,
+                                               AccumulatorFragmentIterator,
+                                               WarpTileIterator,
+                                               SharedLoadIterator,
+                                               OutputOp,
+                                               Padding,
+                                               kFragmentsPerIteration>;
 };
 
-
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 }  // namespace threadblock
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/ft_gemm_configs.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/ft_gemm_configs.h
index 702ff05d40..fa77599f12 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/ft_gemm_configs.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/ft_gemm_configs.h
@@ -35,39 +35,39 @@ limitations under the License. */
 //       in the kernel layout details when doing weight only quantization.
 enum class CutlassTileConfig {
   // Signals that we should run heuristics do choose a config
-  Undefined, // 0
+  Undefined,  // 0
 
   // Signals that we should run heuristics do choose a config
-  ChooseWithHeuristic, // 1
+  ChooseWithHeuristic,  // 1
 
   // SiMT config
-  CtaShape128x128x8_WarpShape64x64x8, // 2
+  CtaShape128x128x8_WarpShape64x64x8,  // 2
 
   // TensorCore configs CTA_N = 128, CTA_K = 64
   // Warp configs for M=16
-  CtaShape16x128x64_WarpShape16x32x64, // 3
-  CtaShape16x256x64_WarpShape16x64x64, // 4
+  CtaShape16x128x64_WarpShape16x32x64,  // 3
+  CtaShape16x256x64_WarpShape16x64x64,  // 4
 
   // Warp configs for M=32
-  CtaShape32x128x64_WarpShape32x32x64, // 5
+  CtaShape32x128x64_WarpShape32x32x64,  // 5
 
   // Warp configs for M=64
-  CtaShape64x128x64_WarpShape32x64x64, // 6
-  CtaShape64x128x64_WarpShape64x32x64, // 7
+  CtaShape64x128x64_WarpShape32x64x64,  // 6
+  CtaShape64x128x64_WarpShape64x32x64,  // 7
 
   // Warp configs for M=128
-  CtaShape128x128x64_WarpShape64x32x64, // 8
-  CtaShape128x128x64_WarpShape128x32x64, // 9
+  CtaShape128x128x64_WarpShape64x32x64,   // 8
+  CtaShape128x128x64_WarpShape128x32x64,  // 9
 
   // configs for large M in encoder
-  CtaShape128x256x64_WarpShape64x64x64, // 10
-  CtaShape256x128x64_WarpShape64x64x64, // 11
+  CtaShape128x256x64_WarpShape64x64x64,  // 10
+  CtaShape256x128x64_WarpShape64x64x64,  // 11
 };
 
 enum class SplitKStyle {
-  NO_SPLIT_K, //0
-  SPLIT_K_SERIAL, //1
-  SPLIT_K_STREAM, //2
+  NO_SPLIT_K,      // 0
+  SPLIT_K_SERIAL,  // 1
+  SPLIT_K_STREAM,  // 2
   // SPLIT_K_PARALLEL // Not supported yet
 };
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/default_dequant_gemm_nf4.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/default_dequant_gemm_nf4.h
index c3ec38968c..12b78b8f03 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/default_dequant_gemm_nf4.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/default_dequant_gemm_nf4.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,25 +18,27 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
 /*! \file
     \brief
-      Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with
-      the appropriate threadblock-scoped epilogue.
+      Default kernel-level GEMM definitions combine threadblock-scoped matrix
+   multiply-add with the appropriate threadblock-scoped epilogue.
 
-      Note, CUTLASS epilogues universally target row-major outputs. Column-major outputs are
-      accommodated by exchanging A and B operands and assuming transposed layouts. Partial
-      specializations here choose 'device::GemmTransposed' to implement this functionality.
+      Note, CUTLASS epilogues universally target row-major outputs. Column-major
+   outputs are accommodated by exchanging A and B operands and assuming
+   transposed layouts. Partial specializations here choose
+   'device::GemmTransposed' to implement this functionality.
 */
 
 #pragma once
@@ -101,8 +103,7 @@ template <
     /// Permute result D
     typename PermuteDLayout = layout::NoPermute,
     ///
-    typename Enable = void
->
+    typename Enable = void>
 struct DefaultDequantGemm;
 
 ///////////////////////////////////////////////////
@@ -152,40 +153,77 @@ template <
     /// Scatter result D by using an index array
     bool ScatterD,
     /// Permute result D
-    typename PermuteDLayout
->
-struct DefaultDequantGemm<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementC,
-                   LayoutC, ElementAccumulator, arch::OpClassTensorOp,
-                   arch::Sm80, ThreadblockShape, WarpShape, InstructionShape,
-                   EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial,
-                   Operator, SharedMemoryClear, GatherA, GatherB, ScatterD, PermuteDLayout> {
-
-  static_assert(platform::is_same<LayoutC, layout::RowMajor>::value
-             || platform::is_same<LayoutC, layout::AffineRankN<2>>::value,
-             "Epilogue in the kernel level must be row major");
+    typename PermuteDLayout>
+struct DefaultDequantGemm<ElementA,
+                          LayoutA,
+                          kAlignmentA,
+                          ElementB,
+                          LayoutB,
+                          kAlignmentB,
+                          ElementC,
+                          LayoutC,
+                          ElementAccumulator,
+                          arch::OpClassTensorOp,
+                          arch::Sm80,
+                          ThreadblockShape,
+                          WarpShape,
+                          InstructionShape,
+                          EpilogueOutputOp,
+                          ThreadblockSwizzle,
+                          Stages,
+                          SplitKSerial,
+                          Operator,
+                          SharedMemoryClear,
+                          GatherA,
+                          GatherB,
+                          ScatterD,
+                          PermuteDLayout> {
+  static_assert(platform::is_same<LayoutC, layout::RowMajor>::value ||
+                    platform::is_same<LayoutC, layout::AffineRankN<2>>::value,
+                "Epilogue in the kernel level must be row major");
 
   /// Define the threadblock-scoped matrix multiply-accumulate
-  using Mma = typename cutlass::gemm::threadblock::DefaultMma<
-      ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB,
-      ElementAccumulator, LayoutC, arch::OpClassTensorOp, arch::Sm80,
-      ThreadblockShape, WarpShape, InstructionShape, Stages,
-      Operator, false, SharedMemoryClear, GatherA, GatherB>::ThreadblockMma;
+  using Mma =
+      typename cutlass::gemm::threadblock::DefaultMma<ElementA,
+                                                      LayoutA,
+                                                      kAlignmentA,
+                                                      ElementB,
+                                                      LayoutB,
+                                                      kAlignmentB,
+                                                      ElementAccumulator,
+                                                      LayoutC,
+                                                      arch::OpClassTensorOp,
+                                                      arch::Sm80,
+                                                      ThreadblockShape,
+                                                      WarpShape,
+                                                      InstructionShape,
+                                                      Stages,
+                                                      Operator,
+                                                      false,
+                                                      SharedMemoryClear,
+                                                      GatherA,
+                                                      GatherB>::ThreadblockMma;
 
   static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
 
   /// Define the epilogue
-  using RegularEpilogue = typename cutlass::epilogue::threadblock::DequantEpilogueTensorOp<
-          ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp,
-          EpilogueOutputOp::kCount, ScatterD, PermuteDLayout>::Epilogue;
+  using RegularEpilogue =
+      typename cutlass::epilogue::threadblock::DequantEpilogueTensorOp<
+          ThreadblockShape,
+          typename Mma::Operator,
+          kPartitionsK,
+          EpilogueOutputOp,
+          EpilogueOutputOp::kCount,
+          ScatterD,
+          PermuteDLayout>::Epilogue;
 
   using Epilogue = RegularEpilogue;
 
   /// Define the kernel-level GEMM operator.
-  using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+  using GemmKernel =
+      kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
 };
 
-
-
 template <
     /// Element type for A matrix operand
     typename ElementA_,
@@ -235,8 +273,7 @@ template <
     /// Scatter result D by using an index array
     bool ScatterD = false,
     /// Permute result D
-    typename PermuteDLayout = layout::NoPermute
->
+    typename PermuteDLayout = layout::NoPermute>
 struct DefaultInt8InterleavedGemm;
 
 /// Partial specialization for Ampere
@@ -275,54 +312,69 @@ template <
     /// epilogue
     bool SplitKSerial,
     /// Operation performed by GEMM
-    typename Operator
->
+    typename Operator>
 struct DefaultInt8InterleavedGemm<int8_t,
-                   LayoutA,
-                   kAlignmentA,
-                   int8_t,
-                   LayoutB,
-                   kAlignmentB,
-                   ElementC,
-                   LayoutC,
-                   ElementAccumulator,
-                   OperatorClass,
-                   ArchTag,
-                   ThreadblockShape,
-                   WarpShape,
-                   InstructionShape,
-                   EpilogueOutputOp,
-                   ThreadblockSwizzle,
-                   Stages,
-                   SplitKSerial,
-                   Operator> {
-
+                                  LayoutA,
+                                  kAlignmentA,
+                                  int8_t,
+                                  LayoutB,
+                                  kAlignmentB,
+                                  ElementC,
+                                  LayoutC,
+                                  ElementAccumulator,
+                                  OperatorClass,
+                                  ArchTag,
+                                  ThreadblockShape,
+                                  WarpShape,
+                                  InstructionShape,
+                                  EpilogueOutputOp,
+                                  ThreadblockSwizzle,
+                                  Stages,
+                                  SplitKSerial,
+                                  Operator> {
   static_assert(platform::is_same<LayoutC, layout::RowMajor>::value,
-             "Epilogue in the kernel level must be row major");
+                "Epilogue in the kernel level must be row major");
 
   /// Define the threadblock-scoped matrix multiply-accumulate
   using Mma = typename cutlass::gemm::threadblock::DefaulInt8Nf4InterleavedMma<
-      int8_t, LayoutA, kAlignmentA, int8_t, LayoutB, kAlignmentB,
-      ElementAccumulator, LayoutC, OperatorClass, ArchTag,
-      ThreadblockShape, WarpShape, InstructionShape, Stages,
-      Operator, false, SharedMemoryClearOption::kNone>::ThreadblockMma;
-
-
+      int8_t,
+      LayoutA,
+      kAlignmentA,
+      int8_t,
+      LayoutB,
+      kAlignmentB,
+      ElementAccumulator,
+      LayoutC,
+      OperatorClass,
+      ArchTag,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      Stages,
+      Operator,
+      false,
+      SharedMemoryClearOption::kNone>::ThreadblockMma;
 
   static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
 
   /// Define the epilogue
-  using RegularEpilogue = typename cutlass::epilogue::threadblock::DequantEpilogueTensorOp<
-          ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp,
-          EpilogueOutputOp::kCount, false, layout::NoPermute>::Epilogue;
+  using RegularEpilogue =
+      typename cutlass::epilogue::threadblock::DequantEpilogueTensorOp<
+          ThreadblockShape,
+          typename Mma::Operator,
+          kPartitionsK,
+          EpilogueOutputOp,
+          EpilogueOutputOp::kCount,
+          false,
+          layout::NoPermute>::Epilogue;
 
   using Epilogue = RegularEpilogue;
 
   /// Define the kernel-level GEMM operator.
-  using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+  using GemmKernel =
+      kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
 };
 
-
 /// Partial specialization for Ampere
 template <
     /// Layout type for A matrix operand
@@ -359,51 +411,67 @@ template <
     /// epilogue
     bool SplitKSerial,
     /// Operation performed by GEMM
-    typename Operator
->
+    typename Operator>
 struct DefaultInt8InterleavedGemm<int8_t,
-                   LayoutA,
-                   kAlignmentA,
-                   cutlass::uint4b_t,
-                   LayoutB,
-                   kAlignmentB,
-                   ElementC,
-                   LayoutC,
-                   ElementAccumulator,
-                   OperatorClass,
-                   ArchTag,
-                   ThreadblockShape,
-                   WarpShape,
-                   InstructionShape,
-                   EpilogueOutputOp,
-                   ThreadblockSwizzle,
-                   Stages,
-                   SplitKSerial,
-                   Operator> {
-
+                                  LayoutA,
+                                  kAlignmentA,
+                                  cutlass::uint4b_t,
+                                  LayoutB,
+                                  kAlignmentB,
+                                  ElementC,
+                                  LayoutC,
+                                  ElementAccumulator,
+                                  OperatorClass,
+                                  ArchTag,
+                                  ThreadblockShape,
+                                  WarpShape,
+                                  InstructionShape,
+                                  EpilogueOutputOp,
+                                  ThreadblockSwizzle,
+                                  Stages,
+                                  SplitKSerial,
+                                  Operator> {
   static_assert(platform::is_same<LayoutC, layout::RowMajor>::value,
-             "Epilogue in the kernel level must be row major");
+                "Epilogue in the kernel level must be row major");
 
   /// Define the threadblock-scoped matrix multiply-accumulate
   using Mma = typename cutlass::gemm::threadblock::DefaulInt8Nf4InterleavedMma<
-      int8_t, LayoutA, kAlignmentA, cutlass::uint4b_t, LayoutB, kAlignmentB,
-      ElementAccumulator, LayoutC, OperatorClass, ArchTag,
-      ThreadblockShape, WarpShape, InstructionShape, Stages,
-      Operator, false, SharedMemoryClearOption::kNone>::ThreadblockMma;
-
-
+      int8_t,
+      LayoutA,
+      kAlignmentA,
+      cutlass::uint4b_t,
+      LayoutB,
+      kAlignmentB,
+      ElementAccumulator,
+      LayoutC,
+      OperatorClass,
+      ArchTag,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      Stages,
+      Operator,
+      false,
+      SharedMemoryClearOption::kNone>::ThreadblockMma;
 
   static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
 
   /// Define the epilogue
-  using RegularEpilogue = typename cutlass::epilogue::threadblock::DequantEpilogueTensorOp<
-          ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp,
-          EpilogueOutputOp::kCount, false, layout::NoPermute>::Epilogue;
+  using RegularEpilogue =
+      typename cutlass::epilogue::threadblock::DequantEpilogueTensorOp<
+          ThreadblockShape,
+          typename Mma::Operator,
+          kPartitionsK,
+          EpilogueOutputOp,
+          EpilogueOutputOp::kCount,
+          false,
+          layout::NoPermute>::Epilogue;
 
   using Epilogue = RegularEpilogue;
 
   /// Define the kernel-level GEMM operator.
-  using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+  using GemmKernel =
+      kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
 };
 ////////////////////////////////////////////////////////////////////////////////
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/default_intA_nf4B_traits.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/default_intA_nf4B_traits.h
index 2a13073f1f..9c69c037f8 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/default_intA_nf4B_traits.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/default_intA_nf4B_traits.h
@@ -44,56 +44,63 @@ namespace cutlass {
 namespace gemm {
 namespace kernel {
 
-template<typename InAType, typename IntBType, typename OutType, typename arch, typename Enable = void>
-struct Int8Nf4GemmArchTraits {
-};
+template <typename InAType,
+          typename IntBType,
+          typename OutType,
+          typename arch,
+          typename Enable = void>
+struct Int8Nf4GemmArchTraits {};
 
-template<typename arch>
+template <typename arch>
 struct Int8Nf4GemmArchTraits<float, float, float, arch> {
-    static constexpr int Stages = 2;
-    using OperatorClass         = cutlass::arch::OpClassSimt;
-    using AccType               = float;
-    using LayoutB               = cutlass::layout::RowMajor;
+  static constexpr int Stages = 2;
+  using OperatorClass = cutlass::arch::OpClassSimt;
+  using AccType = float;
+  using LayoutB = cutlass::layout::RowMajor;
 
-    static constexpr int ElementsPerAccessA = 1;
-    static constexpr int ElementsPerAccessB = 1;
-    static constexpr int ElementsPerAccessC = 1;
-    static constexpr int ThreadblockK       = 8;
-    using InstructionShape                  = cutlass::gemm::GemmShape<1, 1, 1>;
+  static constexpr int ElementsPerAccessA = 1;
+  static constexpr int ElementsPerAccessB = 1;
+  static constexpr int ElementsPerAccessC = 1;
+  static constexpr int ThreadblockK = 8;
+  using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>;
 
-    using Operator = cutlass::arch::OpMultiplyAdd;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
 // ======================= Ampere Traits ==============================
-template<typename IntAType, typename IntBType, typename OutType>
-struct Int8Nf4GemmArchTraits<IntAType,
-                          IntBType,
-                          OutType,
-                          cutlass::arch::Sm80,
-                          typename cutlass::platform::enable_if<
-                            cutlass::platform::is_same<IntBType, int8_t>::value ||
-                            cutlass::platform::is_same<IntBType, cutlass::uint4b_t>::value>::type> {
-private:
-    using LayoutDetails = LayoutDetailsB<IntBType, cutlass::arch::Sm80>;
+template <typename IntAType, typename IntBType, typename OutType>
+struct Int8Nf4GemmArchTraits<
+    IntAType,
+    IntBType,
+    OutType,
+    cutlass::arch::Sm80,
+    typename cutlass::platform::enable_if<
+        cutlass::platform::is_same<IntBType, int8_t>::value ||
+        cutlass::platform::is_same<IntBType, cutlass::uint4b_t>::value>::type> {
+ private:
+  using LayoutDetails = LayoutDetailsB<IntBType, cutlass::arch::Sm80>;
 
-public:
-    static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
+ public:
+  static constexpr int ThreadblockK = LayoutDetails::ThreadblockK;
 
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using AccType       = int32_t;
-    using LayoutB       = typename LayoutDetails::Layout;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using AccType = int32_t;
+  using LayoutB = typename LayoutDetails::Layout;
 
-    static constexpr int ElementsPerAccessA = 128 / cutlass::sizeof_bits<IntAType>::value;
-    static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
-    static constexpr int ElementsPerAccessC = 128 / cutlass::sizeof_bits<OutType>::value;
-    // static_assert(cutlass::platform::is_same<InType, int8_t>::value,
-    //               "input type must be int8_t");
-    // static_assert((ElementsPerAccessA == 16), "=====");
-    // static_assert((ElementsPerAccessB == 16), "=====");
-    // static_assert((ElementsPerAccessC == 8), "=====");
-    using InstructionShape                  = cutlass::gemm::GemmShape<16, 8, 32>;
-    // using InstructionShape                  = cutlass::gemm::GemmShape<16, 8, 16>;
-    using Operator = typename LayoutDetails::Operator;
+  static constexpr int ElementsPerAccessA =
+      128 / cutlass::sizeof_bits<IntAType>::value;
+  static constexpr int ElementsPerAccessB = LayoutDetails::ElementsPerAccess;
+  static constexpr int ElementsPerAccessC =
+      128 / cutlass::sizeof_bits<OutType>::value;
+  // static_assert(cutlass::platform::is_same<InType, int8_t>::value,
+  //               "input type must be int8_t");
+  // static_assert((ElementsPerAccessA == 16), "=====");
+  // static_assert((ElementsPerAccessB == 16), "=====");
+  // static_assert((ElementsPerAccessC == 8), "=====");
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
+  // using InstructionShape                  = cutlass::gemm::GemmShape<16, 8,
+  // 16>;
+  using Operator = typename LayoutDetails::Operator;
 };
 
 }  // namespace kernel
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor_interleaved_nf4.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor_interleaved_nf4.h
index 574f9f2056..922e64d6c6 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor_interleaved_nf4.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor_interleaved_nf4.h
@@ -290,7 +290,8 @@ struct GemmWithEpilogueVisitorInterleavedNf4 {
       //       ceil_div(args.problem_size.k(), args.batch_count), kAlignK);
 
       //   if (gemm_k_size) {
-      //     grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size);
+      //     grid_tiled_shape.k() = ceil_div(args.problem_size.k(),
+      //     gemm_k_size);
       //   }
       // }
 
@@ -318,7 +319,8 @@ struct GemmWithEpilogueVisitorInterleavedNf4 {
 
   /// Determines whether kernel satisfies alignment
   static Status can_implement(cutlass::gemm::GemmCoord const& problem_size) {
-    CUTLASS_TRACE_HOST("GemmWithEpilogueVisitorInterleavedNf4::can_implement()");
+    CUTLASS_TRACE_HOST(
+        "GemmWithEpilogueVisitorInterleavedNf4::can_implement()");
 
     static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
     static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
@@ -440,16 +442,17 @@ struct GemmWithEpilogueVisitorInterleavedNf4 {
           params.ptr_B)[threadblock_tile_offset.k()];
     }
 #endif
-  // if(threadIdx.x==0){
-  //     printf("##### block: %d-%d-%d, offset_k:%d, threadblock_tile_offset.m-n-k():%d-%d-%d, params.gemm_k_size:%d \n",
-  //             blockIdx.x, blockIdx.y, blockIdx.z,
-  //             offset_k,
-  //             threadblock_tile_offset.m(),
-  //             threadblock_tile_offset.n(),
-  //             threadblock_tile_offset.k(),
-  //             params.gemm_k_size
-  //             );
-  //   }
+    // if(threadIdx.x==0){
+    //     printf("##### block: %d-%d-%d, offset_k:%d,
+    //     threadblock_tile_offset.m-n-k():%d-%d-%d, params.gemm_k_size:%d \n",
+    //             blockIdx.x, blockIdx.y, blockIdx.z,
+    //             offset_k,
+    //             threadblock_tile_offset.m(),
+    //             threadblock_tile_offset.n(),
+    //             threadblock_tile_offset.k(),
+    //             params.gemm_k_size
+    //             );
+    //   }
 
     // Compute initial location in logical coordinates
     cutlass::MatrixCoord tb_offset_A{
@@ -460,7 +463,8 @@ struct GemmWithEpilogueVisitorInterleavedNf4 {
     // cutlass::MatrixCoord tb_offset_B{offset_k, threadblock_tile_offset.n() *
     // Mma::Shape::kN};
     // printf("#### kInterleave:%d \n", kInterleave);
-    // printf("###### offset_k : %d; params.gemm_k_size:%d; threadblock_tile_offset.k():%d \n",
+    // printf("###### offset_k : %d; params.gemm_k_size:%d;
+    // threadblock_tile_offset.k():%d \n",
     //         offset_k,
     //         params.gemm_k_size,
     //         threadblock_tile_offset.k()
@@ -470,7 +474,8 @@ struct GemmWithEpilogueVisitorInterleavedNf4 {
         offset_k * kInterleave,
         threadblock_tile_offset.n() * Mma::Shape::kN / kInterleave};
     // if(threadIdx.x==0){
-    //   printf("##### block: %d-%d-%d, tb_offset_B:%d-%d, kInterleave:%d, Mma::IteratorB::Shape::kRow:%d, Mma::Shape::kK:%d \n",
+    //   printf("##### block: %d-%d-%d, tb_offset_B:%d-%d, kInterleave:%d,
+    //   Mma::IteratorB::Shape::kRow:%d, Mma::Shape::kK:%d \n",
     //           blockIdx.x, blockIdx.y, blockIdx.z,
     //           offset_k * kInterleave,
     //           threadblock_tile_offset.n() * Mma::Shape::kN / kInterleave,
@@ -500,13 +505,11 @@ struct GemmWithEpilogueVisitorInterleavedNf4 {
         thread_idx,
         tb_offset_B);
     typename Mma::IteratorNF4LookUpTable iterator_nf4_look_up_table =
-      Mma::IteratorNF4LookUpTable(
-        params.params_nf4_look_up_table,
-      params.ref_nf4_look_up_table.data(),
-        {0,16},
-        threadIdx.x,
-        {0,0}
-      );
+        Mma::IteratorNF4LookUpTable(params.params_nf4_look_up_table,
+                                    params.ref_nf4_look_up_table.data(),
+                                    {0, 16},
+                                    threadIdx.x,
+                                    {0, 0});
 
     // Broadcast the warp_id computed by lane 0 to ensure dependent code
     // is compiled as warp-uniform.
@@ -530,7 +533,12 @@ struct GemmWithEpilogueVisitorInterleavedNf4 {
         (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
     // printf("#### gemm_k_iterations: %d \n", gemm_k_iterations);
     // Compute threadblock-scoped matrix multiply-add
-    mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, iterator_nf4_look_up_table, accumulators);
+    mma(gemm_k_iterations,
+        accumulators,
+        iterator_A,
+        iterator_B,
+        iterator_nf4_look_up_table,
+        accumulators);
     // if(threadIdx.x==0){
     //   printf("##### block: %d-%d-%d, offset-m-n-k:%d-%d-%d \n",
     //           blockIdx.x, blockIdx.y, blockIdx.z,
@@ -552,7 +560,8 @@ struct GemmWithEpilogueVisitorInterleavedNf4 {
         threadblock_tile_offset.n() * Mma::Shape::kN);
 
     // int block_idx = threadblock_tile_offset.m() +
-    //                 threadblock_tile_offset.n() * params.grid_tiled_shape.m();
+    //                 threadblock_tile_offset.n() *
+    //                 params.grid_tiled_shape.m();
 
     //
     // Construct the epilogue visitor
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h
index 35469f7ad9..b09368a848 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h
@@ -29,9 +29,10 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 /*
-  This file exists so that we use the same weight layout for MoE grouped gemm and regular gemm when the weight is
-  quantized. The preprocessing code reads this template to know how to organize the quantized weight matrices
-  to be consumed by CUTLASS.
+  This file exists so that we use the same weight layout for MoE grouped gemm
+  and regular gemm when the weight is quantized. The preprocessing code reads
+  this template to know how to organize the quantized weight matrices to be
+  consumed by CUTLASS.
 
   Note that for int4, ThreadBlockK MUST be 64.
 
@@ -53,79 +54,108 @@ namespace cutlass {
 namespace gemm {
 namespace kernel {
 
-template<typename TypeB, typename Arch, typename Enable = void>
-struct LayoutDetailsB {
-};
+template <typename TypeB, typename Arch, typename Enable = void>
+struct LayoutDetailsB {};
 
-// // Volta specialiations. Volta will dequantize before STS, so we need a different operator
-template<typename TypeB>
+// // Volta specialiations. Volta will dequantize before STS, so we need a
+// different operator
+template <typename TypeB>
 struct LayoutDetailsB<TypeB, arch::Sm70> {
-    static constexpr int ThreadblockK      = 64;
-    using Layout                           = layout::RowMajor;
-    static constexpr int ElementsPerAccess = 8;
-    using Operator                         = cutlass::arch::OpMultiplyAdd;
+  static constexpr int ThreadblockK = 64;
+  using Layout = layout::RowMajor;
+  static constexpr int ElementsPerAccess = 8;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
-// Specializations for Turing+ when B is FP16. These are currently only used for MoE networks.
-// TODO - Switch this to column major for weights since gemms should be more performant.
-template<typename Arch>
-struct LayoutDetailsB<half_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
-    static constexpr int ThreadblockK      = 64;
-    using Layout                           = layout::RowMajor;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<half_t>::value;
-    using Operator                         = cutlass::arch::OpMultiplyAdd;
+// Specializations for Turing+ when B is FP16. These are currently only used for
+// MoE networks.
+// TODO - Switch this to column major for weights since gemms should be more
+// performant.
+template <typename Arch>
+struct LayoutDetailsB<
+    half_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK = 64;
+  using Layout = layout::RowMajor;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<half_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
-template<typename Arch>
-struct LayoutDetailsB<bfloat16_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
-    static constexpr int ThreadblockK      = 64;
-    using Layout                           = layout::RowMajor;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<bfloat16_t>::value;
-    using Operator                         = cutlass::arch::OpMultiplyAdd;
+template <typename Arch>
+struct LayoutDetailsB<
+    bfloat16_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK = 64;
+  using Layout = layout::RowMajor;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<bfloat16_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAdd;
 };
 
-// Specializations for Turing+ when B is quantized. These can use the operator OpMultiplyAddDequantizeInterleavedBToA,
-// which signals that we want to dequantize after loading from smem.
-template<typename Arch>
-struct LayoutDetailsB<uint8_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
-    static constexpr int ThreadblockK = 64;
+// Specializations for Turing+ when B is quantized. These can use the operator
+// OpMultiplyAddDequantizeInterleavedBToA, which signals that we want to
+// dequantize after loading from smem.
+template <typename Arch>
+struct LayoutDetailsB<
+    uint8_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK = 64;
 
-private:
-    static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits<uint8_t>::value;
-    static constexpr int ColumnsInterleaved   = ElementsPerCacheLine / ThreadblockK;
+ private:
+  static constexpr int ElementsPerCacheLine =
+      128 * 8 / sizeof_bits<uint8_t>::value;
+  static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
 
-public:
-    using Layout                           = layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<uint8_t>::value;
-    using Operator                         = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
+ public:
+  using Layout =
+      layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<uint8_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
 };
 
-template<typename Arch>
-struct LayoutDetailsB<uint4b_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
-    static constexpr int ThreadblockK = 64;
+template <typename Arch>
+struct LayoutDetailsB<
+    uint4b_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK = 64;
 
-private:
-    static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits<uint4b_t>::value;
-    static constexpr int ColumnsInterleaved   = ElementsPerCacheLine / ThreadblockK;
+ private:
+  static constexpr int ElementsPerCacheLine =
+      128 * 8 / sizeof_bits<uint4b_t>::value;
+  static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
 
-public:
-    using Layout                           = layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<uint4b_t>::value;
-    using Operator                         = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
+ public:
+  using Layout =
+      layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<uint4b_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
 };
 // For int8 int8 int32 Gemm. Author(zhengzekang)
-template<typename Arch>
-struct LayoutDetailsB<int8_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
-    static constexpr int ThreadblockK = 64;
+template <typename Arch>
+struct LayoutDetailsB<
+    int8_t,
+    Arch,
+    typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK = 64;
 
-private:
-    static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits<int8_t>::value;
-    static constexpr int ColumnsInterleaved   = ElementsPerCacheLine / ThreadblockK;
+ private:
+  static constexpr int ElementsPerCacheLine =
+      128 * 8 / sizeof_bits<int8_t>::value;
+  static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
 
-public:
-    using Layout                           = layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
-    static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<int8_t>::value;
-    using Operator                         = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
+ public:
+  using Layout =
+      layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
+  static constexpr int ElementsPerAccess =
+      128 / cutlass::sizeof_bits<int8_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
 };
 
 }  // namespace kernel
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_mma_nf4_int8_interleaved.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_mma_nf4_int8_interleaved.h
index 540c9a0931..bc8ab30bd1 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_mma_nf4_int8_interleaved.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_mma_nf4_int8_interleaved.h
@@ -74,12 +74,11 @@ template <
     /// Gather operand A by using an index array
     bool GatherA = false,
     /// Gather operand B by using an index array
-    bool GatherB = false
-    >
+    bool GatherB = false>
 struct DefaulInt8Nf4InterleavedMma;
 
 // int8 int8 int32 Gemm specialization. Author(zhengzekang)
-template<
+template <
     /// Layout type for A matrix operand
     typename LayoutA,
     /// Access granularity of A matrix in units of elements
@@ -101,24 +100,6 @@ template<
     /// Operation performed by GEMM
     typename Operator>
 struct DefaulInt8Nf4InterleavedMma<int8_t,
-                  LayoutA,
-                  kAlignmentA,
-                  int8_t,
-                  LayoutB,
-                  kAlignmentB,
-                  ElementAccumulator,
-                  layout::RowMajor,
-                  arch::OpClassTensorOp,
-                  ArchTag,
-                  ThreadblockShape,
-                  WarpShape,
-                  InstructionShape,
-                  2,
-                  Operator> {
-
-private:
-
-    using Mma = Int8Nf4InterleavedMma<int8_t,
                                    LayoutA,
                                    kAlignmentA,
                                    int8_t,
@@ -132,23 +113,39 @@ private:
                                    WarpShape,
                                    InstructionShape,
                                    2,
-                                   Operator>;
+                                   Operator> {
+ private:
+  using Mma = Int8Nf4InterleavedMma<int8_t,
+                                    LayoutA,
+                                    kAlignmentA,
+                                    int8_t,
+                                    LayoutB,
+                                    kAlignmentB,
+                                    ElementAccumulator,
+                                    layout::RowMajor,
+                                    arch::OpClassTensorOp,
+                                    ArchTag,
+                                    ThreadblockShape,
+                                    WarpShape,
+                                    InstructionShape,
+                                    2,
+                                    Operator>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
-template<
+template <
     /// Layout type for A matrix operand
     typename LayoutA,
     /// Access granularity of A matrix in units of elements
@@ -174,26 +171,6 @@ template<
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
 struct DefaulInt8Nf4InterleavedMma<int8_t,
-                  LayoutA,
-                  kAlignmentA,
-                  int8_t,
-                  LayoutB,
-                  kAlignmentB,
-                  ElementAccumulator,
-                  layout::RowMajor,
-                  arch::OpClassTensorOp,
-                  ArchTag,
-                  ThreadblockShape,
-                  WarpShape,
-                  InstructionShape,
-                  kStages,
-                  Operator,
-                  false,
-                  SharedMemoryClear> {
-
-private:
-
-    using Mma = Int8Nf4InterleavedMma<int8_t,
                                    LayoutA,
                                    kAlignmentA,
                                    int8_t,
@@ -208,27 +185,42 @@ private:
                                    InstructionShape,
                                    kStages,
                                    Operator,
-                                   SharedMemoryClear>;
+                                   false,
+                                   SharedMemoryClear> {
+ private:
+  using Mma = Int8Nf4InterleavedMma<int8_t,
+                                    LayoutA,
+                                    kAlignmentA,
+                                    int8_t,
+                                    LayoutB,
+                                    kAlignmentB,
+                                    ElementAccumulator,
+                                    layout::RowMajor,
+                                    arch::OpClassTensorOp,
+                                    ArchTag,
+                                    ThreadblockShape,
+                                    WarpShape,
+                                    InstructionShape,
+                                    kStages,
+                                    Operator,
+                                    SharedMemoryClear>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
-
-
-
 // int8 int4 int32
-template<
+template <
     /// Layout type for A matrix operand
     typename LayoutA,
     /// Access granularity of A matrix in units of elements
@@ -250,24 +242,6 @@ template<
     /// Operation performed by GEMM
     typename Operator>
 struct DefaulInt8Nf4InterleavedMma<int8_t,
-                  LayoutA,
-                  kAlignmentA,
-                  cutlass::uint4b_t,
-                  LayoutB,
-                  kAlignmentB,
-                  ElementAccumulator,
-                  layout::RowMajor,
-                  arch::OpClassTensorOp,
-                  ArchTag,
-                  ThreadblockShape,
-                  WarpShape,
-                  InstructionShape,
-                  2,
-                  Operator> {
-
-private:
-
-    using Mma = Int8Nf4InterleavedMma<int8_t,
                                    LayoutA,
                                    kAlignmentA,
                                    cutlass::uint4b_t,
@@ -281,23 +255,39 @@ private:
                                    WarpShape,
                                    InstructionShape,
                                    2,
-                                   Operator>;
+                                   Operator> {
+ private:
+  using Mma = Int8Nf4InterleavedMma<int8_t,
+                                    LayoutA,
+                                    kAlignmentA,
+                                    cutlass::uint4b_t,
+                                    LayoutB,
+                                    kAlignmentB,
+                                    ElementAccumulator,
+                                    layout::RowMajor,
+                                    arch::OpClassTensorOp,
+                                    ArchTag,
+                                    ThreadblockShape,
+                                    WarpShape,
+                                    InstructionShape,
+                                    2,
+                                    Operator>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
-template<
+template <
     /// Layout type for A matrix operand
     typename LayoutA,
     /// Access granularity of A matrix in units of elements
@@ -323,26 +313,6 @@ template<
     /// Shared memory clear option
     SharedMemoryClearOption SharedMemoryClear>
 struct DefaulInt8Nf4InterleavedMma<int8_t,
-                  LayoutA,
-                  kAlignmentA,
-                  cutlass::uint4b_t,
-                  LayoutB,
-                  kAlignmentB,
-                  ElementAccumulator,
-                  layout::RowMajor,
-                  arch::OpClassTensorOp,
-                  ArchTag,
-                  ThreadblockShape,
-                  WarpShape,
-                  InstructionShape,
-                  kStages,
-                  Operator,
-                  false,
-                  SharedMemoryClear> {
-
-private:
-
-    using Mma = Int8Nf4InterleavedMma<int8_t,
                                    LayoutA,
                                    kAlignmentA,
                                    cutlass::uint4b_t,
@@ -357,20 +327,38 @@ private:
                                    InstructionShape,
                                    kStages,
                                    Operator,
-                                   SharedMemoryClear>;
+                                   false,
+                                   SharedMemoryClear> {
+ private:
+  using Mma = Int8Nf4InterleavedMma<int8_t,
+                                    LayoutA,
+                                    kAlignmentA,
+                                    cutlass::uint4b_t,
+                                    LayoutB,
+                                    kAlignmentB,
+                                    ElementAccumulator,
+                                    layout::RowMajor,
+                                    arch::OpClassTensorOp,
+                                    ArchTag,
+                                    ThreadblockShape,
+                                    WarpShape,
+                                    InstructionShape,
+                                    kStages,
+                                    Operator,
+                                    SharedMemoryClear>;
 
-public:
-    // Define the MmaCore components
-    using MmaCore = typename Mma::MmaCore;
+ public:
+  // Define the MmaCore components
+  using MmaCore = typename Mma::MmaCore;
 
-    // Define iterators over tiles from the A operand
-    using IteratorA = typename Mma::IteratorA;
+  // Define iterators over tiles from the A operand
+  using IteratorA = typename Mma::IteratorA;
 
-    // Define iterators over tiles from the B operand
-    using IteratorB = typename Mma::IteratorB;
+  // Define iterators over tiles from the B operand
+  using IteratorB = typename Mma::IteratorB;
 
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = typename Mma::ThreadblockMma;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
 }  // namespace threadblock
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_nf4_int8_interleaved_mma.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_nf4_int8_interleaved_mma.h
index e3cb9da1bf..5c48abd7cc 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_nf4_int8_interleaved_mma.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_nf4_int8_interleaved_mma.h
@@ -38,57 +38,67 @@ namespace gemm {
 namespace threadblock {
 ////////////////////////////////////////////////////////////////////////////////
 
-// // We need to distinguish here, since we want volta support. It is too much effort
-// // to write shared memory iterators that are probably needed for volta to function
-// // properly. As a result, we allow converters both after the LDG (for volta) and after
+// // We need to distinguish here, since we want volta support. It is too much
+// effort
+// // to write shared memory iterators that are probably needed for volta to
+// function
+// // properly. As a result, we allow converters both after the LDG (for volta)
+// and after
 // // the LDS for Turing+.
-template<
+template <
     /// Iterator for B matrix in global memory
     typename IteratorB,
     /// Warp level Mma
     typename MmaOperator,
     /// Math operation perform by warp level operator
     typename MathOperator>
-struct SetConvertersInt8Nf4Interleaved {
-};
+struct SetConvertersInt8Nf4Interleaved {};
 
 // // Dequantize after LDG, so set transforms accordingly
-template<
+template <
     /// Iterator for B matrix in global memory
     typename IteratorB,
     /// Mma Policy
     typename MmaOperator>
-struct SetConvertersInt8Nf4Interleaved<IteratorB, MmaOperator, arch::OpMultiplyAdd> {
-    using TransformAfterLDG =
-        FastInterleavedAndBiasedNumericArrayConverterNf4<typename MmaOperator::ArchMmaOperator::ElementB,
-                                                      typename IteratorB::Element,
-                                                      IteratorB::Fragment::kElements>;
+struct SetConvertersInt8Nf4Interleaved<IteratorB,
+                                       MmaOperator,
+                                       arch::OpMultiplyAdd> {
+  using TransformAfterLDG = FastInterleavedAndBiasedNumericArrayConverterNf4<
+      typename MmaOperator::ArchMmaOperator::ElementB,
+      typename IteratorB::Element,
+      IteratorB::Fragment::kElements>;
 
-    using TransformAfterLDS = NumericArrayConverter<typename MmaOperator::ArchMmaOperator::ElementB,
-                                                    typename MmaOperator::ArchMmaOperator::ElementB,
-                                                    MmaOperator::FragmentB::kElements>;
+  using TransformAfterLDS =
+      NumericArrayConverter<typename MmaOperator::ArchMmaOperator::ElementB,
+                            typename MmaOperator::ArchMmaOperator::ElementB,
+                            MmaOperator::FragmentB::kElements>;
 };
 
 // Dequantize after LDS, so set transforms accordingly
 
-template<
+template <
     /// Iterator for B matrix in global memory
     typename IteratorB,
     /// Mma Policy
     typename MmaOperator>
-struct SetConvertersInt8Nf4Interleaved<IteratorB, MmaOperator, arch::OpMultiplyAddDequantizeInterleavedBToA> {
-    using TransformAfterLDG =
-        NumericArrayConverter<typename IteratorB::Element, typename IteratorB::Element, IteratorB::Fragment::kElements>;
+struct SetConvertersInt8Nf4Interleaved<
+    IteratorB,
+    MmaOperator,
+    arch::OpMultiplyAddDequantizeInterleavedBToA> {
+  using TransformAfterLDG =
+      NumericArrayConverter<typename IteratorB::Element,
+                            typename IteratorB::Element,
+                            IteratorB::Fragment::kElements>;
 
-    using TransformAfterLDS =
-        FastInterleavedAndBiasedNumericArrayConverterNf4<typename MmaOperator::ArchMmaOperator::ElementB,
-                                                      typename TransformAfterLDG::result_type::Element,
-                                                      MmaOperator::FragmentB::kElements>;
+  using TransformAfterLDS = FastInterleavedAndBiasedNumericArrayConverterNf4<
+      typename MmaOperator::ArchMmaOperator::ElementB,
+      typename TransformAfterLDG::result_type::Element,
+      MmaOperator::FragmentB::kElements>;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template<
+template <
     /// Element type for A matrix operand
     typename ElementA_,
     /// Layout type for A matrix operand
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_nf4_int8_interleaved_mma_multistage.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_nf4_int8_interleaved_mma_multistage.h
index b2bc5d2511..aa3a04d9ac 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_nf4_int8_interleaved_mma_multistage.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/default_nf4_int8_interleaved_mma_multistage.h
@@ -46,7 +46,7 @@ namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template<
+template <
     /// Type for elementA
     typename ElementA,
     /// Layout type for A matrix operand
@@ -78,120 +78,134 @@ template<
     ///
     SharedMemoryClearOption SharedMemoryClear>
 struct Int8Nf4InterleavedMma<ElementA,
-             LayoutA,
-             kAlignmentA,
-             ElementB,
-             LayoutB,
-             kAlignmentB,
-             ElementAccumulator,
-             layout::RowMajor,
-             OperatorClass,
-             ArchTag,
-             ThreadblockShape,
-             WarpShape,
-             InstructionShape,
-             kStages,
-             Operator,
-             SharedMemoryClear,
-             typename platform::enable_if<(ArchTag::kMinComputeCapability >= 80)>::type> {
+                             LayoutA,
+                             kAlignmentA,
+                             ElementB,
+                             LayoutB,
+                             kAlignmentB,
+                             ElementAccumulator,
+                             layout::RowMajor,
+                             OperatorClass,
+                             ArchTag,
+                             ThreadblockShape,
+                             WarpShape,
+                             InstructionShape,
+                             kStages,
+                             Operator,
+                             SharedMemoryClear,
+                             typename platform::enable_if<(
+                                 ArchTag::kMinComputeCapability >= 80)>::type> {
+  static_assert(platform::is_same<ElementA, int8_t>::value,
+                "Element A must be in8t");
 
-    static_assert(platform::is_same<ElementA, int8_t>::value,
-                  "Element A must be in8t");
+  static_assert(
+      platform::is_same<Operator,
+                        arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
+      "Mma multistage must dequantize after ldsm");
 
-    static_assert(platform::is_same<Operator, arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
-                  "Mma multistage must dequantize after ldsm");
+  static_assert(platform::is_same<ElementB, int8_t>::value ||
+                    platform::is_same<ElementB, uint4b_t>::value,
+                "Element B must be int8 or uint4");
+  static_assert(WarpShape::kK != 0, "");
+  static_assert(ThreadblockShape::kK != 0, "");
 
-    static_assert(platform::is_same<ElementB, int8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-                  "Element B must be int8 or uint4");
-    static_assert(WarpShape::kK!=0,"");
-    static_assert(ThreadblockShape::kK!=0,"");
+  static cutlass::arch::CacheOperation::Kind const CacheOpA =
+      ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits<ElementA>::value * kAlignmentA) == 128) ?
-                                                                    cutlass::arch::CacheOperation::Global :
-                                                                    cutlass::arch::CacheOperation::Always;
+  static cutlass::arch::CacheOperation::Kind const CacheOpB =
+      ((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpB = ((sizeof_bits<ElementB>::value * kAlignmentB) == 128) ?
-                                                                    cutlass::arch::CacheOperation::Global :
-                                                                    cutlass::arch::CacheOperation::Always;
+  // Define the MmaCore components
+  // Mma core does not depend on stages, so pass in at least 3 here to mma
+  // multistage pieces are created
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          ElementA,
+                                                          LayoutA,
+                                                          ElementB,
+                                                          LayoutB,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          OperatorClass,
+                                                          std::max(kStages, 3),
+                                                          Operator,
+                                                          false,
+                                                          CacheOpA,
+                                                          CacheOpB>;
 
-    // Define the MmaCore components
-    // Mma core does not depend on stages, so pass in at least 3 here to mma multistage pieces are created
-    using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
-                                                                        WarpShape,
-                                                                        InstructionShape,
-                                                                        ElementA,
-                                                                        LayoutA,
-                                                                        ElementB,
-                                                                        LayoutB,
-                                                                        ElementAccumulator,
-                                                                        layout::RowMajor,
-                                                                        OperatorClass,
-                                                                        std::max(kStages, 3),
-                                                                        Operator,
-                                                                        false,
-                                                                        CacheOpA,
-                                                                        CacheOpB>;
+  // Define iterators over tiles from the A operand
+  using ThreadMapA = typename MmaCore::IteratorThreadMapA;
+  using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
+  using IteratorA =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
+          ElementA,
+          LayoutA,
+          1,
+          ThreadMapA,
+          AccessTypeA>;
 
-    // Define iterators over tiles from the A operand
-    using ThreadMapA  = typename MmaCore::IteratorThreadMapA;
-    using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
-    using IteratorA   = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
-        ElementA,
-        LayoutA,
-        1,
-        ThreadMapA,
-        AccessTypeA>;
+  // Define iterators over tiles from the B operand
+  using ThreadMapB = typename MmaCore::IteratorThreadMapB;
+  using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
+  using IteratorB =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>,
+          ElementB,
+          LayoutB,
+          0,
+          ThreadMapB,
+          AccessTypeB>;
+  static int const kAlignmentNF4LookUpTable = 128 / sizeof_bits<int32_t>::value;
+  using AccessTypeNF4LookUpTable =
+      cutlass::Array<int32_t, kAlignmentNF4LookUpTable>;
+  using IteratorNF4LookUpTableThreadMap =
+      transform::PitchLinearStripminedThreadMap<layout::PitchLinearShape<4, 1>,
+                                                16 / kAlignmentNF4LookUpTable,
+                                                kAlignmentNF4LookUpTable>;
+  using IteratorNF4LookUpTable =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<1, 16>,
+          int32_t,
+          cutlass::layout::RowMajor,
+          0,
+          IteratorNF4LookUpTableThreadMap,
+          AccessTypeNF4LookUpTable>;
 
-    // Define iterators over tiles from the B operand
-    using ThreadMapB  = typename MmaCore::IteratorThreadMapB;
-    using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
-    using IteratorB   = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>,
-        ElementB,
-        LayoutB,
-        0,
-        ThreadMapB,
-        AccessTypeB>;
-    static int const kAlignmentNF4LookUpTable = 128 / sizeof_bits<int32_t>::value;
-    using AccessTypeNF4LookUpTable = cutlass::Array<int32_t, kAlignmentNF4LookUpTable>;
-    using IteratorNF4LookUpTableThreadMap = transform::PitchLinearStripminedThreadMap<
-        layout::PitchLinearShape<4, 1>,
-        16 / kAlignmentNF4LookUpTable,
-        kAlignmentNF4LookUpTable>;
-    using IteratorNF4LookUpTable = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<1, 16>,
-        int32_t,
-        cutlass::layout::RowMajor,
-        0,
-        IteratorNF4LookUpTableThreadMap,
-        AccessTypeNF4LookUpTable>;
+  using SmemIteratorNF4LookUpTable = IteratorNF4LookUpTable;
 
-    using SmemIteratorNF4LookUpTable = IteratorNF4LookUpTable;
+  using Converter = FastInterleavedAndBiasedNumericArrayConverterNf4<
+      ElementA,
+      ElementB,
+      MmaCore::MmaPolicy::Operator::FragmentB::kElements>;
 
-    using Converter = FastInterleavedAndBiasedNumericArrayConverterNf4<ElementA,
-                                                                    ElementB,
-                                                                    MmaCore::MmaPolicy::Operator::FragmentB::kElements>;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::Int8Nf4InterleavedMmaMultistage<typename MmaCore::Shape,
-                                                                       IteratorA,
-                                                                       typename MmaCore::SmemIteratorA,
-                                                                       MmaCore::kCacheOpA,
-                                                                       IteratorB,
-                                                                       typename MmaCore::SmemIteratorB,
-                                                                       MmaCore::kCacheOpB,
-                                                                       IteratorNF4LookUpTable,
-                                                                       SmemIteratorNF4LookUpTable,
-                                                                       ElementAccumulator,
-                                                                       layout::RowMajor,
-                                                                       typename MmaCore::MmaPolicy,
-                                                                       kStages,
-                                                                       Converter,
-                                                                       SharedMemoryClear>;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma =
+      cutlass::gemm::threadblock::Int8Nf4InterleavedMmaMultistage<
+          typename MmaCore::Shape,
+          IteratorA,
+          typename MmaCore::SmemIteratorA,
+          MmaCore::kCacheOpA,
+          IteratorB,
+          typename MmaCore::SmemIteratorB,
+          MmaCore::kCacheOpB,
+          IteratorNF4LookUpTable,
+          SmemIteratorNF4LookUpTable,
+          ElementAccumulator,
+          layout::RowMajor,
+          typename MmaCore::MmaPolicy,
+          kStages,
+          Converter,
+          SharedMemoryClear>;
 };
 
-template<
+template <
     /// Type for element A
     typename ElementA,
     /// Layout type for A matrix operand
@@ -224,140 +238,156 @@ template<
     int RowsPerTile,
     ///
     int ColumnsInterleaved>
-struct Int8Nf4InterleavedMma<ElementA,
-             LayoutA,
-             kAlignmentA,
-             ElementB,
-             layout::ColumnMajorTileInterleave<RowsPerTile, ColumnsInterleaved>,
-             kAlignmentB,
-             ElementAccumulator,
-             layout::RowMajor,
-             OperatorClass,
-             ArchTag,
-             ThreadblockShape,
-             WarpShape,
-             InstructionShape,
-             kStages,
-             Operator,
-             SharedMemoryClear,
-             typename platform::enable_if<(ArchTag::kMinComputeCapability >= 80)>::type> {
+struct Int8Nf4InterleavedMma<
+    ElementA,
+    LayoutA,
+    kAlignmentA,
+    ElementB,
+    layout::ColumnMajorTileInterleave<RowsPerTile, ColumnsInterleaved>,
+    kAlignmentB,
+    ElementAccumulator,
+    layout::RowMajor,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    kStages,
+    Operator,
+    SharedMemoryClear,
+    typename platform::enable_if<(ArchTag::kMinComputeCapability >=
+                                  80)>::type> {
+  static_assert(platform::is_same<ElementA, int8_t>::value, "Element A int8_t");
 
-    static_assert(platform::is_same<ElementA, int8_t>::value ,
-                  "Element A int8_t");
+  static_assert(
+      platform::is_same<Operator,
+                        arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
+      "Mma multistage must dequantize after ldsm");
 
-    static_assert(platform::is_same<Operator, arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
-                  "Mma multistage must dequantize after ldsm");
+  // static_assert(platform::is_same<ElementB, int8_t>::value ||
+  // platform::is_same<ElementB, uint4b_t>::value,
+  //               "Element B must be uint8 or uint4");
+  static_assert(platform::is_same<ElementB, uint4b_t>::value,
+                "Element B must be uint4");
 
-    // static_assert(platform::is_same<ElementB, int8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-    //               "Element B must be uint8 or uint4");
-    static_assert(platform::is_same<ElementB, uint4b_t>::value,
-                  "Element B must be uint4");
+  static cutlass::arch::CacheOperation::Kind const CacheOpA =
+      ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits<ElementA>::value * kAlignmentA) == 128) ?
-                                                                    cutlass::arch::CacheOperation::Global :
-                                                                    cutlass::arch::CacheOperation::Always;
+  static cutlass::arch::CacheOperation::Kind const CacheOpB =
+      ((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
+          ? cutlass::arch::CacheOperation::Global
+          : cutlass::arch::CacheOperation::Always;
 
-    static cutlass::arch::CacheOperation::Kind const CacheOpB = ((sizeof_bits<ElementB>::value * kAlignmentB) == 128) ?
-                                                                    cutlass::arch::CacheOperation::Global :
-                                                                    cutlass::arch::CacheOperation::Always;
+  // Define the MmaCore components
+  // Mma core does not depend on stages, so pass in at least 3 here to mma
+  // multistage pieces are created
+  using MmaCore =
+      typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
+                                                          WarpShape,
+                                                          InstructionShape,
+                                                          ElementA,
+                                                          LayoutA,
+                                                          ElementB,
+                                                          layout::ColumnMajor,
+                                                          ElementAccumulator,
+                                                          layout::RowMajor,
+                                                          OperatorClass,
+                                                          std::max(kStages, 3),
+                                                          Operator,
+                                                          false,
+                                                          CacheOpA,
+                                                          CacheOpB>;
 
-    // Define the MmaCore components
-    // Mma core does not depend on stages, so pass in at least 3 here to mma multistage pieces are created
-    using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<ThreadblockShape,
-                                                                        WarpShape,
-                                                                        InstructionShape,
-                                                                        ElementA,
-                                                                        LayoutA,
-                                                                        ElementB,
-                                                                        layout::ColumnMajor,
-                                                                        ElementAccumulator,
-                                                                        layout::RowMajor,
-                                                                        OperatorClass,
-                                                                        std::max(kStages, 3),
-                                                                        Operator,
-                                                                        false,
-                                                                        CacheOpA,
-                                                                        CacheOpB>;
+  // Define iterators over tiles from the A operand
+  using ThreadMapA = typename MmaCore::IteratorThreadMapA;
+  using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
+  using IteratorA =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
+          ElementA,
+          LayoutA,
+          1,
+          ThreadMapA,
+          AccessTypeA>;
 
-    // Define iterators over tiles from the A operand
-    using ThreadMapA  = typename MmaCore::IteratorThreadMapA;
-    using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
-    using IteratorA   = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
-        ElementA,
-        LayoutA,
-        1,
-        ThreadMapA,
-        AccessTypeA>;
+ private:
+  static_assert(!(MmaCore::Shape::kN % ColumnsInterleaved), "");
+  static_assert(RowsPerTile == MmaCore::Shape::kK, "");
+  // static_assert(ColumnsInterleaved==4 || cutlass::platform::is_same<ElementB,
+  // int8_t>::value, "####");
+  using OriginalThreadMap = typename MmaCore::IteratorThreadMapB;
+  using OriginalWarpArrangement =
+      typename OriginalThreadMap::Detail::WarpThreadArrangement;
+  static_assert(!(OriginalWarpArrangement::kStrided % ColumnsInterleaved), "");
 
-private:
-    static_assert(!(MmaCore::Shape::kN % ColumnsInterleaved), "");
-    static_assert(RowsPerTile == MmaCore::Shape::kK, "");
-    // static_assert(ColumnsInterleaved==4 || cutlass::platform::is_same<ElementB, int8_t>::value, "####");
-    using OriginalThreadMap       = typename MmaCore::IteratorThreadMapB;
-    using OriginalWarpArrangement = typename OriginalThreadMap::Detail::WarpThreadArrangement;
-    static_assert(!(OriginalWarpArrangement::kStrided % ColumnsInterleaved), "");
+  using GmemIteratorShape =
+      MatrixShape<MmaCore::Shape::kK * ColumnsInterleaved,
+                  MmaCore::Shape::kN / ColumnsInterleaved>;
+  using GmemThreadMapB = transform::PitchLinearWarpRakedThreadMap<
+      layout::PitchLinearShape<GmemIteratorShape::kRow,
+                               GmemIteratorShape::kColumn>,
+      OriginalThreadMap::kThreads,
+      layout::PitchLinearShape<
+          OriginalWarpArrangement::kContiguous * ColumnsInterleaved,
+          OriginalWarpArrangement::kStrided / ColumnsInterleaved>,
+      MmaCore::kAccessSizeInBits / sizeof_bits<ElementB>::value>;
 
-    using GmemIteratorShape =
-        MatrixShape<MmaCore::Shape::kK * ColumnsInterleaved,
-        MmaCore::Shape::kN / ColumnsInterleaved>;
-    using GmemThreadMapB = transform::PitchLinearWarpRakedThreadMap<
-        layout::PitchLinearShape<GmemIteratorShape::kRow, GmemIteratorShape::kColumn>,
-        OriginalThreadMap::kThreads,
-        layout::PitchLinearShape<OriginalWarpArrangement::kContiguous * ColumnsInterleaved,
-                                 OriginalWarpArrangement::kStrided / ColumnsInterleaved>,
-        MmaCore::kAccessSizeInBits / sizeof_bits<ElementB>::value>;
+ public:
+  // Define iterators over tiles from the B operand
+  using ThreadMapB = typename MmaCore::IteratorThreadMapB;
+  using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
+  using IteratorB =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          GmemIteratorShape,
+          ElementB,
+          layout::ColumnMajor,
+          0,
+          GmemThreadMapB,
+          AccessTypeB>;
 
-public:
-    // Define iterators over tiles from the B operand
-    using ThreadMapB  = typename MmaCore::IteratorThreadMapB;
-    using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
-    using IteratorB   = cutlass::transform::threadblock::
-        PredicatedTileAccessIterator<GmemIteratorShape,
-        ElementB,
-        layout::ColumnMajor,
-        0,
-        GmemThreadMapB,
-        AccessTypeB>;
+  static int const kAlignmentNF4LookUpTable = 128 / sizeof_bits<int32_t>::value;
+  using AccessTypeNF4LookUpTable =
+      cutlass::Array<int32_t, kAlignmentNF4LookUpTable>;
+  using IteratorNF4LookUpTableThreadMap = transform::
+      PitchLinearStripminedThreadMap<layout::PitchLinearShape<4, 1>, 1, 4>;
+  using IteratorNF4LookUpTable =
+      cutlass::transform::threadblock::PredicatedTileAccessIterator<
+          cutlass::MatrixShape<1, 16>,
+          int32_t,
+          cutlass::layout::RowMajor,
+          0,
+          IteratorNF4LookUpTableThreadMap,
+          AccessTypeNF4LookUpTable>;
 
-    static int const kAlignmentNF4LookUpTable = 128 / sizeof_bits<int32_t>::value;
-    using AccessTypeNF4LookUpTable = cutlass::Array<int32_t, kAlignmentNF4LookUpTable>;
-    using IteratorNF4LookUpTableThreadMap = transform::PitchLinearStripminedThreadMap<
-        layout::PitchLinearShape<4, 1>,
-        1,
-        4>;
-    using IteratorNF4LookUpTable = cutlass::transform::threadblock::PredicatedTileAccessIterator<
-        cutlass::MatrixShape<1, 16>,
-        int32_t,
-        cutlass::layout::RowMajor,
-        0,
-        IteratorNF4LookUpTableThreadMap,
-        AccessTypeNF4LookUpTable>;
+  using SmemIteratorNF4LookUpTable = IteratorNF4LookUpTable;
 
+  // static_assert(MmaCore::MmaPolicy::Operator::FragmentB::kElements==64,"MmaCore::MmaPolicy::Operator::FragmentB::kElements
+  // == 32");
+  using Converter = FastInterleavedAndBiasedNumericArrayConverterNf4<
+      ElementA,
+      ElementB,
+      MmaCore::MmaPolicy::Operator::FragmentB::kElements>;
 
-    using SmemIteratorNF4LookUpTable = IteratorNF4LookUpTable;
-
-    // static_assert(MmaCore::MmaPolicy::Operator::FragmentB::kElements==64,"MmaCore::MmaPolicy::Operator::FragmentB::kElements == 32");
-    using Converter = FastInterleavedAndBiasedNumericArrayConverterNf4<ElementA,
-                                                                    ElementB,
-                                                                    MmaCore::MmaPolicy::Operator::FragmentB::kElements>;
-
-    // Define the threadblock-scoped pipelined matrix multiply
-    using ThreadblockMma = cutlass::gemm::threadblock::Int8Nf4InterleavedMmaMultistage<typename MmaCore::Shape,
-                                                                       IteratorA,
-                                                                       typename MmaCore::SmemIteratorA,
-                                                                       MmaCore::kCacheOpA,
-                                                                       IteratorB,
-                                                                       typename MmaCore::SmemIteratorB,
-                                                                       MmaCore::kCacheOpB,
-                                                                       IteratorNF4LookUpTable,
-                                                                       SmemIteratorNF4LookUpTable,
-                                                                       ElementAccumulator,
-                                                                       layout::RowMajor,
-                                                                       typename MmaCore::MmaPolicy,
-                                                                       kStages,
-                                                                       Converter,
-                                                                       SharedMemoryClear>;
+  // Define the threadblock-scoped pipelined matrix multiply
+  using ThreadblockMma =
+      cutlass::gemm::threadblock::Int8Nf4InterleavedMmaMultistage<
+          typename MmaCore::Shape,
+          IteratorA,
+          typename MmaCore::SmemIteratorA,
+          MmaCore::kCacheOpA,
+          IteratorB,
+          typename MmaCore::SmemIteratorB,
+          MmaCore::kCacheOpB,
+          IteratorNF4LookUpTable,
+          SmemIteratorNF4LookUpTable,
+          ElementAccumulator,
+          layout::RowMajor,
+          typename MmaCore::MmaPolicy,
+          kStages,
+          Converter,
+          SharedMemoryClear>;
 };
 
 }  // namespace threadblock
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_base.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_base.h
index ac47190447..585d88ec85 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_base.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_base.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
@@ -63,46 +64,48 @@ namespace gemm {
 namespace threadblock {
 
 // ////////////////////////////////////////////////////////////////////////////////
-// // SFINAE trick so I can keep the same loop code for Volta and dispatch to the
-// // correct warp level mma. On volta, all data is stored to shared memory as FP16.
-// template<typename WarpMma, int kExpansionFactor = 1>
-// CUTLASS_DEVICE void run_warp_mma(WarpMma&                           warp_mma,
+// // SFINAE trick so I can keep the same loop code for Volta and dispatch to
+// the
+// // correct warp level mma. On volta, all data is stored to shared memory as
+// FP16. template<typename WarpMma, int kExpansionFactor = 1> CUTLASS_DEVICE
+// void run_warp_mma(WarpMma&                           warp_mma,
 //                                  typename WarpMma::FragmentC&       D,
 //                                  typename WarpMma::FragmentA const& A,
 //                                  typename WarpMma::FragmentB const& B,
 //                                  typename WarpMma::FragmentC const& C,
-//                                  const int                          warp_tileB_k_offset)
+//                                  const int warp_tileB_k_offset)
 // {
 //     warp_mma(D, A, B, C);
 // }
 
 // template<typename WarpMma, int kExpansionFactor = WarpMma::kExpansionFactor>
-// CUTLASS_DEVICE void run_warp_mma(WarpMma&                                      warp_mma,
-//                                  typename WarpMma::FragmentC&                  D,
-//                                  typename WarpMma::TransformedFragmentA const& A,
-//                                  typename WarpMma::TransformedFragmentB const& B,
-//                                  typename WarpMma::FragmentC const&            C,
-//                                  const int                                     warp_tileB_k_offset)
+// CUTLASS_DEVICE void run_warp_mma(WarpMma& warp_mma,
+//                                  typename WarpMma::FragmentC& D, typename
+//                                  WarpMma::TransformedFragmentA const& A,
+//                                  typename WarpMma::TransformedFragmentB
+//                                  const& B, typename WarpMma::FragmentC const&
+//                                  C, const int warp_tileB_k_offset)
 // {
 //     warp_mma(D, A, B, C, warp_tileB_k_offset);
 // }
 
-// TODO(zhengzekang): Since we annotate the first implement, we currently hack to `run_ampere_warp_mma` used in A100.
-template<typename WarpMma, int kExpansionFactor = WarpMma::kExpansionFactor>
-CUTLASS_DEVICE void run_ampere_warp_mma(WarpMma&                                      warp_mma,
-                                 typename WarpMma::FragmentC&                  D,
-                                 typename WarpMma::TransformedFragmentA const& A,
-                                 typename WarpMma::TransformedFragmentB const& B,
-                                 typename WarpMma::FragmentC const&            C,
-                                 const int                                     warp_tileB_k_offset)
-{
-    warp_mma(D, A, B, C, warp_tileB_k_offset);
+// TODO(zhengzekang): Since we annotate the first implement, we currently hack
+// to `run_ampere_warp_mma` used in A100.
+template <typename WarpMma, int kExpansionFactor = WarpMma::kExpansionFactor>
+CUTLASS_DEVICE void run_ampere_warp_mma(
+    WarpMma& warp_mma,
+    typename WarpMma::FragmentC& D,
+    typename WarpMma::TransformedFragmentA const& A,
+    typename WarpMma::TransformedFragmentB const& B,
+    typename WarpMma::FragmentC const& C,
+    const int warp_tileB_k_offset) {
+  warp_mma(D, A, B, C, warp_tileB_k_offset);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Structure to compute the matrix product targeting CUDA cores and SIMT math
 /// instructions.
-template<
+template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Policy describing tuning details (concept: MmaPolicy)
@@ -112,135 +115,139 @@ template<
     /// Used for partial specialization
     typename Enable = bool>
 class Int8InterleavedMmaBase {
-public:
-    ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using Shape = Shape_;
+ public:
+  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using Shape = Shape_;
 
-    ///< Policy describing tuning details
-    using Policy = Policy_;
+  ///< Policy describing tuning details
+  using Policy = Policy_;
 
+  //
+  // Dependent types
+  //
+
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
+
+  /// Shape describing the overall GEMM computed from shared memory
+  /// by each warp.
+  using WarpGemm = typename Policy::Operator::Shape;
+
+  /// Shape describing the number of warps filling the CTA
+  using WarpCount = GemmShape<Shape::kM / WarpGemm::kM,
+                              Shape::kN / WarpGemm::kN,
+                              Shape::kK / WarpGemm::kK>;
+
+  /// Number of warp-level GEMM operations
+  static int const kWarpGemmIterations =
+      (WarpGemm::kK / Operator::Policy::MmaShape::kK);
+  static_assert(Operator::IteratorB::InstructionShape::kRow >=
+                    Operator::InstructionShape::kK,
+                "");
+  static constexpr int kNumKIterationsPerWarpBLoad =
+      Operator::IteratorB::InstructionShape::kRow /
+      Operator::InstructionShape::kK;
+  static_assert(!(kWarpGemmIterations % kNumKIterationsPerWarpBLoad), "");
+  static constexpr int kWarpGemmIterationsForB =
+      kWarpGemmIterations / kNumKIterationsPerWarpBLoad;
+
+  /// Number of stages
+  static int const kStages = Stages;
+
+  /// Tensor reference to the A operand
+  using TensorRefA =
+      TensorRef<typename Operator::ElementA, typename Operator::LayoutA>;
+
+  /// Tensor reference to the B operand
+  using TensorRefB =
+      TensorRef<typename Operator::ElementB, typename Operator::LayoutB>;
+
+  //
+  // Nested structs
+  //
+
+  /// Shared storage object needed by threadblock-scoped GEMM
+  class SharedStorage {
+   public:
     //
-    // Dependent types
+    // Type definitions
     //
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+    /// Shape of the A matrix operand in shared memory
+    using ShapeA =
+        MatrixShape<Shape::kM + Policy::SmemPaddingA::kRow,
+                    Shape::kK * kStages + Policy::SmemPaddingA::kColumn>;
 
-    /// Shape describing the overall GEMM computed from shared memory
-    /// by each warp.
-    using WarpGemm = typename Policy::Operator::Shape;
+    /// Shape of the B matrix operand in shared memory
+    using ShapeB = MatrixShape<Shape::kK * kStages + Policy::SmemPaddingB::kRow,
+                               Shape::kN + Policy::SmemPaddingB::kColumn>;
 
-    /// Shape describing the number of warps filling the CTA
-    using WarpCount = GemmShape<Shape::kM / WarpGemm::kM, Shape::kN / WarpGemm::kN, Shape::kK / WarpGemm::kK>;
-
-    /// Number of warp-level GEMM operations
-    static int const kWarpGemmIterations = (WarpGemm::kK / Operator::Policy::MmaShape::kK);
-    static_assert(Operator::IteratorB::InstructionShape::kRow>=Operator::InstructionShape::kK,"");
-    static constexpr int kNumKIterationsPerWarpBLoad =
-        Operator::IteratorB::InstructionShape::kRow / Operator::InstructionShape::kK;
-    static_assert(!(kWarpGemmIterations % kNumKIterationsPerWarpBLoad), "");
-    static constexpr int kWarpGemmIterationsForB = kWarpGemmIterations / kNumKIterationsPerWarpBLoad;
-
-    /// Number of stages
-    static int const kStages = Stages;
-
-    /// Tensor reference to the A operand
-    using TensorRefA = TensorRef<typename Operator::ElementA, typename Operator::LayoutA>;
-
-    /// Tensor reference to the B operand
-    using TensorRefB = TensorRef<typename Operator::ElementB, typename Operator::LayoutB>;
-
-    //
-    // Nested structs
-    //
-
-    /// Shared storage object needed by threadblock-scoped GEMM
-    class SharedStorage {
-    public:
-        //
-        // Type definitions
-        //
-
-        /// Shape of the A matrix operand in shared memory
-        using ShapeA =
-            MatrixShape<Shape::kM + Policy::SmemPaddingA::kRow, Shape::kK * kStages + Policy::SmemPaddingA::kColumn>;
-
-        /// Shape of the B matrix operand in shared memory
-        using ShapeB =
-            MatrixShape<Shape::kK * kStages + Policy::SmemPaddingB::kRow, Shape::kN + Policy::SmemPaddingB::kColumn>;
-
-    public:
-        //
-        // Data members
-        //
-
-        /// Buffer for A operand
-        AlignedBuffer<typename Operator::ElementA, ShapeA::kCount> operand_A;
-
-        /// Buffer for B operand
-        AlignedBuffer<typename Operator::ElementB, ShapeB::kCount> operand_B;
-
-    public:
-        //
-        // Methods
-        //
-
-        /// Returns a layout object for the A matrix
-        CUTLASS_DEVICE
-        static typename Operator::LayoutA LayoutA()
-        {
-            return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn});
-        }
-
-        /// Returns a layout object for the B matrix
-        CUTLASS_HOST_DEVICE
-        static typename Operator::LayoutB LayoutB()
-        {
-            return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn});
-        }
-
-        /// Returns a TensorRef to the A operand
-        CUTLASS_HOST_DEVICE
-        TensorRefA operand_A_ref()
-        {
-            return TensorRefA{operand_A.data(), LayoutA()};
-        }
-
-        /// Returns a TensorRef to the B operand
-        CUTLASS_HOST_DEVICE
-        TensorRefB operand_B_ref()
-        {
-            return TensorRefB{operand_B.data(), LayoutB()};
-        }
-    };
-
-protected:
+   public:
     //
     // Data members
     //
 
-    /// Iterator to load a warp-scoped tile of A operand from shared memory
-    typename Operator::IteratorA warp_tile_iterator_A_;
+    /// Buffer for A operand
+    AlignedBuffer<typename Operator::ElementA, ShapeA::kCount> operand_A;
 
-    /// Iterator to load a warp-scoped tile of B operand from shared memory
-    typename Operator::IteratorB warp_tile_iterator_B_;
+    /// Buffer for B operand
+    AlignedBuffer<typename Operator::ElementB, ShapeB::kCount> operand_B;
 
-public:
-    /// Construct from tensor references
+   public:
+    //
+    // Methods
+    //
+
+    /// Returns a layout object for the A matrix
     CUTLASS_DEVICE
-    Int8InterleavedMmaBase(
-        ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        SharedStorage& shared_storage,
-        ///< ID within the threadblock
-        int thread_idx,
-        ///< ID of warp
-        int warp_idx,
-        ///< ID of each thread within a warp
-        int lane_idx):
-        warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx),
-        warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx)
-    {
+    static typename Operator::LayoutA LayoutA() {
+      return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn});
     }
+
+    /// Returns a layout object for the B matrix
+    CUTLASS_HOST_DEVICE
+    static typename Operator::LayoutB LayoutB() {
+      return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn});
+    }
+
+    /// Returns a TensorRef to the A operand
+    CUTLASS_HOST_DEVICE
+    TensorRefA operand_A_ref() {
+      return TensorRefA{operand_A.data(), LayoutA()};
+    }
+
+    /// Returns a TensorRef to the B operand
+    CUTLASS_HOST_DEVICE
+    TensorRefB operand_B_ref() {
+      return TensorRefB{operand_B.data(), LayoutB()};
+    }
+  };
+
+ protected:
+  //
+  // Data members
+  //
+
+  /// Iterator to load a warp-scoped tile of A operand from shared memory
+  typename Operator::IteratorA warp_tile_iterator_A_;
+
+  /// Iterator to load a warp-scoped tile of B operand from shared memory
+  typename Operator::IteratorB warp_tile_iterator_B_;
+
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  Int8InterleavedMmaBase(
+      ///< Shared storage needed for internal use by threadblock-scoped GEMM
+      SharedStorage& shared_storage,
+      ///< ID within the threadblock
+      int thread_idx,
+      ///< ID of warp
+      int warp_idx,
+      ///< ID of each thread within a warp
+      int lane_idx)
+      : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx),
+        warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {}
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_multistage.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_multistage.h
index 3021605f27..57e19ad92b 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_multistage.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_multistage.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,18 +18,18 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
-
 /* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -72,7 +72,7 @@ namespace threadblock {
 
 /// Structure to compute the matrix product targeting CUDA cores and SIMT math
 /// instructions.
-template<
+template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Iterates over tiles of A operand in global memory
@@ -107,489 +107,532 @@ template<
     SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone,
     /// Used for partial specialization
     typename Enable = bool>
-class Int8InterleavedMmaMultistage: public Int8InterleavedMmaBase<Shape_, Policy_, Stages> {
-public:
-    ///< Base class
-    using Base = Int8InterleavedMmaBase<Shape_, Policy_, Stages>;
-    ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using Shape = Shape_;
-    ///< Iterates over tiles of A operand in global memory
-    using IteratorA = IteratorA_;
-    ///< Iterates over tiles of B operand in global memory
-    using IteratorB = IteratorB_;
-    ///< Data type of accumulator matrix
-    using ElementC = ElementC_;
-    ///< Layout of accumulator matrix
-    using LayoutC = LayoutC_;
-    ///< Policy describing tuning details
-    using Policy = Policy_;
+class Int8InterleavedMmaMultistage
+    : public Int8InterleavedMmaBase<Shape_, Policy_, Stages> {
+ public:
+  ///< Base class
+  using Base = Int8InterleavedMmaBase<Shape_, Policy_, Stages>;
+  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using Shape = Shape_;
+  ///< Iterates over tiles of A operand in global memory
+  using IteratorA = IteratorA_;
+  ///< Iterates over tiles of B operand in global memory
+  using IteratorB = IteratorB_;
+  ///< Data type of accumulator matrix
+  using ElementC = ElementC_;
+  ///< Layout of accumulator matrix
+  using LayoutC = LayoutC_;
+  ///< Policy describing tuning details
+  using Policy = Policy_;
 
+  using SmemIteratorA = SmemIteratorA_;
+  using SmemIteratorB = SmemIteratorB_;
 
-    using SmemIteratorA     = SmemIteratorA_;
-    using SmemIteratorB     = SmemIteratorB_;
+  static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
+  static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
 
-    static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
-    static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
+  using TransformBAfterLDS = TransformBAfterLDS_;
 
-    using TransformBAfterLDS = TransformBAfterLDS_;
+  //
+  // Dependent types
+  //
 
-    //
-    // Dependent types
-    //
+  /// Fragment of accumulator tile
+  using FragmentC = typename Policy::Operator::FragmentC;
 
-    /// Fragment of accumulator tile
-    using FragmentC = typename Policy::Operator::FragmentC;
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+  /// Minimum architecture is Sm80 to support cp.async
+  using ArchTag = arch::Sm80;
 
-    /// Minimum architecture is Sm80 to support cp.async
-    using ArchTag = arch::Sm80;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = Operator::kTransformA;
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = Operator::kTransformA;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = Operator::kTransformB;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = Operator::kTransformB;
+  /// Internal structure exposed for introspection.
+  struct Detail {
+    static_assert(Base::kWarpGemmIterations > 1,
+                  "The pipelined structure requires at least two warp-level "
+                  "GEMM operations.");
+    // static_assert(Base::kWarpGemmIterations==4,"Base::kWarpGemmIterations!=4");
+    /// Number of cp.async instructions to load one stage of operand A
+    static int const AsyncCopyIterationsPerStageA =
+        IteratorA::ThreadMap::Iterations::kCount;
 
-    /// Internal structure exposed for introspection.
-    struct Detail {
+    /// Number of cp.async instructions to load one stage of operand B
+    static int const AsyncCopyIterationsPerStageB =
+        IteratorB::ThreadMap::Iterations::kCount;
 
-        static_assert(Base::kWarpGemmIterations > 1,
-                      "The pipelined structure requires at least two warp-level "
-                      "GEMM operations.");
-        // static_assert(Base::kWarpGemmIterations==4,"Base::kWarpGemmIterations!=4");
-        /// Number of cp.async instructions to load one stage of operand A
-        static int const AsyncCopyIterationsPerStageA = IteratorA::ThreadMap::Iterations::kCount;
+    /// Number of stages
+    static int const kStages = Stages;
 
-        /// Number of cp.async instructions to load one stage of operand B
-        static int const AsyncCopyIterationsPerStageB = IteratorB::ThreadMap::Iterations::kCount;
+    /// Number of cp.async instructions to load on group of operand A
+    static int const kAccessesPerGroupA =
+        (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
 
-        /// Number of stages
-        static int const kStages = Stages;
+    /// Number of cp.async instructions to load on group of operand B
+    static int const kAccessesPerGroupB =
+        (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
+  };
 
-        /// Number of cp.async instructions to load on group of operand A
-        static int const kAccessesPerGroupA =
-            (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
+ private:
+  using WarpFragmentA = typename Operator::FragmentA;
+  using WarpFragmentB = typename Operator::FragmentB;
+  using ElementB = typename IteratorB::Element;
+  using LayoutDetailsForB = kernel::LayoutDetailsB<ElementB, ArchTag>;
 
-        /// Number of cp.async instructions to load on group of operand B
-        static int const kAccessesPerGroupB =
-            (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
-    };
+  static constexpr bool RequiresTileInterleave =
+      layout::IsColumnMajorTileInterleave<
+          typename LayoutDetailsForB::Layout>::value;
+  static_assert(!RequiresTileInterleave ||
+                    (RequiresTileInterleave &&
+                     (Shape::kK == LayoutDetailsForB::ThreadblockK)),
+                "Layout K must match threadblockK");
 
-private:
-    using WarpFragmentA = typename Operator::FragmentA;
-    using WarpFragmentB = typename Operator::FragmentB;
-    using ElementB          = typename IteratorB::Element;
-    using LayoutDetailsForB = kernel::LayoutDetailsB<ElementB, ArchTag>;
+ private:
+  //
+  // Data members
+  //
 
-    static constexpr bool RequiresTileInterleave =
-        layout::IsColumnMajorTileInterleave<typename LayoutDetailsForB::Layout>::value;
-    static_assert(!RequiresTileInterleave || (RequiresTileInterleave && (Shape::kK == LayoutDetailsForB::ThreadblockK)),
-                  "Layout K must match threadblockK");
+  /// Iterator to write threadblock-scoped tile of A operand to shared memory
+  SmemIteratorA smem_iterator_A_;
 
-private:
-    //
-    // Data members
-    //
+  /// Iterator to write threadblock-scoped tile of B operand to shared memory
+  SmemIteratorB smem_iterator_B_;
 
-    /// Iterator to write threadblock-scoped tile of A operand to shared memory
-    SmemIteratorA smem_iterator_A_;
-
-    /// Iterator to write threadblock-scoped tile of B operand to shared memory
-    SmemIteratorB smem_iterator_B_;
-
-public:
-    /// Construct from tensor references
-    CUTLASS_DEVICE
-    Int8InterleavedMmaMultistage(
-        ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        typename Base::SharedStorage& shared_storage,
-        ///< ID within the threadblock
-        int thread_idx,
-        ///< ID of warp
-        int warp_idx,
-        ///< ID of each thread within a warp
-        int lane_idx):
-        Base(shared_storage, thread_idx, warp_idx, lane_idx),
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  Int8InterleavedMmaMultistage(
+      ///< Shared storage needed for internal use by threadblock-scoped GEMM
+      typename Base::SharedStorage& shared_storage,
+      ///< ID within the threadblock
+      int thread_idx,
+      ///< ID of warp
+      int warp_idx,
+      ///< ID of each thread within a warp
+      int lane_idx)
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
         smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
-        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx)
-    {
-        // Compute warp location within threadblock tile by mapping the warp_id to
-        // three coordinates:
-        //   _m: the warp's position within the threadblock along the M dimension
-        //   _n: the warp's position within the threadblock along the N dimension
-        //   _k: the warp's position within the threadblock along the K dimension
+        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) {
+    // Compute warp location within threadblock tile by mapping the warp_id to
+    // three coordinates:
+    //   _m: the warp's position within the threadblock along the M dimension
+    //   _n: the warp's position within the threadblock along the N dimension
+    //   _k: the warp's position within the threadblock along the K dimension
 
-        int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
-        int warp_idx_k  = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
 
-        int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
-        int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
+    int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
+    int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
 
-        // Add per-warp offsets in units of warp-level tiles
-        this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
-        this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+    // Add per-warp offsets in units of warp-level tiles
+    this->warp_tile_iterator_A_.add_tile_offset(
+        {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
+    this->warp_tile_iterator_B_.add_tile_offset(
+        {Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  }
+
+  CUTLASS_DEVICE
+  void copy_tiles_and_advance(IteratorA& iterator_A,
+                              IteratorB& iterator_B,
+                              int group_start_A = 0,
+                              int group_start_B = 0) {
+    iterator_A.set_iteration_index(group_start_A *
+                                   IteratorA::kAccessesPerVector);
+    this->smem_iterator_A_.set_iteration_index(group_start_A);
+
+    // Async Copy for operand A
+    CUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) {
+      if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                this->smem_iterator_A_.get());
+
+        int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
+                              IteratorA::ThreadMap::kElementsPerAccess /
+                              IteratorA::kAccessesPerVector / 8;
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
+          auto gmem_ptr = iterator_A.get();
+
+          if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
+                dst_ptr + v, gmem_ptr, iterator_A.valid());
+          } else {
+            cutlass::arch::cp_async<kSrcBytes, kCacheOpA>(
+                dst_ptr + v, gmem_ptr, iterator_A.valid());
+          }
+
+          ++iterator_A;
+        }
+
+        ++this->smem_iterator_A_;
+      }
     }
 
-    CUTLASS_DEVICE
-    void
-    copy_tiles_and_advance(IteratorA& iterator_A, IteratorB& iterator_B, int group_start_A = 0, int group_start_B = 0)
-    {
-        iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector);
-        this->smem_iterator_A_.set_iteration_index(group_start_A);
+    iterator_B.set_iteration_index(group_start_B *
+                                   IteratorB::kAccessesPerVector);
+    this->smem_iterator_B_.set_iteration_index(group_start_B);
 
-        // Async Copy for operand A
+    // Async Copy for operand B
+    CUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) {
+      if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                this->smem_iterator_B_.get());
+
+        int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value *
+                              IteratorB::ThreadMap::kElementsPerAccess /
+                              IteratorB::kAccessesPerVector / 8;
+        // if((threadIdx.x==0)  && threadIdx.y==0  && threadIdx.z==0 &&
+        //     blockIdx.x==0  && blockIdx.y==0  && blockIdx.z==0){
+        //         printf("### kSrcBytes:%d,
+        //         IteratorB::kAccessesPerVector:%d\n", kSrcBytes,
+        //         IteratorB::kAccessesPerVector);
+        //     }
+        // // static_assert(kSrcBytes==32, "kSrcBytes==32");
+        // if( threadIdx.y==0  && threadIdx.z==0 &&
+        //     blockIdx.x==0  && blockIdx.y==0  && blockIdx.z==0){
+        //     printf("gmem_ptr cp source of thread %d: ",threadIdx.x);
+        //     for(int i=0;i<16;++i){
+        //         printf("%d, ",iterator_B.get()[i]);
+        //     }
+        //     printf("\n");
+        // }
         CUTLASS_PRAGMA_UNROLL
-        for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) {
-            if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) {
-                typename IteratorA::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorA::AccessType*>(this->smem_iterator_A_.get());
+        for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
+          auto gmem_ptr = iterator_B.get();
 
-                int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value
-                                      * IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector / 8;
+          if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
+                dst_ptr + v, gmem_ptr, iterator_B.valid());
+          } else {
+            cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(
+                dst_ptr + v, gmem_ptr, iterator_B.valid());
+          }
+          ++iterator_B;
+        }
+        ++this->smem_iterator_B_;
+      }
+    }
+  }
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
-                    auto gmem_ptr = iterator_A.get();
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  CUTLASS_DEVICE
+  void operator()(
+      ///< problem size of GEMM
+      int gemm_k_iterations,
+      ///< destination accumulator tile
+      FragmentC& accum,
+      ///< iterator over A operand in global memory
+      IteratorA iterator_A,
+      ///< iterator over B operand in global memory
+      IteratorB iterator_B,
+      ///< initial value of accumulator
+      FragmentC const& src_accum) {
+    //
+    // Prologue
+    //
 
-                    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
-                        cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(dst_ptr + v, gmem_ptr, iterator_A.valid());
-                    }
-                    else {
-                        cutlass::arch::cp_async<kSrcBytes, kCacheOpA>(dst_ptr + v, gmem_ptr, iterator_A.valid());
-                    }
+    TransformBAfterLDS lds_converter;
 
-                    ++iterator_A;
-                }
+    // NOTE - switch to ldg.sts
+    // Issue this first, so cp.async.commit_group will commit this load as well.
+    // Note: we do not commit here and this load will commit in the same group
+    // as
+    //       the first load of A.
 
-                ++this->smem_iterator_A_;
-            }
+    // Issue several complete stages
+    CUTLASS_PRAGMA_UNROLL
+    for (int stage = 0; stage < Base::kStages - 1;
+         ++stage, --gemm_k_iterations) {
+      iterator_A.clear_mask(gemm_k_iterations == 0);
+      iterator_B.clear_mask(gemm_k_iterations == 0);
+
+      iterator_A.set_iteration_index(0);
+      this->smem_iterator_A_.set_iteration_index(0);
+
+      // Async Copy for operand A
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                this->smem_iterator_A_.get());
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
+          int const kSrcBytes =
+              sizeof_bits<typename IteratorA::Element>::value *
+              IteratorA::ThreadMap::kElementsPerAccess /
+              IteratorA::kAccessesPerVector / 8;
+
+          int src_bytes = (iterator_A.valid() ? kSrcBytes : 0);
+
+          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
+              dst_ptr + v, iterator_A.get(), iterator_A.valid());
+
+          ++iterator_A;
         }
 
-        iterator_B.set_iteration_index(group_start_B * IteratorB::kAccessesPerVector);
-        this->smem_iterator_B_.set_iteration_index(group_start_B);
+        ++this->smem_iterator_A_;
+      }
+
+      iterator_B.set_iteration_index(0);
+      this->smem_iterator_B_.set_iteration_index(0);
+
+      // Async Copy for operand B
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                this->smem_iterator_B_.get());
 
-        // Async Copy for operand B
         CUTLASS_PRAGMA_UNROLL
-        for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) {
-            if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) {
-                typename IteratorB::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorB::AccessType*>(this->smem_iterator_B_.get());
+        for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
+          int const kSrcBytes =
+              sizeof_bits<typename IteratorB::Element>::value *
+              IteratorB::ThreadMap::kElementsPerAccess /
+              IteratorB::kAccessesPerVector / 8;
 
-                int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value
-                                      * IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector / 8;
-                // if((threadIdx.x==0)  && threadIdx.y==0  && threadIdx.z==0 &&
-                //     blockIdx.x==0  && blockIdx.y==0  && blockIdx.z==0){
-                //         printf("### kSrcBytes:%d, IteratorB::kAccessesPerVector:%d\n", kSrcBytes, IteratorB::kAccessesPerVector);
-                //     }
-                // // static_assert(kSrcBytes==32, "kSrcBytes==32");
-                // if( threadIdx.y==0  && threadIdx.z==0 &&
-                //     blockIdx.x==0  && blockIdx.y==0  && blockIdx.z==0){
-                //     printf("gmem_ptr cp source of thread %d: ",threadIdx.x);
-                //     for(int i=0;i<16;++i){
-                //         printf("%d, ",iterator_B.get()[i]);
-                //     }
-                //     printf("\n");
-                // }
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
-                    auto gmem_ptr = iterator_B.get();
+          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
+              dst_ptr + v, iterator_B.get(), iterator_B.valid());
 
-                    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
-                        cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(dst_ptr + v, gmem_ptr, iterator_B.valid());
-                    }
-                    else {
-                        cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(dst_ptr + v, gmem_ptr, iterator_B.valid());
-                    }
-                    ++iterator_B;
-                }
-                ++this->smem_iterator_B_;
-            }
+          ++iterator_B;
         }
+
+        ++this->smem_iterator_B_;
+      }
+
+      // Move to the next stage
+      iterator_A.add_tile_offset({0, 1});
+      iterator_B.add_tile_offset({1, 0});
+
+      this->smem_iterator_A_.add_tile_offset({0, 1});
+      this->smem_iterator_B_.add_tile_offset({1, 0});
+
+      // Defines the boundary of a stage of cp.async.
+      cutlass::arch::cp_async_fence();
     }
 
-    /// Perform a threadblock-scoped matrix multiply-accumulate
-    CUTLASS_DEVICE
-    void operator()(
-        ///< problem size of GEMM
-        int gemm_k_iterations,
-        ///< destination accumulator tile
-        FragmentC& accum,
-        ///< iterator over A operand in global memory
-        IteratorA iterator_A,
-        ///< iterator over B operand in global memory
-        IteratorB iterator_B,
-        ///< initial value of accumulator
-        FragmentC const& src_accum)
-    {
+    // Perform accumulation in the 'd' output operand
+    accum = src_accum;
 
-        //
-        // Prologue
-        //
+    //
+    // Clear the remaining tiles of SMEM. This is a functional requirement for
+    // some kernels so that all accumulator elements outside the GEMM footprint
+    // are zero.
+    //
 
-        TransformBAfterLDS lds_converter;
+    if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) {
+      /// Iterator to write threadblock-scoped tile of A operand to shared
+      /// memory
+      SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
 
-        // NOTE - switch to ldg.sts
-        // Issue this first, so cp.async.commit_group will commit this load as well.
-        // Note: we do not commit here and this load will commit in the same group as
-        //       the first load of A.
+      typename IteratorA::AccessType zero_A;
+      zero_A.clear();
 
-        // Issue several complete stages
-        CUTLASS_PRAGMA_UNROLL
-        for (int stage = 0; stage < Base::kStages - 1; ++stage, --gemm_k_iterations) {
+      last_smem_iterator_A.set_iteration_index(0);
 
-            iterator_A.clear_mask(gemm_k_iterations == 0);
-            iterator_B.clear_mask(gemm_k_iterations == 0);
+      // Async Copy for operand A
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                last_smem_iterator_A.get());
 
-            iterator_A.set_iteration_index(0);
-            this->smem_iterator_A_.set_iteration_index(0);
+        *dst_ptr = zero_A;
 
-            // Async Copy for operand A
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
-                typename IteratorA::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorA::AccessType*>(this->smem_iterator_A_.get());
+        ++last_smem_iterator_A;
+      }
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
-                    int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value
-                                          * IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector
-                                          / 8;
+      /// Iterator to write threadblock-scoped tile of B operand to shared
+      /// memory
+      SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_);
+      typename IteratorB::AccessType zero_B;
 
-                    int src_bytes = (iterator_A.valid() ? kSrcBytes : 0);
+      zero_B.clear();
+      last_smem_iterator_B.set_iteration_index(0);
 
-                    cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
-                        dst_ptr + v, iterator_A.get(), iterator_A.valid());
+      // Async Copy for operand B
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                last_smem_iterator_B.get());
 
-                    ++iterator_A;
-                }
+        *dst_ptr = zero_B;
 
-                ++this->smem_iterator_A_;
-            }
+        ++last_smem_iterator_B;
+      }
+    }
 
-            iterator_B.set_iteration_index(0);
-            this->smem_iterator_B_.set_iteration_index(0);
+    // Waits until kStages-2 stages have committed.
+    cutlass::arch::cp_async_wait<Base::kStages - 2>();
+    __syncthreads();
 
-            // Async Copy for operand B
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
-                typename IteratorB::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorB::AccessType*>(this->smem_iterator_B_.get());
+    // Pair of fragments used to overlap shared memory loads and math
+    // instructions
+    WarpFragmentA warp_frag_A[2];
+    WarpFragmentB warp_frag_B[2];
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
-                    int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value
-                                          * IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector
-                                          / 8;
+    Operator warp_mma;
 
-                    cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
-                        dst_ptr + v, iterator_B.get(), iterator_B.valid());
+    this->warp_tile_iterator_A_.set_kgroup_index(0);
+    this->warp_tile_iterator_B_.set_kgroup_index(0);
 
-                    ++iterator_B;
-                }
+    this->warp_tile_iterator_A_.load(warp_frag_A[0]);
+    this->warp_tile_iterator_B_.load(warp_frag_B[0]);
 
-                ++this->smem_iterator_B_;
-            }
+    ++this->warp_tile_iterator_A_;
+    ++this->warp_tile_iterator_B_;
 
-            // Move to the next stage
-            iterator_A.add_tile_offset({0, 1});
-            iterator_B.add_tile_offset({1, 0});
+    iterator_A.clear_mask(gemm_k_iterations == 0);
+    iterator_B.clear_mask(gemm_k_iterations == 0);
 
-            this->smem_iterator_A_.add_tile_offset({0, 1});
-            this->smem_iterator_B_.add_tile_offset({1, 0});
+    int smem_write_stage_idx = Base::kStages - 1;
+    int smem_read_stage_idx = 0;
 
-            // Defines the boundary of a stage of cp.async.
-            cutlass::arch::cp_async_fence();
-        }
+    //
+    // Mainloop
+    //
 
-        // Perform accumulation in the 'd' output operand
-        accum = src_accum;
-
-        //
-        // Clear the remaining tiles of SMEM. This is a functional requirement for some kernels
-        // so that all accumulator elements outside the GEMM footprint are zero.
-        //
-
-        if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) {
-
-            /// Iterator to write threadblock-scoped tile of A operand to shared memory
-            SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
-
-            typename IteratorA::AccessType zero_A;
-            zero_A.clear();
-
-            last_smem_iterator_A.set_iteration_index(0);
-
-            // Async Copy for operand A
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
-
-                typename IteratorA::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorA::AccessType*>(last_smem_iterator_A.get());
-
-                *dst_ptr = zero_A;
-
-                ++last_smem_iterator_A;
-            }
-
-            /// Iterator to write threadblock-scoped tile of B operand to shared memory
-            SmemIteratorB                  last_smem_iterator_B(this->smem_iterator_B_);
-            typename IteratorB::AccessType zero_B;
-
-            zero_B.clear();
-            last_smem_iterator_B.set_iteration_index(0);
-
-            // Async Copy for operand B
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
-
-                typename IteratorB::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorB::AccessType*>(last_smem_iterator_B.get());
-
-                *dst_ptr = zero_B;
-
-                ++last_smem_iterator_B;
-            }
-        }
-
-        // Waits until kStages-2 stages have committed.
-        cutlass::arch::cp_async_wait<Base::kStages - 2>();
-        __syncthreads();
-
-        // Pair of fragments used to overlap shared memory loads and math
-        // instructions
-        WarpFragmentA                       warp_frag_A[2];
-        WarpFragmentB                       warp_frag_B[2];
-
-        Operator warp_mma;
-
-        this->warp_tile_iterator_A_.set_kgroup_index(0);
-        this->warp_tile_iterator_B_.set_kgroup_index(0);
-
-        this->warp_tile_iterator_A_.load(warp_frag_A[0]);
-        this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+    CUTLASS_GEMM_LOOP
+    for (; gemm_k_iterations > (-Base::kStages + 1);) {
+      //
+      // Loop over GEMM K dimension
+      //
 
+      // Computes a warp-level GEMM on data held in shared memory
+      // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate
+      CUTLASS_PRAGMA_UNROLL
+      for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
+           ++warp_mma_k) {
+        // Load warp-level tiles from shared memory, wrapping to k offset if
+        // this is the last group as the case may be.
+        this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) %
+                                                     Base::kWarpGemmIterations);
+        this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
         ++this->warp_tile_iterator_A_;
-        ++this->warp_tile_iterator_B_;
+        // static_assert(Base::kNumKIterationsPerWarpBLoad==1,"Base::kNumKIterationsPerWarpBLoad!=1");
+        // static_assert(Base::kWarpGemmIterationsForB==4,"Base::kWarpGemmIterationsForB!=4");
+        const int warp_tileB_k_compute_offset =
+            warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
+        const int warp_tileB_k_load_offset =
+            warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
+        if (warp_tileB_k_compute_offset ==
+            Base::kNumKIterationsPerWarpBLoad - 1) {
+          this->warp_tile_iterator_B_.set_kgroup_index(
+              (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
+          this->warp_tile_iterator_B_.load(
+              warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
+          ++this->warp_tile_iterator_B_;
+        }
+        // TODO(wangbojun) lds_converter can be remove for int8 B input
+        typename TransformBAfterLDS::result_type converted_frag_B =
+            lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
 
-        iterator_A.clear_mask(gemm_k_iterations == 0);
-        iterator_B.clear_mask(gemm_k_iterations == 0);
+        // TODO(zhengzekang)
+        // run_warp_mma(
+        run_ampere_warp_mma(warp_mma,
+                            accum,
+                            warp_frag_A[warp_mma_k % 2],
+                            converted_frag_B,
+                            accum,
+                            warp_tileB_k_compute_offset);
+        // if(threadIdx.x==0  && threadIdx.y==0  && threadIdx.z==0 &&
+        //     blockIdx.x==0  && blockIdx.y==0  && blockIdx.z==0){
+        //     printf("### run_warp_mma: "
+        //         "%d \n",
+        //         reinterpret_cast<int32_t &>(accum));
+        // }
+        // Issue global->shared copies for the this stage
+        if (warp_mma_k < Base::kWarpGemmIterations - 1) {
+          int group_start_iteration_A, group_start_iteration_B;
 
-        int smem_write_stage_idx = Base::kStages - 1;
-        int smem_read_stage_idx  = 0;
+          group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
+          group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
 
-        //
-        // Mainloop
-        //
-
-        CUTLASS_GEMM_LOOP
-        for (; gemm_k_iterations > (-Base::kStages + 1);) {
-            //
-            // Loop over GEMM K dimension
-            //
-
-            // Computes a warp-level GEMM on data held in shared memory
-            // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate
-            CUTLASS_PRAGMA_UNROLL
-            for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) {
-
-                // Load warp-level tiles from shared memory, wrapping to k offset if
-                // this is the last group as the case may be.
-                this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
-                this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
-                ++this->warp_tile_iterator_A_;
-                // static_assert(Base::kNumKIterationsPerWarpBLoad==1,"Base::kNumKIterationsPerWarpBLoad!=1");
-                // static_assert(Base::kWarpGemmIterationsForB==4,"Base::kWarpGemmIterationsForB!=4");
-                const int warp_tileB_k_compute_offset = warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
-                const int warp_tileB_k_load_offset    = warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
-                if (warp_tileB_k_compute_offset == Base::kNumKIterationsPerWarpBLoad - 1) {
-                    this->warp_tile_iterator_B_.set_kgroup_index((warp_tileB_k_load_offset + 1)
-                                                                 % Base::kWarpGemmIterationsForB);
-                    this->warp_tile_iterator_B_.load(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
-                    ++this->warp_tile_iterator_B_;
-                }
-                // TODO(wangbojun) lds_converter can be remove for int8 B input
-                typename TransformBAfterLDS::result_type converted_frag_B =
-                    lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
-
-                // TODO(zhengzekang)
-                // run_warp_mma(
-                run_ampere_warp_mma(
-                    warp_mma, accum, warp_frag_A[warp_mma_k % 2], converted_frag_B, accum, warp_tileB_k_compute_offset);
-                // if(threadIdx.x==0  && threadIdx.y==0  && threadIdx.z==0 &&
-                //     blockIdx.x==0  && blockIdx.y==0  && blockIdx.z==0){
-                //     printf("### run_warp_mma: "
-                //         "%d \n",
-                //         reinterpret_cast<int32_t &>(accum));
-                // }
-                // Issue global->shared copies for the this stage
-                if (warp_mma_k < Base::kWarpGemmIterations - 1) {
-                    int group_start_iteration_A, group_start_iteration_B;
-
-                    group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
-                    group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
-
-                    copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
-                }
-
-                if (warp_mma_k + 2 == Base::kWarpGemmIterations) {
-                    int group_start_iteration_A, group_start_iteration_B;
-                    group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
-                    group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
-
-                    copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
-
-                    // Inserts a memory fence between stages of cp.async instructions.
-                    cutlass::arch::cp_async_fence();
-
-                    // Waits until kStages-2 stages have committed.
-                    arch::cp_async_wait<Base::kStages - 2>();
-                    __syncthreads();
-
-                    // Move to the next stage
-                    iterator_A.add_tile_offset({0, 1});
-                    iterator_B.add_tile_offset({1, 0});
-
-                    this->smem_iterator_A_.add_tile_offset({0, 1});
-                    this->smem_iterator_B_.add_tile_offset({1, 0});
-
-                    // Add negative offsets to return iterators to the 'start' of the
-                    // circular buffer in shared memory
-                    if (smem_write_stage_idx == (Base::kStages - 1)) {
-                        this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
-                        this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
-                        smem_write_stage_idx = 0;
-                    }
-                    else {
-                        ++smem_write_stage_idx;
-                    }
-
-                    if (smem_read_stage_idx == (Base::kStages - 1)) {
-                        this->warp_tile_iterator_A_.add_tile_offset(
-                            {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
-                        this->warp_tile_iterator_B_.add_tile_offset(
-                            {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0});
-                        smem_read_stage_idx = 0;
-                    }
-                    else {
-                        ++smem_read_stage_idx;
-                    }
-
-                    --gemm_k_iterations;
-                    iterator_A.clear_mask(gemm_k_iterations == 0);
-                    iterator_B.clear_mask(gemm_k_iterations == 0);
-                }
-            }
+          copy_tiles_and_advance(iterator_A,
+                                 iterator_B,
+                                 group_start_iteration_A,
+                                 group_start_iteration_B);
         }
 
-        if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
-            // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop
-            cutlass::arch::cp_async_fence();
-            cutlass::arch::cp_async_wait<0>();
-            __syncthreads();
+        if (warp_mma_k + 2 == Base::kWarpGemmIterations) {
+          int group_start_iteration_A, group_start_iteration_B;
+          group_start_iteration_A =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
+          group_start_iteration_B =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
+
+          copy_tiles_and_advance(iterator_A,
+                                 iterator_B,
+                                 group_start_iteration_A,
+                                 group_start_iteration_B);
+
+          // Inserts a memory fence between stages of cp.async instructions.
+          cutlass::arch::cp_async_fence();
+
+          // Waits until kStages-2 stages have committed.
+          arch::cp_async_wait<Base::kStages - 2>();
+          __syncthreads();
+
+          // Move to the next stage
+          iterator_A.add_tile_offset({0, 1});
+          iterator_B.add_tile_offset({1, 0});
+
+          this->smem_iterator_A_.add_tile_offset({0, 1});
+          this->smem_iterator_B_.add_tile_offset({1, 0});
+
+          // Add negative offsets to return iterators to the 'start' of the
+          // circular buffer in shared memory
+          if (smem_write_stage_idx == (Base::kStages - 1)) {
+            this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
+            this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
+            smem_write_stage_idx = 0;
+          } else {
+            ++smem_write_stage_idx;
+          }
+
+          if (smem_read_stage_idx == (Base::kStages - 1)) {
+            this->warp_tile_iterator_A_.add_tile_offset(
+                {0,
+                 -Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterations});
+            this->warp_tile_iterator_B_.add_tile_offset(
+                {-Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterationsForB,
+                 0});
+            smem_read_stage_idx = 0;
+          } else {
+            ++smem_read_stage_idx;
+          }
+
+          --gemm_k_iterations;
+          iterator_A.clear_mask(gemm_k_iterations == 0);
+          iterator_B.clear_mask(gemm_k_iterations == 0);
         }
+      }
     }
+
+    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+      // commit and drain all pending and predicated LDGSTS pnz from the GEMM
+      // mainloop
+      cutlass::arch::cp_async_fence();
+      cutlass::arch::cp_async_wait<0>();
+      __syncthreads();
+    }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_pipelined.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_pipelined.h
index 2b7312f365..f7c9cac33f 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_pipelined.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/int8_mma_pipelined.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
@@ -74,18 +75,21 @@ namespace threadblock {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions.
-template<
+/// Structure to compute the matrix product targeting CUDA cores and SIMT math
+/// instructions.
+template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Iterates over tiles of A operand in global memory
-    //  (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator)
+    //  (concept: ReadableTileIterator | ForwardTileIterator |
+    //  MaskedTileIterator)
     typename IteratorA_,
     /// Iterates over tiles of A operand in shared memory
     /// (concept: WriteableTileIterator | RandomAccessTileIterator)
     typename SmemIteratorA_,
     /// Iterates over tiles of B operand in global memory
-    //  (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator)
+    //  (concept: ReadableTileIterator | ForwardTileIterator |
+    //  MaskedTileIterator)
     typename IteratorB_,
     /// Iterates over tiles of B operand in shared memory
     /// (concept: WriteableTileIterator | RandomAccessTileIterator)
@@ -102,250 +106,275 @@ template<
     typename TransformBAfterLDS_,
     /// Used for partial specialization
     typename Enable = bool>
-class Int8InterleavedMmaPipelined: public Int8InterleavedMmaBase<Shape_, Policy_, 2> {
-public:
-    ///< Base class
-    using Base = Int8InterleavedMmaBase<Shape_, Policy_, 2>;
-    using Shape     = Shape_;      ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using IteratorA = IteratorA_;  ///< Iterates over tiles of A operand in global memory
-    using IteratorB = IteratorB_;  ///< Iterates over tiles of B operand in global memory
-    using ElementC  = ElementC_;   ///< Data type of accumulator matrix
-    using LayoutC   = LayoutC_;    ///< Layout of accumulator matrix
-    using Policy    = Policy_;     ///< Policy describing tuning details
+class Int8InterleavedMmaPipelined
+    : public Int8InterleavedMmaBase<Shape_, Policy_, 2> {
+ public:
+  ///< Base class
+  using Base = Int8InterleavedMmaBase<Shape_, Policy_, 2>;
+  using Shape =
+      Shape_;  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using IteratorA =
+      IteratorA_;  ///< Iterates over tiles of A operand in global memory
+  using IteratorB =
+      IteratorB_;  ///< Iterates over tiles of B operand in global memory
+  using ElementC = ElementC_;  ///< Data type of accumulator matrix
+  using LayoutC = LayoutC_;    ///< Layout of accumulator matrix
+  using Policy = Policy_;      ///< Policy describing tuning details
 
-    using SmemIteratorA     = SmemIteratorA_;
-    using SmemIteratorB     = SmemIteratorB_;
+  using SmemIteratorA = SmemIteratorA_;
+  using SmemIteratorB = SmemIteratorB_;
 
-    using TransformBAfterLDG = TransformBAfterLDG_;
-    using TransformBAfterLDS = TransformBAfterLDS_;
+  using TransformBAfterLDG = TransformBAfterLDG_;
+  using TransformBAfterLDS = TransformBAfterLDS_;
 
-    //
-    // Dependent types
-    //
+  //
+  // Dependent types
+  //
 
-    /// Fragment of operand A loaded from global memory
-    using FragmentA = typename IteratorA::Fragment;
+  /// Fragment of operand A loaded from global memory
+  using FragmentA = typename IteratorA::Fragment;
 
-    /// Fragment of operand B loaded from global memory
-    using FragmentB = typename IteratorB::Fragment;
+  /// Fragment of operand B loaded from global memory
+  using FragmentB = typename IteratorB::Fragment;
 
-    /// Fragment of accumulator tile
-    using FragmentC = typename Policy::Operator::FragmentC;
+  /// Fragment of accumulator tile
+  using FragmentC = typename Policy::Operator::FragmentC;
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
 
-    /// Obtain the arch tag from the warp-level operator
-    using ArchTag = typename Policy::Operator::ArchTag;
+  /// Obtain the arch tag from the warp-level operator
+  using ArchTag = typename Policy::Operator::ArchTag;
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = Operator::kTransformA;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = Operator::kTransformA;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = Operator::kTransformB;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = Operator::kTransformB;
 
-    // staticaly assert kStages for Int8InterleavedMmaPipelined is two (Double-buffered pipeline)
-    static_assert((Base::kStages == 2), "Int8InterleavedMmaPipelined requires kStages set to value 2");
+  // staticaly assert kStages for Int8InterleavedMmaPipelined is two
+  // (Double-buffered pipeline)
+  static_assert((Base::kStages == 2),
+                "Int8InterleavedMmaPipelined requires kStages set to value 2");
 
-private:
-    using WarpFragmentA = typename Operator::FragmentA;
-    using WarpFragmentB = typename Operator::FragmentB;
-    using ElementB          = typename IteratorB::Element;
-    using LayoutDetailsForB = kernel::LayoutDetailsB<ElementB, ArchTag>;
+ private:
+  using WarpFragmentA = typename Operator::FragmentA;
+  using WarpFragmentB = typename Operator::FragmentB;
+  using ElementB = typename IteratorB::Element;
+  using LayoutDetailsForB = kernel::LayoutDetailsB<ElementB, ArchTag>;
 
-    static constexpr bool RequiresTileInterleave =
-        layout::IsColumnMajorTileInterleave<typename LayoutDetailsForB::Layout>::value;
-    static_assert(!RequiresTileInterleave || (RequiresTileInterleave && (Shape::kK == LayoutDetailsForB::ThreadblockK)),
-                  "Layout K must match threadblockK");
+  static constexpr bool RequiresTileInterleave =
+      layout::IsColumnMajorTileInterleave<
+          typename LayoutDetailsForB::Layout>::value;
+  static_assert(!RequiresTileInterleave ||
+                    (RequiresTileInterleave &&
+                     (Shape::kK == LayoutDetailsForB::ThreadblockK)),
+                "Layout K must match threadblockK");
 
-protected:
-    /// Iterator to write threadblock-scoped tile of A operand to shared memory
-    SmemIteratorA smem_iterator_A_;
+ protected:
+  /// Iterator to write threadblock-scoped tile of A operand to shared memory
+  SmemIteratorA smem_iterator_A_;
 
-    /// Iterator to write threadblock-scoped tile of B operand to shared memory
-    SmemIteratorB smem_iterator_B_;
+  /// Iterator to write threadblock-scoped tile of B operand to shared memory
+  SmemIteratorB smem_iterator_B_;
 
-public:
-    /// Construct from tensor references
-    CUTLASS_DEVICE
-    Int8InterleavedMmaPipelined(typename Base::SharedStorage&
-                       shared_storage,  ///< Shared storage needed for internal use by threadblock-scoped GEMM
-                   int thread_idx,      ///< ID within the threadblock
-                   int warp_idx,        ///< ID of warp
-                   int lane_idx         ///< ID of each thread within a warp
-                   ):
-        Base(shared_storage, thread_idx, warp_idx, lane_idx),
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  Int8InterleavedMmaPipelined(
+      typename Base::SharedStorage&
+          shared_storage,  ///< Shared storage needed for internal use by
+                           ///< threadblock-scoped GEMM
+      int thread_idx,      ///< ID within the threadblock
+      int warp_idx,        ///< ID of warp
+      int lane_idx         ///< ID of each thread within a warp
+      )
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
         smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
-        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx)
-    {
+        smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) {
+    // Compute warp location within threadblock tile by mapping the warp_id to
+    // three coordinates:
+    //   _m: the warp's position within the threadblock along the M dimension
+    //   _n: the warp's position within the threadblock along the N dimension
+    //   _k: the warp's position within the threadblock along the K dimension
 
-        // Compute warp location within threadblock tile by mapping the warp_id to
-        // three coordinates:
-        //   _m: the warp's position within the threadblock along the M dimension
-        //   _n: the warp's position within the threadblock along the N dimension
-        //   _k: the warp's position within the threadblock along the K dimension
+    int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
 
-        int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
-        int warp_idx_k  = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
+    int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
 
-        int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
-        int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
+    // Add per-warp offsets in units of warp-level tiles
+    this->warp_tile_iterator_A_.add_tile_offset(
+        {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
+    this->warp_tile_iterator_B_.add_tile_offset(
+        {Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+  }
 
-        // Add per-warp offsets in units of warp-level tiles
-        this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
-        this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
-    }
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  CUTLASS_DEVICE
+  void operator()(
+      int gemm_k_iterations,  ///< number of iterations of the mainloop
+      FragmentC& accum,       ///< destination accumulator tile
+      IteratorA iterator_A,   ///< iterator over A operand in global memory
+      IteratorB iterator_B,   ///< iterator over B operand in global memory
+      FragmentC const& src_accum) {  ///< source accumulator tile
 
-    /// Perform a threadblock-scoped matrix multiply-accumulate
-    CUTLASS_DEVICE
-    void operator()(int              gemm_k_iterations,  ///< number of iterations of the mainloop
-                    FragmentC&       accum,              ///< destination accumulator tile
-                    IteratorA        iterator_A,         ///< iterator over A operand in global memory
-                    IteratorB        iterator_B,         ///< iterator over B operand in global memory
-                    FragmentC const& src_accum)
-    {  ///< source accumulator tile
+    //
+    // Prologue
+    //
+    TransformBAfterLDG ldg_converter;
+    TransformBAfterLDS lds_converter;
 
-        //
-        // Prologue
-        //
-        TransformBAfterLDG ldg_converter;
-        TransformBAfterLDS lds_converter;
+    using TransformA = NumericArrayConverter<typename WarpFragmentA::Element,
+                                             typename FragmentA::Element,
+                                             FragmentA::kElements>;
 
-        using TransformA =
-            NumericArrayConverter<typename WarpFragmentA::Element, typename FragmentA::Element, FragmentA::kElements>;
+    // These transforms are mainly to handle when we have bfloat activations and
+    // weights in GMEM and want to issue HMMA on architectures older than
+    // Ampere. We will convert to FP16 before STS.
+    TransformA transformA;
 
-        // These transforms are mainly to handle when we have bfloat activations and weights in GMEM and want
-        // to issue HMMA on architectures older than Ampere. We will convert to FP16 before STS.
-        TransformA     transformA;
+    // Perform accumulation in the 'd' output operand
+    accum = src_accum;
 
-        // Perform accumulation in the 'd' output operand
-        accum = src_accum;
+    FragmentA tb_frag_A;
+    FragmentB tb_frag_B;
 
-        FragmentA     tb_frag_A;
-        FragmentB     tb_frag_B;
+    tb_frag_A.clear();
+    tb_frag_B.clear();
 
+    // The last kblock is loaded in the prolog
+    iterator_A.load(tb_frag_A);
+    iterator_B.load(tb_frag_B);
 
-        tb_frag_A.clear();
-        tb_frag_B.clear();
+    ++iterator_A;
+    ++iterator_B;
 
-        // The last kblock is loaded in the prolog
-        iterator_A.load(tb_frag_A);
-        iterator_B.load(tb_frag_B);
+    this->smem_iterator_A_.store(transformA(tb_frag_A));
+    this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
 
-        ++iterator_A;
-        ++iterator_B;
+    ++this->smem_iterator_A_;
+    ++this->smem_iterator_B_;
 
-        this->smem_iterator_A_.store(transformA(tb_frag_A));
-        this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
+    __syncthreads();
 
-        ++this->smem_iterator_A_;
-        ++this->smem_iterator_B_;
+    // Pair of fragments used to overlap shared memory loads and math
+    // instructions
+    WarpFragmentA warp_frag_A[2];
+    WarpFragmentB warp_frag_B[2];
 
-        __syncthreads();
+    this->warp_tile_iterator_A_.set_kgroup_index(0);
+    this->warp_tile_iterator_B_.set_kgroup_index(0);
 
-        // Pair of fragments used to overlap shared memory loads and math instructions
-        WarpFragmentA warp_frag_A[2];
-        WarpFragmentB warp_frag_B[2];
+    this->warp_tile_iterator_A_.load(warp_frag_A[0]);
+    this->warp_tile_iterator_B_.load(warp_frag_B[0]);
 
-        this->warp_tile_iterator_A_.set_kgroup_index(0);
-        this->warp_tile_iterator_B_.set_kgroup_index(0);
+    ++this->warp_tile_iterator_A_;
+    ++this->warp_tile_iterator_B_;
 
-        this->warp_tile_iterator_A_.load(warp_frag_A[0]);
-        this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+    Operator warp_mma;
 
-        ++this->warp_tile_iterator_A_;
-        ++this->warp_tile_iterator_B_;
+    int smem_write_stage_idx = 1;
 
-        Operator warp_mma;
+    // Avoid reading out of bounds
+    iterator_A.clear_mask(gemm_k_iterations <= 1);
+    iterator_B.clear_mask(gemm_k_iterations <= 1);
 
-        int smem_write_stage_idx = 1;
+    // Issue loads during the first warp-level matrix multiply-add *AFTER*
+    // issuing shared memory loads (which have the tighest latency requirement).
 
-        // Avoid reading out of bounds
-        iterator_A.clear_mask(gemm_k_iterations <= 1);
-        iterator_B.clear_mask(gemm_k_iterations <= 1);
+    //
+    // Mainloop
+    //
 
-        // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing
-        // shared memory loads (which have the tighest latency requirement).
+    // Note: The main loop does not support Base::kWarpGemmIterations == 2.
+    CUTLASS_GEMM_LOOP
+    for (; gemm_k_iterations > 0; --gemm_k_iterations) {
+      //
+      // Loop over GEMM K dimension
+      //
 
-        //
-        // Mainloop
-        //
+      CUTLASS_PRAGMA_UNROLL
+      for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
+           ++warp_mma_k) {
+        // Load warp-level tiles from shared memory, wrapping to k offset if
+        // this is the last group as the case may be.
 
-        // Note: The main loop does not support Base::kWarpGemmIterations == 2.
-        CUTLASS_GEMM_LOOP
-        for (; gemm_k_iterations > 0; --gemm_k_iterations) {
-            //
-            // Loop over GEMM K dimension
-            //
+        if (warp_mma_k == Base::kWarpGemmIterations - 1) {
+          // Write fragments to shared memory
+          this->smem_iterator_A_.store(transformA(tb_frag_A));
 
-            CUTLASS_PRAGMA_UNROLL
-            for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) {
+          this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
 
-                // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group
-                // as the case may be.
+          __syncthreads();
 
-                if (warp_mma_k == Base::kWarpGemmIterations - 1) {
+          ++this->smem_iterator_A_;
+          ++this->smem_iterator_B_;
 
-                    // Write fragments to shared memory
-                    this->smem_iterator_A_.store(transformA(tb_frag_A));
+          // Add negative offsets to return iterators to the 'start' of the
+          // circular buffer in shared memory
+          if (smem_write_stage_idx == 1) {
+            this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
+            this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
+          } else {
+            this->warp_tile_iterator_A_.add_tile_offset(
+                {0,
+                 -Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterations});
+            this->warp_tile_iterator_B_.add_tile_offset(
+                {-Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterationsForB,
+                 0});
+          }
 
-                    this->smem_iterator_B_.store(ldg_converter(tb_frag_B));
-
-                    __syncthreads();
-
-                    ++this->smem_iterator_A_;
-                    ++this->smem_iterator_B_;
-
-                    // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory
-                    if (smem_write_stage_idx == 1) {
-                        this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
-                        this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
-                    }
-                    else {
-                        this->warp_tile_iterator_A_.add_tile_offset(
-                            {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
-                        this->warp_tile_iterator_B_.add_tile_offset(
-                            {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0});
-                    }
-
-                    smem_write_stage_idx ^= 1;
-                }
-
-                this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
-                this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
-                ++this->warp_tile_iterator_A_;
-
-                const int warp_tileB_k_compute_offset = warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
-                const int warp_tileB_k_load_offset    = warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
-                // We are just about to finish computing on a fragment of B, so initiate the load for the next fragment.
-                if (warp_tileB_k_compute_offset == Base::kNumKIterationsPerWarpBLoad - 1) {
-                    this->warp_tile_iterator_B_.set_kgroup_index((warp_tileB_k_load_offset + 1)
-                                                                 % Base::kWarpGemmIterationsForB);
-                    this->warp_tile_iterator_B_.load(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
-                    ++this->warp_tile_iterator_B_;
-                }
-
-                if (warp_mma_k == 0) {
-
-                    iterator_A.load(tb_frag_A);
-                    iterator_B.load(tb_frag_B);
-
-                    ++iterator_A;
-                    ++iterator_B;
-
-                    // Avoid reading out of bounds if this was the last loop iteration
-                    iterator_A.clear_mask(gemm_k_iterations <= 2);
-                    iterator_B.clear_mask(gemm_k_iterations <= 2);
-                }
-
-                typename TransformBAfterLDS::result_type converted_frag_B =
-                    lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
-                // run_warp_mma(
-                run_ampere_warp_mma(
-                    warp_mma, accum, warp_frag_A[warp_mma_k % 2], converted_frag_B, accum, warp_tileB_k_compute_offset);
-            }
+          smem_write_stage_idx ^= 1;
         }
+
+        this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) %
+                                                     Base::kWarpGemmIterations);
+        this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
+        ++this->warp_tile_iterator_A_;
+
+        const int warp_tileB_k_compute_offset =
+            warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
+        const int warp_tileB_k_load_offset =
+            warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
+        // We are just about to finish computing on a fragment of B, so initiate
+        // the load for the next fragment.
+        if (warp_tileB_k_compute_offset ==
+            Base::kNumKIterationsPerWarpBLoad - 1) {
+          this->warp_tile_iterator_B_.set_kgroup_index(
+              (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
+          this->warp_tile_iterator_B_.load(
+              warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
+          ++this->warp_tile_iterator_B_;
+        }
+
+        if (warp_mma_k == 0) {
+          iterator_A.load(tb_frag_A);
+          iterator_B.load(tb_frag_B);
+
+          ++iterator_A;
+          ++iterator_B;
+
+          // Avoid reading out of bounds if this was the last loop iteration
+          iterator_A.clear_mask(gemm_k_iterations <= 2);
+          iterator_B.clear_mask(gemm_k_iterations <= 2);
+        }
+
+        typename TransformBAfterLDS::result_type converted_frag_B =
+            lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
+        // run_warp_mma(
+        run_ampere_warp_mma(warp_mma,
+                            accum,
+                            warp_frag_A[warp_mma_k % 2],
+                            converted_frag_B,
+                            accum,
+                            warp_tileB_k_compute_offset);
+      }
     }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/nf4_int8_mma_base.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/nf4_int8_mma_base.h
index 413e53a9a3..5048bc52e0 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/nf4_int8_mma_base.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/nf4_int8_mma_base.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,14 +18,15 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
@@ -65,7 +66,7 @@ namespace threadblock {
 ////////////////////////////////////////////////////////////////////////////////
 /// Structure to compute the matrix product targeting CUDA cores and SIMT math
 /// instructions.
-template<
+template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Policy describing tuning details (concept: MmaPolicy)
@@ -75,138 +76,142 @@ template<
     /// Used for partial specialization
     typename Enable = bool>
 class Int8Nf4InterleavedMmaBase {
-public:
-    ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using Shape = Shape_;
+ public:
+  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using Shape = Shape_;
 
-    ///< Policy describing tuning details
-    using Policy = Policy_;
+  ///< Policy describing tuning details
+  using Policy = Policy_;
 
+  //
+  // Dependent types
+  //
+
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
+
+  /// Shape describing the overall GEMM computed from shared memory
+  /// by each warp.
+  using WarpGemm = typename Policy::Operator::Shape;
+
+  /// Shape describing the number of warps filling the CTA
+  using WarpCount = GemmShape<Shape::kM / WarpGemm::kM,
+                              Shape::kN / WarpGemm::kN,
+                              Shape::kK / WarpGemm::kK>;
+
+  /// Number of warp-level GEMM operations
+  static int const kWarpGemmIterations =
+      (WarpGemm::kK / Operator::Policy::MmaShape::kK);
+  static_assert(Operator::IteratorB::InstructionShape::kRow >=
+                    Operator::InstructionShape::kK,
+                "");
+  static constexpr int kNumKIterationsPerWarpBLoad =
+      Operator::IteratorB::InstructionShape::kRow /
+      Operator::InstructionShape::kK;
+  static_assert(!(kWarpGemmIterations % kNumKIterationsPerWarpBLoad), "");
+  static constexpr int kWarpGemmIterationsForB =
+      kWarpGemmIterations / kNumKIterationsPerWarpBLoad;
+
+  /// Number of stages
+  static int const kStages = Stages;
+
+  /// Tensor reference to the A operand
+  using TensorRefA =
+      TensorRef<typename Operator::ElementA, typename Operator::LayoutA>;
+
+  /// Tensor reference to the B operand
+  using TensorRefB =
+      TensorRef<typename Operator::ElementB, typename Operator::LayoutB>;
+
+  //
+  // Nested structs
+  //
+
+  /// Shared storage object needed by threadblock-scoped GEMM
+  class SharedStorage {
+   public:
     //
-    // Dependent types
+    // Type definitions
     //
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+    /// Shape of the A matrix operand in shared memory
+    using ShapeA =
+        MatrixShape<Shape::kM + Policy::SmemPaddingA::kRow,
+                    Shape::kK * kStages + Policy::SmemPaddingA::kColumn>;
 
-    /// Shape describing the overall GEMM computed from shared memory
-    /// by each warp.
-    using WarpGemm = typename Policy::Operator::Shape;
+    /// Shape of the B matrix operand in shared memory
+    using ShapeB = MatrixShape<Shape::kK * kStages + Policy::SmemPaddingB::kRow,
+                               Shape::kN + Policy::SmemPaddingB::kColumn>;
 
-    /// Shape describing the number of warps filling the CTA
-    using WarpCount = GemmShape<Shape::kM / WarpGemm::kM, Shape::kN / WarpGemm::kN, Shape::kK / WarpGemm::kK>;
-
-    /// Number of warp-level GEMM operations
-    static int const kWarpGemmIterations = (WarpGemm::kK / Operator::Policy::MmaShape::kK);
-    static_assert(Operator::IteratorB::InstructionShape::kRow>=Operator::InstructionShape::kK,"");
-    static constexpr int kNumKIterationsPerWarpBLoad =
-        Operator::IteratorB::InstructionShape::kRow / Operator::InstructionShape::kK;
-    static_assert(!(kWarpGemmIterations % kNumKIterationsPerWarpBLoad), "");
-    static constexpr int kWarpGemmIterationsForB = kWarpGemmIterations / kNumKIterationsPerWarpBLoad;
-
-    /// Number of stages
-    static int const kStages = Stages;
-
-    /// Tensor reference to the A operand
-    using TensorRefA = TensorRef<typename Operator::ElementA, typename Operator::LayoutA>;
-
-    /// Tensor reference to the B operand
-    using TensorRefB = TensorRef<typename Operator::ElementB, typename Operator::LayoutB>;
-
-    //
-    // Nested structs
-    //
-
-    /// Shared storage object needed by threadblock-scoped GEMM
-    class SharedStorage {
-    public:
-        //
-        // Type definitions
-        //
-
-        /// Shape of the A matrix operand in shared memory
-        using ShapeA =
-            MatrixShape<Shape::kM + Policy::SmemPaddingA::kRow, Shape::kK * kStages + Policy::SmemPaddingA::kColumn>;
-
-        /// Shape of the B matrix operand in shared memory
-        using ShapeB =
-            MatrixShape<Shape::kK * kStages + Policy::SmemPaddingB::kRow, Shape::kN + Policy::SmemPaddingB::kColumn>;
-
-    public:
-        //
-        // Data members
-        //
-
-        /// Buffer for A operand
-        AlignedBuffer<typename Operator::ElementA, ShapeA::kCount> operand_A;
-
-        /// Buffer for B operand
-        AlignedBuffer<typename Operator::ElementB, ShapeB::kCount> operand_B;
-
-        /// Buffer to hold scales for threadblock
-        AlignedBuffer<int32_t, 16> operand_nf4_look_up_table;
-
-    public:
-        //
-        // Methods
-        //
-
-        /// Returns a layout object for the A matrix
-        CUTLASS_DEVICE
-        static typename Operator::LayoutA LayoutA()
-        {
-            return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn});
-        }
-
-        /// Returns a layout object for the B matrix
-        CUTLASS_HOST_DEVICE
-        static typename Operator::LayoutB LayoutB()
-        {
-            return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn});
-        }
-
-        /// Returns a TensorRef to the A operand
-        CUTLASS_HOST_DEVICE
-        TensorRefA operand_A_ref()
-        {
-            return TensorRefA{operand_A.data(), LayoutA()};
-        }
-
-        /// Returns a TensorRef to the B operand
-        CUTLASS_HOST_DEVICE
-        TensorRefB operand_B_ref()
-        {
-            return TensorRefB{operand_B.data(), LayoutB()};
-        }
-    };
-
-protected:
+   public:
     //
     // Data members
     //
 
-    /// Iterator to load a warp-scoped tile of A operand from shared memory
-    typename Operator::IteratorA warp_tile_iterator_A_;
+    /// Buffer for A operand
+    AlignedBuffer<typename Operator::ElementA, ShapeA::kCount> operand_A;
 
-    /// Iterator to load a warp-scoped tile of B operand from shared memory
-    typename Operator::IteratorB warp_tile_iterator_B_;
+    /// Buffer for B operand
+    AlignedBuffer<typename Operator::ElementB, ShapeB::kCount> operand_B;
 
-public:
-    /// Construct from tensor references
+    /// Buffer to hold scales for threadblock
+    AlignedBuffer<int32_t, 16> operand_nf4_look_up_table;
+
+   public:
+    //
+    // Methods
+    //
+
+    /// Returns a layout object for the A matrix
     CUTLASS_DEVICE
-    Int8Nf4InterleavedMmaBase(
-        ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        SharedStorage& shared_storage,
-        ///< ID within the threadblock
-        int thread_idx,
-        ///< ID of warp
-        int warp_idx,
-        ///< ID of each thread within a warp
-        int lane_idx):
-        warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx),
-        warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx)
-    {
+    static typename Operator::LayoutA LayoutA() {
+      return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn});
     }
+
+    /// Returns a layout object for the B matrix
+    CUTLASS_HOST_DEVICE
+    static typename Operator::LayoutB LayoutB() {
+      return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn});
+    }
+
+    /// Returns a TensorRef to the A operand
+    CUTLASS_HOST_DEVICE
+    TensorRefA operand_A_ref() {
+      return TensorRefA{operand_A.data(), LayoutA()};
+    }
+
+    /// Returns a TensorRef to the B operand
+    CUTLASS_HOST_DEVICE
+    TensorRefB operand_B_ref() {
+      return TensorRefB{operand_B.data(), LayoutB()};
+    }
+  };
+
+ protected:
+  //
+  // Data members
+  //
+
+  /// Iterator to load a warp-scoped tile of A operand from shared memory
+  typename Operator::IteratorA warp_tile_iterator_A_;
+
+  /// Iterator to load a warp-scoped tile of B operand from shared memory
+  typename Operator::IteratorB warp_tile_iterator_B_;
+
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  Int8Nf4InterleavedMmaBase(
+      ///< Shared storage needed for internal use by threadblock-scoped GEMM
+      SharedStorage& shared_storage,
+      ///< ID within the threadblock
+      int thread_idx,
+      ///< ID of warp
+      int warp_idx,
+      ///< ID of each thread within a warp
+      int lane_idx)
+      : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx),
+        warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {}
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/nf4_int8_mma_multistage.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/nf4_int8_mma_multistage.h
index d124b09cb9..c28e0ed844 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/nf4_int8_mma_multistage.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/threadblock/nf4_int8_mma_multistage.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,18 +18,18 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 
-
 /* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -66,8 +66,9 @@ limitations under the License. */
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 template <typename T>
-[[gnu::warning("your type here")]]
-bool print_type() { return false; }
+[[gnu::warning("your type here")]] bool print_type() {
+  return false;
+}
 
 namespace cutlass {
 namespace gemm {
@@ -77,7 +78,7 @@ namespace threadblock {
 
 /// Structure to compute the matrix product targeting CUDA cores and SIMT math
 /// instructions.
-template<
+template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Iterates over tiles of A operand in global memory
@@ -116,467 +117,601 @@ template<
     SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone,
     /// Used for partial specialization
     typename Enable = bool>
-class Int8Nf4InterleavedMmaMultistage: public Int8Nf4InterleavedMmaBase<Shape_, Policy_, Stages> {
-public:
-    ///< Base class
-    using Base = Int8Nf4InterleavedMmaBase<Shape_, Policy_, Stages>;
-    ///< Size of the Gemm problem - concept: gemm::GemmShape<>
-    using Shape = Shape_;
-    ///< Iterates over tiles of A operand in global memory
-    using IteratorA = IteratorA_;
-    ///< Iterates over tiles of B operand in global memory
-    using IteratorB = IteratorB_;
-    ///< Iterates over tiles of nf4 look up table in global memory
-    using IteratorNF4LookUpTable = IteratorNF4LookUpTable_;
-    using ElementNF4LookUpTable = typename IteratorNF4LookUpTable::Element;
-    using LayoutNF4LookUpTable = typename IteratorNF4LookUpTable::Layout;
+class Int8Nf4InterleavedMmaMultistage
+    : public Int8Nf4InterleavedMmaBase<Shape_, Policy_, Stages> {
+ public:
+  ///< Base class
+  using Base = Int8Nf4InterleavedMmaBase<Shape_, Policy_, Stages>;
+  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
+  using Shape = Shape_;
+  ///< Iterates over tiles of A operand in global memory
+  using IteratorA = IteratorA_;
+  ///< Iterates over tiles of B operand in global memory
+  using IteratorB = IteratorB_;
+  ///< Iterates over tiles of nf4 look up table in global memory
+  using IteratorNF4LookUpTable = IteratorNF4LookUpTable_;
+  using ElementNF4LookUpTable = typename IteratorNF4LookUpTable::Element;
+  using LayoutNF4LookUpTable = typename IteratorNF4LookUpTable::Layout;
 
+  ///< Data type of accumulator matrix
+  using ElementC = ElementC_;
+  ///< Layout of accumulator matrix
+  using LayoutC = LayoutC_;
+  ///< Policy describing tuning details
+  using Policy = Policy_;
 
-    ///< Data type of accumulator matrix
-    using ElementC = ElementC_;
-    ///< Layout of accumulator matrix
-    using LayoutC = LayoutC_;
-    ///< Policy describing tuning details
-    using Policy = Policy_;
+  using SmemIteratorA = SmemIteratorA_;
+  using SmemIteratorB = SmemIteratorB_;
+  using SmemIteratorNF4LookUpTable = SmemIteratorNF4LookUpTable_;
+  static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
+  static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
 
+  using TransformBAfterLDS = TransformBAfterLDS_;
 
-    using SmemIteratorA     = SmemIteratorA_;
-    using SmemIteratorB     = SmemIteratorB_;
-    using SmemIteratorNF4LookUpTable = SmemIteratorNF4LookUpTable_;
-    static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
-    static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
+  //
+  // Dependent types
+  //
 
-    using TransformBAfterLDS = TransformBAfterLDS_;
+  /// Fragment of accumulator tile
+  using FragmentC = typename Policy::Operator::FragmentC;
 
-    //
-    // Dependent types
-    //
+  /// Warp-level Mma
+  using Operator = typename Policy::Operator;
 
-    /// Fragment of accumulator tile
-    using FragmentC = typename Policy::Operator::FragmentC;
+  /// Minimum architecture is Sm80 to support cp.async
+  using ArchTag = arch::Sm80;
 
-    /// Warp-level Mma
-    using Operator = typename Policy::Operator;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = Operator::kTransformA;
 
-    /// Minimum architecture is Sm80 to support cp.async
-    using ArchTag = arch::Sm80;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = Operator::kTransformB;
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = Operator::kTransformA;
+  /// Internal structure exposed for introspection.
+  struct Detail {
+    static_assert(Base::kWarpGemmIterations > 1,
+                  "The pipelined structure requires at least two warp-level "
+                  "GEMM operations.");
+    // static_assert(Base::kWarpGemmIterations==4,"Base::kWarpGemmIterations!=4");
+    /// Number of cp.async instructions to load one stage of operand A
+    static int const AsyncCopyIterationsPerStageA =
+        IteratorA::ThreadMap::Iterations::kCount;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = Operator::kTransformB;
+    /// Number of cp.async instructions to load one stage of operand B
+    static int const AsyncCopyIterationsPerStageB =
+        IteratorB::ThreadMap::Iterations::kCount;
 
-    /// Internal structure exposed for introspection.
-    struct Detail {
+    /// Number of stages
+    static int const kStages = Stages;
 
-        static_assert(Base::kWarpGemmIterations > 1,
-                      "The pipelined structure requires at least two warp-level "
-                      "GEMM operations.");
-        // static_assert(Base::kWarpGemmIterations==4,"Base::kWarpGemmIterations!=4");
-        /// Number of cp.async instructions to load one stage of operand A
-        static int const AsyncCopyIterationsPerStageA = IteratorA::ThreadMap::Iterations::kCount;
+    /// Number of cp.async instructions to load on group of operand A
+    static int const kAccessesPerGroupA =
+        (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
 
-        /// Number of cp.async instructions to load one stage of operand B
-        static int const AsyncCopyIterationsPerStageB = IteratorB::ThreadMap::Iterations::kCount;
+    /// Number of cp.async instructions to load on group of operand B
+    static int const kAccessesPerGroupB =
+        (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) /
+        Base::kWarpGemmIterations;
+  };
 
-        /// Number of stages
-        static int const kStages = Stages;
+ private:
+  using WarpFragmentA = typename Operator::FragmentA;
+  using WarpFragmentB = typename Operator::FragmentB;
+  using ElementB = typename IteratorB::Element;
+  using LayoutDetailsForB = kernel::LayoutDetailsB<ElementB, ArchTag>;
 
-        /// Number of cp.async instructions to load on group of operand A
-        static int const kAccessesPerGroupA =
-            (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
+  static constexpr bool RequiresTileInterleave =
+      layout::IsColumnMajorTileInterleave<
+          typename LayoutDetailsForB::Layout>::value;
+  static_assert(!RequiresTileInterleave ||
+                    (RequiresTileInterleave &&
+                     (Shape::kK == LayoutDetailsForB::ThreadblockK)),
+                "Layout K must match threadblockK");
 
-        /// Number of cp.async instructions to load on group of operand B
-        static int const kAccessesPerGroupB =
-            (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
-    };
+ private:
+  //
+  // Data members
+  //
 
-private:
-    using WarpFragmentA = typename Operator::FragmentA;
-    using WarpFragmentB = typename Operator::FragmentB;
-    using ElementB          = typename IteratorB::Element;
-    using LayoutDetailsForB = kernel::LayoutDetailsB<ElementB, ArchTag>;
+  /// Iterator to write threadblock-scoped tile of A operand to shared memory
+  SmemIteratorA smem_iterator_A_;
 
-    static constexpr bool RequiresTileInterleave =
-        layout::IsColumnMajorTileInterleave<typename LayoutDetailsForB::Layout>::value;
-    static_assert(!RequiresTileInterleave || (RequiresTileInterleave && (Shape::kK == LayoutDetailsForB::ThreadblockK)),
-                  "Layout K must match threadblockK");
+  /// Iterator to write threadblock-scoped tile of B operand to shared memory
+  SmemIteratorB smem_iterator_B_;
 
-private:
-    //
-    // Data members
-    //
+  SmemIteratorNF4LookUpTable smem_iterator_nf4_look_up_table_;
 
-    /// Iterator to write threadblock-scoped tile of A operand to shared memory
-    SmemIteratorA smem_iterator_A_;
-
-    /// Iterator to write threadblock-scoped tile of B operand to shared memory
-    SmemIteratorB smem_iterator_B_;
-
-    SmemIteratorNF4LookUpTable smem_iterator_nf4_look_up_table_;
-
-public:
-    /// Construct from tensor references
-    CUTLASS_DEVICE
-    Int8Nf4InterleavedMmaMultistage(
-        ///< Shared storage needed for internal use by threadblock-scoped GEMM
-        typename Base::SharedStorage& shared_storage,
-        ///< ID within the threadblock
-        int thread_idx,
-        ///< ID of warp
-        int warp_idx,
-        ///< ID of each thread within a warp
-        int lane_idx):
-        Base(shared_storage, thread_idx, warp_idx, lane_idx),
+ public:
+  /// Construct from tensor references
+  CUTLASS_DEVICE
+  Int8Nf4InterleavedMmaMultistage(
+      ///< Shared storage needed for internal use by threadblock-scoped GEMM
+      typename Base::SharedStorage& shared_storage,
+      ///< ID within the threadblock
+      int thread_idx,
+      ///< ID of warp
+      int warp_idx,
+      ///< ID of each thread within a warp
+      int lane_idx)
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
         smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
         smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx),
-        smem_iterator_nf4_look_up_table_(LayoutNF4LookUpTable(16),
-                                         shared_storage.operand_nf4_look_up_table.data(),
-                                         {1, 16},
-                                         thread_idx)
-{
-        // Compute warp location within threadblock tile by mapping the warp_id to
-        // three coordinates:
-        //   _m: the warp's position within the threadblock along the M dimension
-        //   _n: the warp's position within the threadblock along the N dimension
-        //   _k: the warp's position within the threadblock along the K dimension
+        smem_iterator_nf4_look_up_table_(
+            LayoutNF4LookUpTable(16),
+            shared_storage.operand_nf4_look_up_table.data(),
+            {1, 16},
+            thread_idx) {
+    // Compute warp location within threadblock tile by mapping the warp_id to
+    // three coordinates:
+    //   _m: the warp's position within the threadblock along the M dimension
+    //   _n: the warp's position within the threadblock along the N dimension
+    //   _k: the warp's position within the threadblock along the K dimension
 
-        int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
-        int warp_idx_k  = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
+    int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);
 
-        int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
-        int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
+    int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
+    int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;
 
-        // Add per-warp offsets in units of warp-level tiles
-        this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
-        this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
+    // Add per-warp offsets in units of warp-level tiles
+    this->warp_tile_iterator_A_.add_tile_offset(
+        {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
+    this->warp_tile_iterator_B_.add_tile_offset(
+        {Base::kWarpGemmIterationsForB * warp_idx_k, warp_idx_n});
 
+    // if((threadIdx.x % 32) == 0){
+    //     printf("#### %d-%d-%d-%d-%d-%d, gmem_ptr_nf4_look_up_table:%p,
+    //     kSrcBytesNf4:%d \n",
+    //             blockIdx.x, blockIdx.y, blockIdx.z,
+    //             threadIdx.x, threadIdx.y, threadIdx.z,
+    //             gmem_ptr_nf4_look_up_table,
+    //             kSrcBytesNf4);
+    // }
+    // cutlass::arch::cp_async_zfill<kSrcBytesNf4,
+    // cutlass::arch::CacheOperation::Global>(
+    //                 dst_ptr_nf4_look_up_table, gmem_ptr_nf4_look_up_table,
+    //                 iterator_nf4_look_up_table.valid());
+  }
 
-        // if((threadIdx.x % 32) == 0){
-        //     printf("#### %d-%d-%d-%d-%d-%d, gmem_ptr_nf4_look_up_table:%p, kSrcBytesNf4:%d \n",
-        //             blockIdx.x, blockIdx.y, blockIdx.z,
-        //             threadIdx.x, threadIdx.y, threadIdx.z,
-        //             gmem_ptr_nf4_look_up_table,
-        //             kSrcBytesNf4);
-        // }
-        // cutlass::arch::cp_async_zfill<kSrcBytesNf4, cutlass::arch::CacheOperation::Global>(
-        //                 dst_ptr_nf4_look_up_table, gmem_ptr_nf4_look_up_table, iterator_nf4_look_up_table.valid());
+  CUTLASS_DEVICE
+  void copy_tiles_and_advance(IteratorA& iterator_A,
+                              IteratorB& iterator_B,
+                              int group_start_A = 0,
+                              int group_start_B = 0) {
+    iterator_A.set_iteration_index(group_start_A *
+                                   IteratorA::kAccessesPerVector);
+    this->smem_iterator_A_.set_iteration_index(group_start_A);
+
+    // Async Copy for operand A
+    CUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) {
+      if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                this->smem_iterator_A_.get());
+
+        int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
+                              IteratorA::ThreadMap::kElementsPerAccess /
+                              IteratorA::kAccessesPerVector / 8;
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
+          auto gmem_ptr = iterator_A.get();
+
+          if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
+                dst_ptr + v, gmem_ptr, iterator_A.valid());
+          } else {
+            cutlass::arch::cp_async<kSrcBytes, kCacheOpA>(
+                dst_ptr + v, gmem_ptr, iterator_A.valid());
+          }
+
+          ++iterator_A;
+        }
+
+        ++this->smem_iterator_A_;
+      }
     }
 
-    CUTLASS_DEVICE
-    void
-    copy_tiles_and_advance(IteratorA& iterator_A, IteratorB& iterator_B, int group_start_A = 0, int group_start_B = 0)
-    {
-        iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector);
-        this->smem_iterator_A_.set_iteration_index(group_start_A);
+    iterator_B.set_iteration_index(group_start_B *
+                                   IteratorB::kAccessesPerVector);
+    this->smem_iterator_B_.set_iteration_index(group_start_B);
+
+    // Async Copy for operand B
+    CUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) {
+      if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                this->smem_iterator_B_.get());
+
+        int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value *
+                              IteratorB::ThreadMap::kElementsPerAccess /
+                              IteratorB::kAccessesPerVector / 8;
 
-        // Async Copy for operand A
         CUTLASS_PRAGMA_UNROLL
-        for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) {
-            if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) {
-                typename IteratorA::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorA::AccessType*>(this->smem_iterator_A_.get());
+        for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
+          auto gmem_ptr = iterator_B.get();
 
-                int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value
-                                      * IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector / 8;
+          if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
+                dst_ptr + v, gmem_ptr, iterator_B.valid());
+          } else {
+            cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(
+                dst_ptr + v, gmem_ptr, iterator_B.valid());
+          }
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
-                    auto gmem_ptr = iterator_A.get();
+          // if(true && (threadIdx.x||threadIdx.y||threadIdx.z)==0){
+          //     int32_t* print_ptr =
+          //     reinterpret_cast<int32_t*>(iterator_B.get()); int32_t*
+          //     print_ptr_smem = reinterpret_cast<int32_t*>(dst_ptr+v); if
+          //     (iterator_B.valid())
+          //     {
+          //         printf("gmem_ptr cp source of thread %d-%d-%d;%d-%d-%d:
+          //         %p:%x-%x-%x-%x=>%x,%x,%x,%x \n",
+          //             blockIdx.x,blockIdx.y,blockIdx.z,
+          //             threadIdx.x,threadIdx.y,threadIdx.z,
+          //             iterator_B.get(),
+          //             static_cast<uint32_t>(print_ptr[0]),
+          //             static_cast<uint32_t>(print_ptr[1]),
+          //             static_cast<uint32_t>(print_ptr[2]),
+          //             static_cast<uint32_t>(print_ptr[3]),
+          //             static_cast<uint32_t>(print_ptr_smem[0]),
+          //             static_cast<uint32_t>(print_ptr_smem[1]),
+          //             static_cast<uint32_t>(print_ptr_smem[2]),
+          //             static_cast<uint32_t>(print_ptr_smem[3]));
+          //     }
+          // }
+          ++iterator_B;
+        }
+        ++this->smem_iterator_B_;
+      }
+    }
+  }
 
-                    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
-                        cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(dst_ptr + v, gmem_ptr, iterator_A.valid());
-                    }
-                    else {
-                        cutlass::arch::cp_async<kSrcBytes, kCacheOpA>(dst_ptr + v, gmem_ptr, iterator_A.valid());
-                    }
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  CUTLASS_DEVICE
+  void operator()(
+      ///< problem size of GEMM
+      int gemm_k_iterations,
+      ///< destination accumulator tile
+      FragmentC& accum,
+      ///< iterator over A operand in global memory
+      IteratorA iterator_A,
+      ///< iterator over B operand in global memory
+      IteratorB iterator_B,
+      IteratorNF4LookUpTable iterator_nf4_look_up_table,
+      ///< initial value of accumulator
+      FragmentC const& src_accum) {
+    // printf("gemm_k_iterations:%d\n", gemm_k_iterations);
 
-                    ++iterator_A;
-                }
+    //
+    // Prologue
+    //
 
-                ++this->smem_iterator_A_;
-            }
+    // use share memory to get look_up_table of nf4;
+
+    // __shared__ uint32_t shared_look_up_table[16];
+
+    // int lane_idx=threadIdx.x%32;
+    // int warp_idx=threadIdx.x/32;
+    // if(lane_idx<16){
+    //     shared_look_up_table[lane_idx]=lane_idx;
+    // }
+
+    // __shared__ uint32_t shared_look_up_table[32][32];
+    // if(warp_idx==0){
+    //     CUTLASS_PRAGMA_UNROLL
+    //     for(int ii=0;ii<16;++ii){
+    //         shared_look_up_table[lane_idx][ii]=ii;
+    //     }
+    // }
+
+    /// load look up table to smem here
+    // __shared__ int32_t nf4_smem_look_up_table[16];
+
+    // int32_t* gmem_ptr_nf4_look_up_table =
+    // reinterpret_cast<int32_t*>(iterator_nf4_look_up_table.get());
+    // // smem look up table
+    // int32_t* dst_ptr_nf4_look_up_table =
+    // reinterpret_cast<int32_t*>(nf4_smem_look_up_table);
+
+    // if(lane_idx == 0){
+    //     int4* dst_ptr_nf4_look_up_table_int4 =
+    //     reinterpret_cast<int4*>(nf4_smem_look_up_table);
+    //     dst_ptr_nf4_look_up_table_int4[lane_idx] =
+    //     *(reinterpret_cast<int4*>(gmem_ptr_nf4_look_up_table) + lane_idx);
+    // }
+    // __syncthreads();
+    // // reg look up table
+    // cutlass::Array<uint32_t, 4, 1> reg_look_up_table;
+    // CUTLASS_PRAGMA_UNROLL
+    // for(int ii=0;ii<4;++ii){
+    //     reg_look_up_table[ii]=*(reinterpret_cast<int32_t*>(dst_ptr_nf4_look_up_table)
+    //     + ii);
+    // }
+
+    TransformBAfterLDS lds_converter;
+
+    // NOTE - switch to ldg.sts
+    // Issue this first, so cp.async.commit_group will commit this load as well.
+    // Note: we do not commit here and this load will commit in the same group
+    // as
+    //       the first load of A.
+
+    // Issue several complete stages
+    CUTLASS_PRAGMA_UNROLL
+    for (int stage = 0; stage < Base::kStages - 1;
+         ++stage, --gemm_k_iterations) {
+      iterator_A.clear_mask(gemm_k_iterations == 0);
+      iterator_B.clear_mask(gemm_k_iterations == 0);
+
+      iterator_A.set_iteration_index(0);
+      this->smem_iterator_A_.set_iteration_index(0);
+
+      // Async Copy for operand A
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                this->smem_iterator_A_.get());
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
+          int const kSrcBytes =
+              sizeof_bits<typename IteratorA::Element>::value *
+              IteratorA::ThreadMap::kElementsPerAccess /
+              IteratorA::kAccessesPerVector / 8;
+
+          int src_bytes = (iterator_A.valid() ? kSrcBytes : 0);
+
+          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
+              dst_ptr + v, iterator_A.get(), iterator_A.valid());
+
+          ++iterator_A;
         }
 
-        iterator_B.set_iteration_index(group_start_B * IteratorB::kAccessesPerVector);
-        this->smem_iterator_B_.set_iteration_index(group_start_B);
+        ++this->smem_iterator_A_;
+      }
+
+      iterator_B.set_iteration_index(0);
+      // print_type<IteratorB>();
+      this->smem_iterator_B_.set_iteration_index(0);
+
+      // Async Copy for operand B
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                this->smem_iterator_B_.get());
 
-        // Async Copy for operand B
         CUTLASS_PRAGMA_UNROLL
-        for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) {
-            if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) {
-                typename IteratorB::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorB::AccessType*>(this->smem_iterator_B_.get());
+        for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
+          int const kSrcBytes =
+              sizeof_bits<typename IteratorB::Element>::value *
+              IteratorB::ThreadMap::kElementsPerAccess /
+              IteratorB::kAccessesPerVector / 8;
 
-                int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value
-                                      * IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector / 8;
+          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
+              dst_ptr + v, iterator_B.get(), iterator_B.valid());
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
-                    auto gmem_ptr = iterator_B.get();
-
-                    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
-                        cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(dst_ptr + v, gmem_ptr, iterator_B.valid());
-                    }
-                    else {
-                        cutlass::arch::cp_async<kSrcBytes, kCacheOpB>(dst_ptr + v, gmem_ptr, iterator_B.valid());
-                    }
-
-                    // if(true && (threadIdx.x||threadIdx.y||threadIdx.z)==0){
-                    //     int32_t* print_ptr = reinterpret_cast<int32_t*>(iterator_B.get());
-                    //     int32_t* print_ptr_smem = reinterpret_cast<int32_t*>(dst_ptr+v);
-                    //     if (iterator_B.valid())
-                    //     {
-                    //         printf("gmem_ptr cp source of thread %d-%d-%d;%d-%d-%d: %p:%x-%x-%x-%x=>%x,%x,%x,%x \n",
-                    //             blockIdx.x,blockIdx.y,blockIdx.z,
-                    //             threadIdx.x,threadIdx.y,threadIdx.z,
-                    //             iterator_B.get(),
-                    //             static_cast<uint32_t>(print_ptr[0]),
-                    //             static_cast<uint32_t>(print_ptr[1]),
-                    //             static_cast<uint32_t>(print_ptr[2]),
-                    //             static_cast<uint32_t>(print_ptr[3]),
-                    //             static_cast<uint32_t>(print_ptr_smem[0]),
-                    //             static_cast<uint32_t>(print_ptr_smem[1]),
-                    //             static_cast<uint32_t>(print_ptr_smem[2]),
-                    //             static_cast<uint32_t>(print_ptr_smem[3]));
-                    //     }
-                    // }
-                    ++iterator_B;
-                }
-                ++this->smem_iterator_B_;
-            }
+          // if(true && (threadIdx.x||threadIdx.y||threadIdx.z)==0){
+          //     int32_t* print_ptr =
+          //     reinterpret_cast<int32_t*>(iterator_B.get()); int32_t*
+          //     print_ptr_smem = reinterpret_cast<int32_t*>(dst_ptr+v); if
+          //     (iterator_B.valid())
+          //     {
+          //         printf("gmem_ptr cp source of thread %d-%d-%d;%d-%d-%d:
+          //         %p:%x-%x-%x-%x=>%x,%x,%x,%x \n",
+          //             blockIdx.x,blockIdx.y,blockIdx.z,
+          //             threadIdx.x,threadIdx.y,threadIdx.z,
+          //             iterator_B.get(),
+          //             static_cast<uint32_t>(print_ptr[0]),
+          //             static_cast<uint32_t>(print_ptr[1]),
+          //             static_cast<uint32_t>(print_ptr[2]),
+          //             static_cast<uint32_t>(print_ptr[3]),
+          //             static_cast<uint32_t>(print_ptr_smem[0]),
+          //             static_cast<uint32_t>(print_ptr_smem[1]),
+          //             static_cast<uint32_t>(print_ptr_smem[2]),
+          //             static_cast<uint32_t>(print_ptr_smem[3]));
+          //     }
+          // }
+          ++iterator_B;
         }
+
+        ++this->smem_iterator_B_;
+      }
+
+      // Move to the next stage
+      iterator_A.add_tile_offset({0, 1});
+      iterator_B.add_tile_offset({1, 0});
+
+      this->smem_iterator_A_.add_tile_offset({0, 1});
+      this->smem_iterator_B_.add_tile_offset({1, 0});
+
+      // Defines the boundary of a stage of cp.async.
+      cutlass::arch::cp_async_fence();
     }
 
-    /// Perform a threadblock-scoped matrix multiply-accumulate
-    CUTLASS_DEVICE
-    void operator()(
-        ///< problem size of GEMM
-        int gemm_k_iterations,
-        ///< destination accumulator tile
-        FragmentC& accum,
-        ///< iterator over A operand in global memory
-        IteratorA iterator_A,
-        ///< iterator over B operand in global memory
-        IteratorB iterator_B,
-        IteratorNF4LookUpTable iterator_nf4_look_up_table,
-        ///< initial value of accumulator
-        FragmentC const& src_accum)
-    {
+    // Perform accumulation in the 'd' output operand
+    accum = src_accum;
 
-        // printf("gemm_k_iterations:%d\n", gemm_k_iterations);
+    //
+    // Clear the remaining tiles of SMEM. This is a functional requirement for
+    // some kernels so that all accumulator elements outside the GEMM footprint
+    // are zero.
+    //
 
-        //
-        // Prologue
-        //
+    if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) {
+      /// Iterator to write threadblock-scoped tile of A operand to shared
+      /// memory
+      SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
 
-        // use share memory to get look_up_table of nf4;
+      typename IteratorA::AccessType zero_A;
+      zero_A.clear();
 
-        // __shared__ uint32_t shared_look_up_table[16];
+      last_smem_iterator_A.set_iteration_index(0);
 
-        // int lane_idx=threadIdx.x%32;
-        // int warp_idx=threadIdx.x/32;
-        // if(lane_idx<16){
-        //     shared_look_up_table[lane_idx]=lane_idx;
-        // }
+      // Async Copy for operand A
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
+        typename IteratorA::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorA::AccessType*>(
+                last_smem_iterator_A.get());
 
-        // __shared__ uint32_t shared_look_up_table[32][32];
-        // if(warp_idx==0){
-        //     CUTLASS_PRAGMA_UNROLL
-        //     for(int ii=0;ii<16;++ii){
-        //         shared_look_up_table[lane_idx][ii]=ii;
-        //     }
-        // }
+        *dst_ptr = zero_A;
 
-        /// load look up table to smem here
-        // __shared__ int32_t nf4_smem_look_up_table[16];
+        ++last_smem_iterator_A;
+      }
 
-        // int32_t* gmem_ptr_nf4_look_up_table = reinterpret_cast<int32_t*>(iterator_nf4_look_up_table.get());
-        // // smem look up table
-        // int32_t* dst_ptr_nf4_look_up_table = reinterpret_cast<int32_t*>(nf4_smem_look_up_table);
+      /// Iterator to write threadblock-scoped tile of B operand to shared
+      /// memory
+      SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_);
+      typename IteratorB::AccessType zero_B;
 
-        // if(lane_idx == 0){
-        //     int4* dst_ptr_nf4_look_up_table_int4 = reinterpret_cast<int4*>(nf4_smem_look_up_table);
-        //     dst_ptr_nf4_look_up_table_int4[lane_idx] = *(reinterpret_cast<int4*>(gmem_ptr_nf4_look_up_table) + lane_idx);
-        // }
-        // __syncthreads();
-        // // reg look up table
-        // cutlass::Array<uint32_t, 4, 1> reg_look_up_table;
-        // CUTLASS_PRAGMA_UNROLL
-        // for(int ii=0;ii<4;++ii){
-        //     reg_look_up_table[ii]=*(reinterpret_cast<int32_t*>(dst_ptr_nf4_look_up_table) + ii);
-        // }
+      zero_B.clear();
+      last_smem_iterator_B.set_iteration_index(0);
 
-        TransformBAfterLDS lds_converter;
+      // Async Copy for operand B
+      CUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
+        typename IteratorB::AccessType* dst_ptr =
+            reinterpret_cast<typename IteratorB::AccessType*>(
+                last_smem_iterator_B.get());
 
-        // NOTE - switch to ldg.sts
-        // Issue this first, so cp.async.commit_group will commit this load as well.
-        // Note: we do not commit here and this load will commit in the same group as
-        //       the first load of A.
+        *dst_ptr = zero_B;
 
-        // Issue several complete stages
-        CUTLASS_PRAGMA_UNROLL
-        for (int stage = 0; stage < Base::kStages - 1; ++stage, --gemm_k_iterations) {
+        ++last_smem_iterator_B;
+      }
+    }
 
-            iterator_A.clear_mask(gemm_k_iterations == 0);
-            iterator_B.clear_mask(gemm_k_iterations == 0);
+    // Waits until kStages-2 stages have committed.
+    cutlass::arch::cp_async_wait<Base::kStages - 2>();
+    __syncthreads();
 
-            iterator_A.set_iteration_index(0);
-            this->smem_iterator_A_.set_iteration_index(0);
+    // Pair of fragments used to overlap shared memory loads and math
+    // instructions
+    WarpFragmentA warp_frag_A[2];
+    WarpFragmentB warp_frag_B[2];
+    typename TransformBAfterLDS::result_type converted_frag_B_buffer[2];
+    Operator warp_mma;
 
-            // Async Copy for operand A
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
-                typename IteratorA::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorA::AccessType*>(this->smem_iterator_A_.get());
+    this->warp_tile_iterator_A_.set_kgroup_index(0);
+    this->warp_tile_iterator_B_.set_kgroup_index(0);
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
-                    int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value
-                                          * IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector
-                                          / 8;
+    this->warp_tile_iterator_B_.load(warp_frag_B[0]);
+    converted_frag_B_buffer[0] = lds_converter(warp_frag_B[0]);
+    this->warp_tile_iterator_A_.load(warp_frag_A[0]);
 
-                    int src_bytes = (iterator_A.valid() ? kSrcBytes : 0);
+    // if((threadIdx.x||threadIdx.y||threadIdx.z)==0){
+    //     uint32_t* frag_b_reg_ptr =
+    //     reinterpret_cast<uint32_t*>(&warp_frag_B[0]); printf("####
+    //     warp_frag_b_load [0] bid:%d-%d-%d,"
+    //                 " frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
+    //             blockIdx.x,blockIdx.y,blockIdx.z,
+    //             frag_b_reg_ptr[0],
+    //             frag_b_reg_ptr[1],
+    //             frag_b_reg_ptr[2],
+    //             frag_b_reg_ptr[3],
+    //             frag_b_reg_ptr[4],
+    //             frag_b_reg_ptr[5],
+    //             frag_b_reg_ptr[6],
+    //             frag_b_reg_ptr[7]
+    //             );
+    // }
+    ++this->warp_tile_iterator_A_;
+    ++this->warp_tile_iterator_B_;
 
-                    cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
-                        dst_ptr + v, iterator_A.get(), iterator_A.valid());
+    iterator_A.clear_mask(gemm_k_iterations == 0);
+    iterator_B.clear_mask(gemm_k_iterations == 0);
 
-                    ++iterator_A;
-                }
+    int smem_write_stage_idx = Base::kStages - 1;
+    int smem_read_stage_idx = 0;
 
-                ++this->smem_iterator_A_;
-            }
+    //
+    // Mainloop
+    //
 
-            iterator_B.set_iteration_index(0);
-            // print_type<IteratorB>();
-            this->smem_iterator_B_.set_iteration_index(0);
+    __syncthreads();
+    CUTLASS_GEMM_LOOP
+    for (; gemm_k_iterations > (-Base::kStages + 1);) {
+      //
+      // Loop over GEMM K dimension
+      //
 
-            // Async Copy for operand B
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
-                typename IteratorB::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorB::AccessType*>(this->smem_iterator_B_.get());
+      // Computes a warp-level GEMM on data held in shared memory
+      // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate
+      CUTLASS_PRAGMA_UNROLL
+      for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
+           ++warp_mma_k) {
+        // Load warp-level tiles from shared memory, wrapping to k offset if
+        // this is the last group as the case may be.
+        this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) %
+                                                     Base::kWarpGemmIterations);
+        this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
+        ++this->warp_tile_iterator_A_;
+        // static_assert(Base::kNumKIterationsPerWarpBLoad==1,"Base::kNumKIterationsPerWarpBLoad!=1");
+        // static_assert(Base::kWarpGemmIterationsForB==4,"Base::kWarpGemmIterationsForB!=4");
+        const int warp_tileB_k_compute_offset =
+            warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
+        const int warp_tileB_k_load_offset =
+            warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
+        if (warp_tileB_k_compute_offset ==
+            Base::kNumKIterationsPerWarpBLoad - 1) {
+          this->warp_tile_iterator_B_.set_kgroup_index(
+              (warp_tileB_k_load_offset + 1) % Base::kWarpGemmIterationsForB);
+          this->warp_tile_iterator_B_.load(
+              warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
+          converted_frag_B_buffer[(warp_tileB_k_load_offset + 1) % 2] =
+              lds_converter(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
 
-                CUTLASS_PRAGMA_UNROLL
-                for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
-                    int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value
-                                          * IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector
-                                          / 8;
-
-                    cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
-                        dst_ptr + v, iterator_B.get(), iterator_B.valid());
-
-                    // if(true && (threadIdx.x||threadIdx.y||threadIdx.z)==0){
-                    //     int32_t* print_ptr = reinterpret_cast<int32_t*>(iterator_B.get());
-                    //     int32_t* print_ptr_smem = reinterpret_cast<int32_t*>(dst_ptr+v);
-                    //     if (iterator_B.valid())
-                    //     {
-                    //         printf("gmem_ptr cp source of thread %d-%d-%d;%d-%d-%d: %p:%x-%x-%x-%x=>%x,%x,%x,%x \n",
-                    //             blockIdx.x,blockIdx.y,blockIdx.z,
-                    //             threadIdx.x,threadIdx.y,threadIdx.z,
-                    //             iterator_B.get(),
-                    //             static_cast<uint32_t>(print_ptr[0]),
-                    //             static_cast<uint32_t>(print_ptr[1]),
-                    //             static_cast<uint32_t>(print_ptr[2]),
-                    //             static_cast<uint32_t>(print_ptr[3]),
-                    //             static_cast<uint32_t>(print_ptr_smem[0]),
-                    //             static_cast<uint32_t>(print_ptr_smem[1]),
-                    //             static_cast<uint32_t>(print_ptr_smem[2]),
-                    //             static_cast<uint32_t>(print_ptr_smem[3]));
-                    //     }
-                    // }
-                    ++iterator_B;
-                }
-
-                ++this->smem_iterator_B_;
-            }
-
-            // Move to the next stage
-            iterator_A.add_tile_offset({0, 1});
-            iterator_B.add_tile_offset({1, 0});
-
-            this->smem_iterator_A_.add_tile_offset({0, 1});
-            this->smem_iterator_B_.add_tile_offset({1, 0});
-
-            // Defines the boundary of a stage of cp.async.
-            cutlass::arch::cp_async_fence();
+          ++this->warp_tile_iterator_B_;
+          // if((threadIdx.x||threadIdx.y||threadIdx.z)==0){
+          //     uint32_t* frag_b_reg_ptr =
+          //     reinterpret_cast<uint32_t*>(&warp_frag_B[(warp_tileB_k_load_offset
+          //     + 1) % 2]); printf("#### warp_frag_b load [%d] bid:%d-%d-%d,"
+          //                " frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
+          //             ((warp_tileB_k_load_offset + 1) % 2),
+          //             blockIdx.x,blockIdx.y,blockIdx.z,
+          //             frag_b_reg_ptr[0],
+          //             frag_b_reg_ptr[1],
+          //             frag_b_reg_ptr[2],
+          //             frag_b_reg_ptr[3],
+          //             frag_b_reg_ptr[4],
+          //             frag_b_reg_ptr[5],
+          //             frag_b_reg_ptr[6],
+          //             frag_b_reg_ptr[7]
+          //             );
+          // }
         }
+        // TODO(wangbojun) lds_converter can be remove for int8 B input
+        // int4
+        // typename TransformBAfterLDS::result_type converted_frag_B =
+        //     lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
 
-        // Perform accumulation in the 'd' output operand
-        accum = src_accum;
+        // typename TransformBAfterLDS::result_type converted_frag_B =
+        //     lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2],
+        //     reinterpret_cast<int32_t*>(nf4_smem_look_up_table));
 
-        //
-        // Clear the remaining tiles of SMEM. This is a functional requirement for some kernels
-        // so that all accumulator elements outside the GEMM footprint are zero.
-        //
+        // typename TransformBAfterLDS::result_type converted_frag_B =
+        //     lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2],
+        //     reg_look_up_table);
 
-        if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) {
-
-            /// Iterator to write threadblock-scoped tile of A operand to shared memory
-            SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_);
-
-            typename IteratorA::AccessType zero_A;
-            zero_A.clear();
-
-            last_smem_iterator_A.set_iteration_index(0);
-
-            // Async Copy for operand A
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
-
-                typename IteratorA::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorA::AccessType*>(last_smem_iterator_A.get());
-
-                *dst_ptr = zero_A;
-
-                ++last_smem_iterator_A;
-            }
-
-            /// Iterator to write threadblock-scoped tile of B operand to shared memory
-            SmemIteratorB                  last_smem_iterator_B(this->smem_iterator_B_);
-            typename IteratorB::AccessType zero_B;
-
-            zero_B.clear();
-            last_smem_iterator_B.set_iteration_index(0);
-
-            // Async Copy for operand B
-            CUTLASS_PRAGMA_UNROLL
-            for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
-
-                typename IteratorB::AccessType* dst_ptr =
-                    reinterpret_cast<typename IteratorB::AccessType*>(last_smem_iterator_B.get());
-
-                *dst_ptr = zero_B;
-
-                ++last_smem_iterator_B;
-            }
-        }
-
-        // Waits until kStages-2 stages have committed.
-        cutlass::arch::cp_async_wait<Base::kStages - 2>();
-        __syncthreads();
-
-        // Pair of fragments used to overlap shared memory loads and math
-        // instructions
-        WarpFragmentA                       warp_frag_A[2];
-        WarpFragmentB                       warp_frag_B[2];
-        typename TransformBAfterLDS::result_type converted_frag_B_buffer[2];
-        Operator warp_mma;
-
-        this->warp_tile_iterator_A_.set_kgroup_index(0);
-        this->warp_tile_iterator_B_.set_kgroup_index(0);
-
-        this->warp_tile_iterator_B_.load(warp_frag_B[0]);
-        converted_frag_B_buffer[0] =
-            lds_converter(warp_frag_B[0]);
-        this->warp_tile_iterator_A_.load(warp_frag_A[0]);
+        // typename TransformBAfterLDS::result_type converted_frag_B =
+        //     lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2],
+        //     shared_look_up_table, warp_idx, lane_idx);
 
         // if((threadIdx.x||threadIdx.y||threadIdx.z)==0){
-        //     uint32_t* frag_b_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_B[0]);
-        //     printf("#### warp_frag_b_load [0] bid:%d-%d-%d,"
-        //                 " frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
+        //     uint32_t* frag_b_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&warp_frag_B[(warp_tileB_k_load_offset)
+        //     % 2]); uint32_t* converted_frag_B_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&converted_frag_B); printf("####
+        //     after lds_converter bid:%d-%d-%d"
+        //             " frag_b_reg_ptr[%d]:%x-%x-%x-%x-%x-%x-%x-%x"
+        //             " converted_frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
         //             blockIdx.x,blockIdx.y,blockIdx.z,
+        //             ((warp_tileB_k_load_offset) % 2),
         //             frag_b_reg_ptr[0],
         //             frag_b_reg_ptr[1],
         //             frag_b_reg_ptr[2],
@@ -584,308 +719,243 @@ public:
         //             frag_b_reg_ptr[4],
         //             frag_b_reg_ptr[5],
         //             frag_b_reg_ptr[6],
-        //             frag_b_reg_ptr[7]
-        //             );
+        //             frag_b_reg_ptr[7],
+        //             converted_frag_B_reg_ptr[0],
+        //             converted_frag_B_reg_ptr[1],
+        //             converted_frag_B_reg_ptr[2],
+        //             converted_frag_B_reg_ptr[3],
+        //             converted_frag_B_reg_ptr[4],
+        //             converted_frag_B_reg_ptr[5],
+        //             converted_frag_B_reg_ptr[6],
+        //             converted_frag_B_reg_ptr[7]
+        //     );
         // }
-        ++this->warp_tile_iterator_A_;
-        ++this->warp_tile_iterator_B_;
 
-        iterator_A.clear_mask(gemm_k_iterations == 0);
-        iterator_B.clear_mask(gemm_k_iterations == 0);
+        // bool  ::print_type< ::cutlass::Array<
+        // ::cutlass::integer_subbyte<(int)4, (bool)0> , (int)64, (bool)0> >
+        // ()") print_type<WarpFragmentB>(); bool  ::print_type<
+        // ::cutlass::Array<signed char, (int)64, (bool)0> > ()") from a
+        // print_type<TransformBAfterLDS::result_type>();
+        // cutlass::Array<signed char, (int)32, (bool)0>
+        // print_type<WarpFragmentA>();
 
-        int smem_write_stage_idx = Base::kStages - 1;
-        int smem_read_stage_idx  = 0;
+        // print_type<FragmentC>();
+        // TODO(zhengzekang)
+        // run_warp_mma(
 
-        //
-        // Mainloop
-        //
+        // if(true){
+        //     uint32_t none_zero = 0;
+        //     // uint32_t* converted_frag_B_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&converted_frag_B); uint32_t*
+        //     converted_frag_B_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&warp_frag_B[warp_mma_k % 2]);
+        //     uint32_t* frag_a_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&warp_frag_A[warp_mma_k % 2]);
+        //     CUTLASS_PRAGMA_UNROLL
+        //     for(int ii=0;ii<warp_frag_B[warp_mma_k % 2].size()/8;++ii){
+        //         none_zero|=*(converted_frag_B_reg_ptr+ii);
+        //     }
+        //     CUTLASS_PRAGMA_UNROLL
+        //     for(int ii=0;ii<warp_frag_A[warp_mma_k % 2].size()/4;++ii){
+        //         none_zero|=*(frag_a_reg_ptr+ii);
+        //     }
 
-        __syncthreads();
-        CUTLASS_GEMM_LOOP
-        for (; gemm_k_iterations > (-Base::kStages + 1);) {
-            //
-            // Loop over GEMM K dimension
-            //
+        //     CUTLASS_PRAGMA_UNROLL
+        //     for(int none_zero_i = 16; none_zero_i>0;none_zero_i/=2){
+        //         none_zero|= __shfl_xor_sync(-1,none_zero,none_zero_i);
+        //     }
 
-            // Computes a warp-level GEMM on data held in shared memory
-            // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate
-            CUTLASS_PRAGMA_UNROLL
-            for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) {
+        //     if(none_zero!=0){
+        //         printf("## before mma ## bidtid:%d-%d-%d-%d-%d-%d,
+        //         warp_mma_k:%d, frag_B_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x;
+        //         frag_a_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x"
+        //                 " accu:
+        //                 %d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d
+        //                 \n", blockIdx.x,blockIdx.y,blockIdx.z, warp_mma_k,
+        //                 threadIdx.x,threadIdx.y,threadIdx.z,
+        //                 converted_frag_B_reg_ptr[0],
+        //                 converted_frag_B_reg_ptr[1],
+        //                 converted_frag_B_reg_ptr[2],
+        //                 converted_frag_B_reg_ptr[3],
+        //                 converted_frag_B_reg_ptr[4],
+        //                 converted_frag_B_reg_ptr[5],
+        //                 converted_frag_B_reg_ptr[6],
+        //                 converted_frag_B_reg_ptr[7],
+        //                 frag_a_reg_ptr[0],
+        //                 frag_a_reg_ptr[1],
+        //                 frag_a_reg_ptr[2],
+        //                 frag_a_reg_ptr[3],
+        //                 frag_a_reg_ptr[4],
+        //                 frag_a_reg_ptr[5],
+        //                 frag_a_reg_ptr[6],
+        //                 frag_a_reg_ptr[7],
+        //                 accum[0],
+        //                 accum[1],
+        //                 accum[2],
+        //                 accum[3],
+        //                 accum[4],
+        //                 accum[5],
+        //                 accum[6],
+        //                 accum[7],
+        //                 accum[8],
+        //                 accum[9],
+        //                 accum[10],
+        //                 accum[11],
+        //                 accum[12],
+        //                 accum[13],
+        //                 accum[14],
+        //                 accum[15],
+        //                 accum[16],
+        //                 accum[17],
+        //                 accum[18],
+        //                 accum[19],
+        //                 accum[20],
+        //                 accum[21],
+        //                 accum[22],
+        //                 accum[23],
+        //                 accum[24],
+        //                 accum[25],
+        //                 accum[26],
+        //                 accum[27],
+        //                 accum[28],
+        //                 accum[29],
+        //                 accum[30],
+        //                 accum[31]
+        //                );
+        //     }
+        // }
+        run_ampere_warp_mma(
+            warp_mma,
+            accum,
+            warp_frag_A[warp_mma_k % 2],
+            converted_frag_B_buffer[warp_tileB_k_load_offset % 2],
+            accum,
+            warp_tileB_k_compute_offset);
+        // auto tmp = static_cast<int32_t>(warp_frag_B[warp_tileB_k_load_offset
+        // % 2]); if(threadIdx.x==0  && threadIdx.y==0  && threadIdx.z==0 &&
+        //     blockIdx.x==0  && blockIdx.y==0  && blockIdx.z==0){
+        //     printf("### run_warp_mma: "
+        //         "%d \n",
+        //         reinterpret_cast<int32_t &>(accum));
+        // }
+        // if(true){
+        //     uint32_t none_zero = 0;
+        //     uint32_t* converted_frag_B_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&converted_frag_B);
+        //     // uint32_t* converted_frag_B_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&warp_frag_B[warp_mma_k % 2]);
+        //     uint32_t* frag_a_reg_ptr =
+        //     reinterpret_cast<uint32_t*>(&warp_frag_A[warp_mma_k % 2]);
+        //     CUTLASS_PRAGMA_UNROLL
+        //     for(int ii=0;ii<warp_frag_B[warp_mma_k % 2].size()/8;++ii){
+        //         none_zero|=*(converted_frag_B_reg_ptr+ii);
+        //     }
+        //     CUTLASS_PRAGMA_UNROLL
+        //     for(int ii=0;ii<warp_frag_A[warp_mma_k % 2].size()/4;++ii){
+        //         none_zero|=*(frag_a_reg_ptr+ii);
+        //     }
+        //     CUTLASS_PRAGMA_UNROLL
+        //     for(int none_zero_i = 16; none_zero_i>0;none_zero_i/=2){
+        //         none_zero|= __shfl_xor_sync(-1,none_zero,none_zero_i);
+        //     }
 
-                // Load warp-level tiles from shared memory, wrapping to k offset if
-                // this is the last group as the case may be.
-                this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
-                this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]);
-                ++this->warp_tile_iterator_A_;
-                // static_assert(Base::kNumKIterationsPerWarpBLoad==1,"Base::kNumKIterationsPerWarpBLoad!=1");
-                // static_assert(Base::kWarpGemmIterationsForB==4,"Base::kWarpGemmIterationsForB!=4");
-                const int warp_tileB_k_compute_offset = warp_mma_k % Base::kNumKIterationsPerWarpBLoad;
-                const int warp_tileB_k_load_offset    = warp_mma_k / Base::kNumKIterationsPerWarpBLoad;
-                if (warp_tileB_k_compute_offset == Base::kNumKIterationsPerWarpBLoad - 1) {
-                    this->warp_tile_iterator_B_.set_kgroup_index((warp_tileB_k_load_offset + 1)
-                                                                 % Base::kWarpGemmIterationsForB);
-                    this->warp_tile_iterator_B_.load(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
-                    converted_frag_B_buffer[(warp_tileB_k_load_offset + 1) % 2] =
-                        lds_converter(warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
+        //     // if(none_zero!=0){
+        //     if((blockIdx.y||blockIdx.z||threadIdx.x||threadIdx.y||threadIdx.z)==0){
 
-                    ++this->warp_tile_iterator_B_;
-                    // if((threadIdx.x||threadIdx.y||threadIdx.z)==0){
-                    //     uint32_t* frag_b_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_B[(warp_tileB_k_load_offset + 1) % 2]);
-                    //     printf("#### warp_frag_b load [%d] bid:%d-%d-%d,"
-                    //                " frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
-                    //             ((warp_tileB_k_load_offset + 1) % 2),
-                    //             blockIdx.x,blockIdx.y,blockIdx.z,
-                    //             frag_b_reg_ptr[0],
-                    //             frag_b_reg_ptr[1],
-                    //             frag_b_reg_ptr[2],
-                    //             frag_b_reg_ptr[3],
-                    //             frag_b_reg_ptr[4],
-                    //             frag_b_reg_ptr[5],
-                    //             frag_b_reg_ptr[6],
-                    //             frag_b_reg_ptr[7]
-                    //             );
-                    // }
-                }
-                // TODO(wangbojun) lds_converter can be remove for int8 B input
-                // int4
-                // typename TransformBAfterLDS::result_type converted_frag_B =
-                //     lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2]);
+        //         printf("## after mma ## bidtid:%d-%d-%d-%d-%d-%d,
+        //         warp_mma_k:%d, gemm_k_iterations:%d,
+        //         Base::kWarpGemmIterations:%d,"
+        //                 " converted_frag_B_reg_ptr:%x; frag_a_reg_ptr:%x"
+        //                 " accu: %d \n",
+        //                 blockIdx.x,blockIdx.y,blockIdx.z,
+        //                 threadIdx.x,threadIdx.y,threadIdx.z,
+        //                 warp_mma_k,
+        //                 gemm_k_iterations,
+        //                 Base::kWarpGemmIterations,
+        //                 converted_frag_B_reg_ptr[0],
+        //                 frag_a_reg_ptr[0],
+        //                 accum[0]
+        //                );
+        //     }
+        // }
+        // Issue global->shared copies for the this stage
+        if (warp_mma_k < Base::kWarpGemmIterations - 1) {
+          int group_start_iteration_A, group_start_iteration_B;
 
-                // typename TransformBAfterLDS::result_type converted_frag_B =
-                //     lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2], reinterpret_cast<int32_t*>(nf4_smem_look_up_table));
+          group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
+          group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
 
-                // typename TransformBAfterLDS::result_type converted_frag_B =
-                //     lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2], reg_look_up_table);
-
-                // typename TransformBAfterLDS::result_type converted_frag_B =
-                //     lds_converter(warp_frag_B[warp_tileB_k_load_offset % 2], shared_look_up_table, warp_idx, lane_idx);
-
-                // if((threadIdx.x||threadIdx.y||threadIdx.z)==0){
-                //     uint32_t* frag_b_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_B[(warp_tileB_k_load_offset) % 2]);
-                //     uint32_t* converted_frag_B_reg_ptr = reinterpret_cast<uint32_t*>(&converted_frag_B);
-                //     printf("#### after lds_converter bid:%d-%d-%d"
-                //             " frag_b_reg_ptr[%d]:%x-%x-%x-%x-%x-%x-%x-%x"
-                //             " converted_frag_b_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x  \n",
-                //             blockIdx.x,blockIdx.y,blockIdx.z,
-                //             ((warp_tileB_k_load_offset) % 2),
-                //             frag_b_reg_ptr[0],
-                //             frag_b_reg_ptr[1],
-                //             frag_b_reg_ptr[2],
-                //             frag_b_reg_ptr[3],
-                //             frag_b_reg_ptr[4],
-                //             frag_b_reg_ptr[5],
-                //             frag_b_reg_ptr[6],
-                //             frag_b_reg_ptr[7],
-                //             converted_frag_B_reg_ptr[0],
-                //             converted_frag_B_reg_ptr[1],
-                //             converted_frag_B_reg_ptr[2],
-                //             converted_frag_B_reg_ptr[3],
-                //             converted_frag_B_reg_ptr[4],
-                //             converted_frag_B_reg_ptr[5],
-                //             converted_frag_B_reg_ptr[6],
-                //             converted_frag_B_reg_ptr[7]
-                //     );
-                // }
-
-                // bool  ::print_type< ::cutlass::Array< ::cutlass::integer_subbyte<(int)4, (bool)0> , (int)64, (bool)0> > ()")
-                // print_type<WarpFragmentB>();
-                // bool  ::print_type< ::cutlass::Array<signed char, (int)64, (bool)0> > ()") from a
-                // print_type<TransformBAfterLDS::result_type>();
-                // cutlass::Array<signed char, (int)32, (bool)0>
-                // print_type<WarpFragmentA>();
-
-                // print_type<FragmentC>();
-                // TODO(zhengzekang)
-                // run_warp_mma(
-
-
-                // if(true){
-                //     uint32_t none_zero = 0;
-                //     // uint32_t* converted_frag_B_reg_ptr = reinterpret_cast<uint32_t*>(&converted_frag_B);
-                //     uint32_t* converted_frag_B_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_B[warp_mma_k % 2]);
-                //     uint32_t* frag_a_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_A[warp_mma_k % 2]);
-                //     CUTLASS_PRAGMA_UNROLL
-                //     for(int ii=0;ii<warp_frag_B[warp_mma_k % 2].size()/8;++ii){
-                //         none_zero|=*(converted_frag_B_reg_ptr+ii);
-                //     }
-                //     CUTLASS_PRAGMA_UNROLL
-                //     for(int ii=0;ii<warp_frag_A[warp_mma_k % 2].size()/4;++ii){
-                //         none_zero|=*(frag_a_reg_ptr+ii);
-                //     }
-
-                //     CUTLASS_PRAGMA_UNROLL
-                //     for(int none_zero_i = 16; none_zero_i>0;none_zero_i/=2){
-                //         none_zero|= __shfl_xor_sync(-1,none_zero,none_zero_i);
-                //     }
-
-                //     if(none_zero!=0){
-                //         printf("## before mma ## bidtid:%d-%d-%d-%d-%d-%d, warp_mma_k:%d, frag_B_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x; frag_a_reg_ptr:%x-%x-%x-%x-%x-%x-%x-%x"
-                //                 " accu: %d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d-%d \n",
-                //                 blockIdx.x,blockIdx.y,blockIdx.z,
-                //                 warp_mma_k,
-                //                 threadIdx.x,threadIdx.y,threadIdx.z,
-                //                 converted_frag_B_reg_ptr[0],
-                //                 converted_frag_B_reg_ptr[1],
-                //                 converted_frag_B_reg_ptr[2],
-                //                 converted_frag_B_reg_ptr[3],
-                //                 converted_frag_B_reg_ptr[4],
-                //                 converted_frag_B_reg_ptr[5],
-                //                 converted_frag_B_reg_ptr[6],
-                //                 converted_frag_B_reg_ptr[7],
-                //                 frag_a_reg_ptr[0],
-                //                 frag_a_reg_ptr[1],
-                //                 frag_a_reg_ptr[2],
-                //                 frag_a_reg_ptr[3],
-                //                 frag_a_reg_ptr[4],
-                //                 frag_a_reg_ptr[5],
-                //                 frag_a_reg_ptr[6],
-                //                 frag_a_reg_ptr[7],
-                //                 accum[0],
-                //                 accum[1],
-                //                 accum[2],
-                //                 accum[3],
-                //                 accum[4],
-                //                 accum[5],
-                //                 accum[6],
-                //                 accum[7],
-                //                 accum[8],
-                //                 accum[9],
-                //                 accum[10],
-                //                 accum[11],
-                //                 accum[12],
-                //                 accum[13],
-                //                 accum[14],
-                //                 accum[15],
-                //                 accum[16],
-                //                 accum[17],
-                //                 accum[18],
-                //                 accum[19],
-                //                 accum[20],
-                //                 accum[21],
-                //                 accum[22],
-                //                 accum[23],
-                //                 accum[24],
-                //                 accum[25],
-                //                 accum[26],
-                //                 accum[27],
-                //                 accum[28],
-                //                 accum[29],
-                //                 accum[30],
-                //                 accum[31]
-                //                );
-                //     }
-                // }
-                run_ampere_warp_mma(
-                    warp_mma, accum, warp_frag_A[warp_mma_k % 2], converted_frag_B_buffer[warp_tileB_k_load_offset % 2], accum, warp_tileB_k_compute_offset);
-                // auto tmp = static_cast<int32_t>(warp_frag_B[warp_tileB_k_load_offset % 2]);
-                // if(threadIdx.x==0  && threadIdx.y==0  && threadIdx.z==0 &&
-                //     blockIdx.x==0  && blockIdx.y==0  && blockIdx.z==0){
-                //     printf("### run_warp_mma: "
-                //         "%d \n",
-                //         reinterpret_cast<int32_t &>(accum));
-                // }
-                // if(true){
-                //     uint32_t none_zero = 0;
-                //     uint32_t* converted_frag_B_reg_ptr = reinterpret_cast<uint32_t*>(&converted_frag_B);
-                //     // uint32_t* converted_frag_B_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_B[warp_mma_k % 2]);
-                //     uint32_t* frag_a_reg_ptr = reinterpret_cast<uint32_t*>(&warp_frag_A[warp_mma_k % 2]);
-                //     CUTLASS_PRAGMA_UNROLL
-                //     for(int ii=0;ii<warp_frag_B[warp_mma_k % 2].size()/8;++ii){
-                //         none_zero|=*(converted_frag_B_reg_ptr+ii);
-                //     }
-                //     CUTLASS_PRAGMA_UNROLL
-                //     for(int ii=0;ii<warp_frag_A[warp_mma_k % 2].size()/4;++ii){
-                //         none_zero|=*(frag_a_reg_ptr+ii);
-                //     }
-                //     CUTLASS_PRAGMA_UNROLL
-                //     for(int none_zero_i = 16; none_zero_i>0;none_zero_i/=2){
-                //         none_zero|= __shfl_xor_sync(-1,none_zero,none_zero_i);
-                //     }
-
-                //     // if(none_zero!=0){
-                //     if((blockIdx.y||blockIdx.z||threadIdx.x||threadIdx.y||threadIdx.z)==0){
-
-                //         printf("## after mma ## bidtid:%d-%d-%d-%d-%d-%d, warp_mma_k:%d, gemm_k_iterations:%d, Base::kWarpGemmIterations:%d,"
-                //                 " converted_frag_B_reg_ptr:%x; frag_a_reg_ptr:%x"
-                //                 " accu: %d \n",
-                //                 blockIdx.x,blockIdx.y,blockIdx.z,
-                //                 threadIdx.x,threadIdx.y,threadIdx.z,
-                //                 warp_mma_k,
-                //                 gemm_k_iterations,
-                //                 Base::kWarpGemmIterations,
-                //                 converted_frag_B_reg_ptr[0],
-                //                 frag_a_reg_ptr[0],
-                //                 accum[0]
-                //                );
-                //     }
-                // }
-                // Issue global->shared copies for the this stage
-                if (warp_mma_k < Base::kWarpGemmIterations - 1) {
-                    int group_start_iteration_A, group_start_iteration_B;
-
-                    group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
-                    group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
-
-                    copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
-                }
-
-                if (warp_mma_k + 2 == Base::kWarpGemmIterations) {
-                    int group_start_iteration_A, group_start_iteration_B;
-                    group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
-                    group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
-
-                    copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
-
-                    // Inserts a memory fence between stages of cp.async instructions.
-                    cutlass::arch::cp_async_fence();
-
-                    // Waits until kStages-2 stages have committed.
-                    arch::cp_async_wait<Base::kStages - 2>();
-                    __syncthreads();
-
-                    // Move to the next stage
-                    iterator_A.add_tile_offset({0, 1});
-                    iterator_B.add_tile_offset({1, 0});
-
-                    this->smem_iterator_A_.add_tile_offset({0, 1});
-                    this->smem_iterator_B_.add_tile_offset({1, 0});
-
-                    // Add negative offsets to return iterators to the 'start' of the
-                    // circular buffer in shared memory
-                    if (smem_write_stage_idx == (Base::kStages - 1)) {
-                        this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
-                        this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
-                        smem_write_stage_idx = 0;
-                    }
-                    else {
-                        ++smem_write_stage_idx;
-                    }
-
-                    if (smem_read_stage_idx == (Base::kStages - 1)) {
-                        this->warp_tile_iterator_A_.add_tile_offset(
-                            {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations});
-                        this->warp_tile_iterator_B_.add_tile_offset(
-                            {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterationsForB, 0});
-                        smem_read_stage_idx = 0;
-                    }
-                    else {
-                        ++smem_read_stage_idx;
-                    }
-
-                    --gemm_k_iterations;
-                    iterator_A.clear_mask(gemm_k_iterations == 0);
-                    iterator_B.clear_mask(gemm_k_iterations == 0);
-                }
-            }
+          copy_tiles_and_advance(iterator_A,
+                                 iterator_B,
+                                 group_start_iteration_A,
+                                 group_start_iteration_B);
         }
 
-        if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
-            // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop
-            cutlass::arch::cp_async_fence();
-            cutlass::arch::cp_async_wait<0>();
-            __syncthreads();
+        if (warp_mma_k + 2 == Base::kWarpGemmIterations) {
+          int group_start_iteration_A, group_start_iteration_B;
+          group_start_iteration_A =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupA;
+          group_start_iteration_B =
+              (warp_mma_k + 1) * Detail::kAccessesPerGroupB;
+
+          copy_tiles_and_advance(iterator_A,
+                                 iterator_B,
+                                 group_start_iteration_A,
+                                 group_start_iteration_B);
+
+          // Inserts a memory fence between stages of cp.async instructions.
+          cutlass::arch::cp_async_fence();
+
+          // Waits until kStages-2 stages have committed.
+          arch::cp_async_wait<Base::kStages - 2>();
+          __syncthreads();
+
+          // Move to the next stage
+          iterator_A.add_tile_offset({0, 1});
+          iterator_B.add_tile_offset({1, 0});
+
+          this->smem_iterator_A_.add_tile_offset({0, 1});
+          this->smem_iterator_B_.add_tile_offset({1, 0});
+
+          // Add negative offsets to return iterators to the 'start' of the
+          // circular buffer in shared memory
+          if (smem_write_stage_idx == (Base::kStages - 1)) {
+            this->smem_iterator_A_.add_tile_offset({0, -Base::kStages});
+            this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0});
+            smem_write_stage_idx = 0;
+          } else {
+            ++smem_write_stage_idx;
+          }
+
+          if (smem_read_stage_idx == (Base::kStages - 1)) {
+            this->warp_tile_iterator_A_.add_tile_offset(
+                {0,
+                 -Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterations});
+            this->warp_tile_iterator_B_.add_tile_offset(
+                {-Base::kStages * Policy::kPartitionsK *
+                     Base::kWarpGemmIterationsForB,
+                 0});
+            smem_read_stage_idx = 0;
+          } else {
+            ++smem_read_stage_idx;
+          }
+
+          --gemm_k_iterations;
+          iterator_A.clear_mask(gemm_k_iterations == 0);
+          iterator_B.clear_mask(gemm_k_iterations == 0);
         }
+      }
     }
+
+    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
+      // commit and drain all pending and predicated LDGSTS pnz from the GEMM
+      // mainloop
+      cutlass::arch::cp_async_fence();
+      cutlass::arch::cp_async_wait<0>();
+      __syncthreads();
+    }
+  }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/default_mma_tensor_op.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/default_mma_tensor_op.h
index 97fca6da67..1e11b8e10f 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/default_mma_tensor_op.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/default_mma_tensor_op.h
@@ -122,10 +122,9 @@ struct DefaultMmaTensorOp<WarpShape_,
                                                       AccumulatorsInRowMajor>;
 };
 
-
-
-// Specialization for int8 int8 int32 gemm, we use instruction shape<16, 8, 32> for better performance.
-template<
+// Specialization for int8 int8 int32 gemm, we use instruction shape<16, 8, 32>
+// for better performance.
+template <
     /// Shape of one matrix production operation (concept: GemmShape)
     typename WarpShape_,
     /// Shape of one matrix production operation (concept: GemmShape)
@@ -152,48 +151,51 @@ struct DefaultMmaTensorOp<WarpShape_,
                           arch::OpMultiplyAddDequantizeInterleavedBToA,
                           PartitionsK,
                           AccumulatorsInRowMajor> {
-private:
-    // Shape for computing the FP16s
-    using ComputeInstructionShape = InstructionShape_;
+ private:
+  // Shape for computing the FP16s
+  using ComputeInstructionShape = InstructionShape_;
 
-    // Chosen so we get K=32.
-    static constexpr int LoadInstructionK = 32 * sizeof_bits<int8_t>::value / sizeof_bits<int8_t>::value;
+  // Chosen so we get K=32.
+  static constexpr int LoadInstructionK =
+      32 * sizeof_bits<int8_t>::value / sizeof_bits<int8_t>::value;
 
-    // Shape for loading the narrow data type from shared memory
-    using LoadInstructionShape = GemmShape<InstructionShape_::kM, InstructionShape_::kN, LoadInstructionK>;
+  // Shape for loading the narrow data type from shared memory
+  using LoadInstructionShape =
+      GemmShape<InstructionShape_::kM, InstructionShape_::kN, LoadInstructionK>;
 
-public:
-    using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<cutlass::arch::Mma<InstructionShape_,
-                                                                             32,
-                                                                             int8_t,
-                                                                             cutlass::layout::RowMajor,
-                                                                             int8_t,
-                                                                             cutlass::layout::ColumnMajor,
-                                                                             int32_t,
-                                                                             cutlass::layout::RowMajor,
-                                                                            //  cutlass::layout::RowMajorTensorOpMultiplicandCrosswise<32,128>,
-                                                                             arch::OpMultiplyAdd>,
-                                                                             cutlass::MatrixShape<1, 1>>;
+ public:
+  using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<
+      cutlass::arch::Mma<
+          InstructionShape_,
+          32,
+          int8_t,
+          cutlass::layout::RowMajor,
+          int8_t,
+          cutlass::layout::ColumnMajor,
+          int32_t,
+          cutlass::layout::RowMajor,
+          //  cutlass::layout::RowMajorTensorOpMultiplicandCrosswise<32,128>,
+          arch::OpMultiplyAdd>,
+      cutlass::MatrixShape<1, 1>>;
 
-    // Define the warp-level tensor op
-    using Type = cutlass::gemm::warp::MmaTensorOpComputeBWithF16<WarpShape_,
-                                                                 int8_t,
-                                                                 LayoutA,
-                                                                 int8_t,
-                                                                 LayoutB,
-                                                                 int32_t,
-                                                                 LayoutC,
-                                                                 Policy,
-                                                                 LoadInstructionShape,
-                                                                 PartitionsK,
-                                                                 AccumulatorsInRowMajor>;
+  // Define the warp-level tensor op
+  using Type =
+      cutlass::gemm::warp::MmaTensorOpComputeBWithF16<WarpShape_,
+                                                      int8_t,
+                                                      LayoutA,
+                                                      int8_t,
+                                                      LayoutB,
+                                                      int32_t,
+                                                      LayoutC,
+                                                      Policy,
+                                                      LoadInstructionShape,
+                                                      PartitionsK,
+                                                      AccumulatorsInRowMajor>;
 };
 
-
-
-
-// Specialization for int8 int8 int32 gemm, we use instruction shape<16, 8, 32> for better performance.
-template<
+// Specialization for int8 int8 int32 gemm, we use instruction shape<16, 8, 32>
+// for better performance.
+template <
     /// Shape of one matrix production operation (concept: GemmShape)
     typename WarpShape_,
     /// Shape of one matrix production operation (concept: GemmShape)
@@ -220,41 +222,46 @@ struct DefaultMmaTensorOp<WarpShape_,
                           arch::OpMultiplyAddDequantizeInterleavedBToA,
                           PartitionsK,
                           AccumulatorsInRowMajor> {
-private:
-    // Shape for computing the FP16s
-    using ComputeInstructionShape = InstructionShape_;
+ private:
+  // Shape for computing the FP16s
+  using ComputeInstructionShape = InstructionShape_;
 
-    // Chosen so we get K=64.
-    static constexpr int LoadInstructionK = 16 * sizeof_bits<int8_t>::value / sizeof_bits<cutlass::uint4b_t>::value;
+  // Chosen so we get K=64.
+  static constexpr int LoadInstructionK =
+      16 * sizeof_bits<int8_t>::value / sizeof_bits<cutlass::uint4b_t>::value;
 
-    // Shape for loading the narrow data type from shared memory
-    using LoadInstructionShape = GemmShape<InstructionShape_::kM, InstructionShape_::kN, LoadInstructionK>;
+  // Shape for loading the narrow data type from shared memory
+  using LoadInstructionShape =
+      GemmShape<InstructionShape_::kM, InstructionShape_::kN, LoadInstructionK>;
 
-public:
-    using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<cutlass::arch::Mma<InstructionShape_,
-                                                                             32,
-                                                                             int8_t,
-                                                                             cutlass::layout::RowMajor,
-                                                                             int8_t,
-                                                                             cutlass::layout::ColumnMajor,
-                                                                             int32_t,
-                                                                             cutlass::layout::RowMajor,
-                                                                            //  cutlass::layout::RowMajorTensorOpMultiplicandCrosswise<32,128>,
-                                                                             arch::OpMultiplyAdd>,
-                                                                             cutlass::MatrixShape<1, 1>>;
+ public:
+  using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<
+      cutlass::arch::Mma<
+          InstructionShape_,
+          32,
+          int8_t,
+          cutlass::layout::RowMajor,
+          int8_t,
+          cutlass::layout::ColumnMajor,
+          int32_t,
+          cutlass::layout::RowMajor,
+          //  cutlass::layout::RowMajorTensorOpMultiplicandCrosswise<32,128>,
+          arch::OpMultiplyAdd>,
+      cutlass::MatrixShape<1, 1>>;
 
-    // Define the warp-level tensor op
-    using Type = cutlass::gemm::warp::MmaTensorOpComputeBWithF16<WarpShape_,
-                                                                 int8_t,
-                                                                 LayoutA,
-                                                                 cutlass::uint4b_t,
-                                                                 LayoutB,
-                                                                 int32_t,
-                                                                 LayoutC,
-                                                                 Policy,
-                                                                 LoadInstructionShape,
-                                                                 PartitionsK,
-                                                                 AccumulatorsInRowMajor>;
+  // Define the warp-level tensor op
+  using Type =
+      cutlass::gemm::warp::MmaTensorOpComputeBWithF16<WarpShape_,
+                                                      int8_t,
+                                                      LayoutA,
+                                                      cutlass::uint4b_t,
+                                                      LayoutB,
+                                                      int32_t,
+                                                      LayoutC,
+                                                      Policy,
+                                                      LoadInstructionShape,
+                                                      PartitionsK,
+                                                      AccumulatorsInRowMajor>;
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/mma_tensorop_compute_B_with_f16.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/mma_tensorop_compute_B_with_f16.h
index 7279541fe9..3476b3816c 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/mma_tensorop_compute_B_with_f16.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/mma_tensorop_compute_B_with_f16.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,19 +18,20 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*! \file
-    \brief Templates implementing warp-level matrix multiply-accumulate operations targeting
-      Tensor Cores.
+    \brief Templates implementing warp-level matrix multiply-accumulate
+   operations targeting Tensor Cores.
 */
 
 #pragma once
@@ -63,8 +64,9 @@ namespace gemm {
 namespace warp {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
-/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions.
-template<
+/// Structure to compute the matrix product targeting CUDA cores and SIMT math
+/// instructions.
+template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Data type of A elements
@@ -91,222 +93,229 @@ template<
     /// Used for partial specialization
     typename Enable = bool>
 class MmaTensorOpComputeBWithF16 {
-public:
-    /// Shape of warp-level matrix operation (concept: GemmShape)
-    using Shape = Shape_;
+ public:
+  /// Shape of warp-level matrix operation (concept: GemmShape)
+  using Shape = Shape_;
 
-    /// Data type of multiplicand A
-    using ElementA = ElementA_;
+  /// Data type of multiplicand A
+  using ElementA = ElementA_;
 
-    /// Layout of multiplicand A
-    using LayoutA = LayoutA_;
+  /// Layout of multiplicand A
+  using LayoutA = LayoutA_;
 
-    /// Data type of multiplicand B
-    using ElementB = ElementB_;
+  /// Data type of multiplicand B
+  using ElementB = ElementB_;
 
-    /// Layout of multiplicand B
-    using LayoutB = LayoutB_;
+  /// Layout of multiplicand B
+  using LayoutB = LayoutB_;
 
-    /// Data type of accumulator matrix C
-    using ElementC = ElementC_;
+  /// Data type of accumulator matrix C
+  using ElementC = ElementC_;
 
-    /// Layout of accumulator matrix C
-    using LayoutC = LayoutC_;
+  /// Layout of accumulator matrix C
+  using LayoutC = LayoutC_;
 
-    /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
-    using Policy = Policy_;
+  /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
+  using Policy = Policy_;
 
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    using ArchMmaOperator = typename Policy::Operator;
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  using ArchMmaOperator = typename Policy::Operator;
 
-    /// Indicates math operator
-    using MathOperator = typename ArchMmaOperator::Operator;
+  /// Indicates math operator
+  using MathOperator = typename ArchMmaOperator::Operator;
 
-    /// Architecture tag from underlying instruction
-    using ArchTag = typename ArchMmaOperator::ArchTag;
-    static_assert((platform::is_same<typename ArchMmaOperator::ElementA, half_t>::value
-                   && platform::is_same<typename ArchMmaOperator::ElementB, half_t>::value)
-                      || (platform::is_same<typename ArchMmaOperator::ElementA, bfloat16_t>::value
-                          && platform::is_same<typename ArchMmaOperator::ElementB, bfloat16_t>::value
-                          && ArchTag::kMinComputeCapability >= 80),
-                  "MmaTensorOpCvtBToA only supports underlying HMMA");
+  /// Architecture tag from underlying instruction
+  using ArchTag = typename ArchMmaOperator::ArchTag;
+  static_assert(
+      (platform::is_same<typename ArchMmaOperator::ElementA, half_t>::value &&
+       platform::is_same<typename ArchMmaOperator::ElementB, half_t>::value) ||
+          (platform::is_same<typename ArchMmaOperator::ElementA,
+                             bfloat16_t>::value &&
+           platform::is_same<typename ArchMmaOperator::ElementB,
+                             bfloat16_t>::value &&
+           ArchTag::kMinComputeCapability >= 80),
+      "MmaTensorOpCvtBToA only supports underlying HMMA");
 
-    static_assert(platform::is_same<ElementA, half_t>::value
-                      || (platform::is_same<ElementA, bfloat16_t>::value && ArchTag::kMinComputeCapability >= 80),
-                  "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on Ampere+");
+  static_assert(platform::is_same<ElementA, half_t>::value ||
+                    (platform::is_same<ElementA, bfloat16_t>::value &&
+                     ArchTag::kMinComputeCapability >= 80),
+                "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on Ampere+");
 
-    // Indicates class of matrix operator
-    using OperatorClass = arch::OpClassTensorOp;
+  // Indicates class of matrix operator
+  using OperatorClass = arch::OpClassTensorOp;
 
-    /// Shape of underlying instruction
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  /// Shape of underlying instruction
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    /// Instruction shape to override shared memory iterators with
-    using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
+  /// Instruction shape to override shared memory iterators with
+  using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
 
-    static_assert(SharedMemoryInstructionShape::kM == InstructionShape::kM,
-                  "M dimension of compute instruction must match load");
-    static_assert(SharedMemoryInstructionShape::kN == InstructionShape::kN,
-                  "N dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kM == InstructionShape::kM,
+                "M dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kN == InstructionShape::kN,
+                "N dimension of compute instruction must match load");
 
-    static constexpr int kExpansionFactor = SharedMemoryInstructionShape::kK / InstructionShape::kK;
+  static constexpr int kExpansionFactor =
+      SharedMemoryInstructionShape::kK / InstructionShape::kK;
 
-    static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
+  static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = ComplexTransform::kNone;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = ComplexTransform::kNone;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
 
-    /// Number of threads participating in warp-level matrix product
-    static int const kThreadCount = 32;
+  /// Number of threads participating in warp-level matrix product
+  static int const kThreadCount = 32;
 
-    /// Number of partitions along K dimension
-    static int const kPartitionsK = PartitionsK_;
+  /// Number of partitions along K dimension
+  static int const kPartitionsK = PartitionsK_;
 
-public:
-    /// Iterates over the A operand in memory
-    using IteratorA = MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kM, Shape::kK>,
-                                                          Operand::kA,
-                                                          ElementA,
-                                                          LayoutA,
-                                                          MatrixShape<InstructionShape::kM, InstructionShape::kK>,
-                                                          Policy::OpDelta::kRow,
-                                                          kThreadCount,
-                                                          kPartitionsK>;
+ public:
+  /// Iterates over the A operand in memory
+  using IteratorA = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kM, Shape::kK>,
+      Operand::kA,
+      ElementA,
+      LayoutA,
+      MatrixShape<InstructionShape::kM, InstructionShape::kK>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
 
-    /// Storage for A tile
-    using FragmentA = typename IteratorA::Fragment;
+  /// Storage for A tile
+  using FragmentA = typename IteratorA::Fragment;
 
-    /// Storage for transformed A tile
-    using TransformedFragmentA = Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
+  /// Storage for transformed A tile
+  using TransformedFragmentA =
+      Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
 
-    /// Iterates over the B operand in memory
-    using IteratorB =
-        MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kK, Shape::kN>,
-                                            Operand::kB,
-                                            ElementB,
-                                            LayoutB,
-                                            MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>,
-                                            Policy::OpDelta::kRow,
-                                            kThreadCount,
-                                            kPartitionsK>;
+  /// Iterates over the B operand in memory
+  using IteratorB = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kK, Shape::kN>,
+      Operand::kB,
+      ElementB,
+      LayoutB,
+      MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
 
-    /// Storage for B tile
-    using FragmentB = typename IteratorB::Fragment;
+  /// Storage for B tile
+  using FragmentB = typename IteratorB::Fragment;
 
-    /// Storage for transformed B tile
-    using TransformedFragmentB = Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
+  /// Storage for transformed B tile
+  using TransformedFragmentB =
+      Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
 
-    /// Iterates over the C operand in memory
-    using IteratorC = MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>,
-                                                         ElementC,
-                                                         LayoutC,
-                                                         typename ArchMmaOperator::Shape,
-                                                         typename Policy::OpDelta>;
+  /// Iterates over the C operand in memory
+  using IteratorC =
+      MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>,
+                                         ElementC,
+                                         LayoutC,
+                                         typename ArchMmaOperator::Shape,
+                                         typename Policy::OpDelta>;
 
-    /// Storage for C tile
-    using FragmentC = typename IteratorC::Fragment;
+  /// Storage for C tile
+  using FragmentC = typename IteratorC::Fragment;
 
-    /// Number of mma operations performed
-    using MmaIterations = MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM,
-                                      (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN>;
+  /// Number of mma operations performed
+  using MmaIterations =
+      MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) /
+                      ArchMmaOperator::Shape::kM,
+                  (Shape::kN + ArchMmaOperator::Shape::kN - 1) /
+                      ArchMmaOperator::Shape::kN>;
 
-public:
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    ArchMmaOperator mma;
+ public:
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  ArchMmaOperator mma;
 
-public:
-    //
-    // Methods
-    //
+ public:
+  //
+  // Methods
+  //
 
-    /// Ctor
-    CUTLASS_DEVICE
-    MmaTensorOpComputeBWithF16() {}
+  /// Ctor
+  CUTLASS_DEVICE
+  MmaTensorOpComputeBWithF16() {}
 
-    /// Performs a warp-level matrix multiply-accumulate operation
-    CUTLASS_DEVICE
-    void operator()(FragmentC&                  D,
-                    TransformedFragmentA const& A,
-                    TransformedFragmentB const& B,
-                    FragmentC const&            C,
-                    const int                   warp_tileB_k_offset) const
-    {
+  /// Performs a warp-level matrix multiply-accumulate operation
+  CUTLASS_DEVICE
+  void operator()(FragmentC& D,
+                  TransformedFragmentA const& A,
+                  TransformedFragmentB const& B,
+                  FragmentC const& C,
+                  const int warp_tileB_k_offset) const {
+    using MmaOperandA = typename ArchMmaOperator::FragmentA;
+    using MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using MmaOperandC = typename ArchMmaOperator::FragmentC;
 
-        using MmaOperandA = typename ArchMmaOperator::FragmentA;
-        using MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using MmaOperandC = typename ArchMmaOperator::FragmentC;
+    static_assert(
+        TransformedFragmentB::kElements ==
+            MmaOperandB::kElements * kExpansionFactor * MmaIterations::kColumn,
+        "Each thread should have a pack of mma registers for each column "
+        "iteration AND for the expanded K dim of B");
 
-        static_assert(
-            TransformedFragmentB::kElements == MmaOperandB::kElements * kExpansionFactor * MmaIterations::kColumn,
-            "Each thread should have a pack of mma registers for each column iteration AND for the expanded K dim of B");
+    D = C;
 
-        D = C;
-
-        MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
-        MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
-        MmaOperandC*       ptr_D = reinterpret_cast<MmaOperandC*>(&D);
+    MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
+    MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
+    MmaOperandC* ptr_D = reinterpret_cast<MmaOperandC*>(&D);
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
-        // Serpentine visitation order maximizing reuse of Rb
-        CUTLASS_PRAGMA_UNROLL
-        for (int n = 0; n < MmaIterations::kColumn; ++n) {
+    // Serpentine visitation order maximizing reuse of Rb
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < MmaIterations::kColumn; ++n) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int m = 0; m < MmaIterations::kRow; ++m) {
+        int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
 
-            CUTLASS_PRAGMA_UNROLL
-            for (int m = 0; m < MmaIterations::kRow; ++m) {
-
-                int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
-
-                int n_offsetB = warp_tileB_k_offset + kExpansionFactor * n;
-                if (AccumulatorsInRowMajor) {  // matrix B is reordered
-                    mma(ptr_D[n + m_serpentine * MmaIterations::kColumn],
-                        ptr_A[m_serpentine],
-                        ptr_B[n_offsetB],
-                        ptr_D[n + m_serpentine * MmaIterations::kColumn]);
-                }
-                else {
-                    mma(ptr_D[m_serpentine + n * MmaIterations::kRow],
-                        ptr_A[m_serpentine],
-                        ptr_B[n_offsetB],
-                        ptr_D[m_serpentine + n * MmaIterations::kRow]);
-                }
-            }
+        int n_offsetB = warp_tileB_k_offset + kExpansionFactor * n;
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n + m_serpentine * MmaIterations::kColumn],
+              ptr_A[m_serpentine],
+              ptr_B[n_offsetB],
+              ptr_D[n + m_serpentine * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m_serpentine + n * MmaIterations::kRow],
+              ptr_A[m_serpentine],
+              ptr_B[n_offsetB],
+              ptr_D[m_serpentine + n * MmaIterations::kRow]);
         }
-#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
-        // Serpentine visitation order maximizing reuse of Ra
-        CUTLASS_PRAGMA_UNROLL
-        for (int m = 0; m < MmaIterations::kRow; ++m) {
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int n = 0; n < MmaIterations::kColumn; ++n) {
-
-                int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
-
-                int n_serpentine_offsetB = warp_tileB_k_offset + kExpansionFactor * n_serpentine;
-                if (AccumulatorsInRowMajor) {  // matrix B is reordered
-                    mma(ptr_D[n_serpentine + m * MmaIterations::kColumn],
-                        ptr_A[m],
-                        ptr_B[n_serpentine_offsetB],
-                        ptr_D[n_serpentine + m * MmaIterations::kColumn]);
-                }
-                else {
-                    mma(ptr_D[m + n_serpentine * MmaIterations::kRow],
-                        ptr_A[m],
-                        ptr_B[n_serpentine_offsetB],
-                        ptr_D[m + n_serpentine * MmaIterations::kRow]);
-                }
-            }
-        }
-#else
-        assert(0);
-#endif
+      }
     }
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+    // Serpentine visitation order maximizing reuse of Ra
+    CUTLASS_PRAGMA_UNROLL
+    for (int m = 0; m < MmaIterations::kRow; ++m) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < MmaIterations::kColumn; ++n) {
+        int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
+
+        int n_serpentine_offsetB =
+            warp_tileB_k_offset + kExpansionFactor * n_serpentine;
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n_serpentine + m * MmaIterations::kColumn],
+              ptr_A[m],
+              ptr_B[n_serpentine_offsetB],
+              ptr_D[n_serpentine + m * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m + n_serpentine * MmaIterations::kRow],
+              ptr_A[m],
+              ptr_B[n_serpentine_offsetB],
+              ptr_D[m + n_serpentine * MmaIterations::kRow]);
+        }
+      }
+    }
+#else
+    assert(0);
+#endif
+  }
 };
 
 // Specialization for int8 int8 int32. Author(zhengzekang)
-template<
+template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Layout of A matrix (concept: MatrixLayout)
@@ -328,233 +337,235 @@ template<
     bool AccumulatorsInRowMajor,
     /// Used for partial specialization
     typename Enable>
-class MmaTensorOpComputeBWithF16<
-    Shape_,
-    int8_t,
-    LayoutA_,
-    int8_t,
-    LayoutB_,
-    ElementC_,
-    LayoutC_,
-    Policy_,
-    SharedMemoryInstructionShape_,
-    PartitionsK_,
-    AccumulatorsInRowMajor,
-    Enable> {
-public:
-    /// Shape of warp-level matrix operation (concept: GemmShape)
-    using Shape = Shape_;
+class MmaTensorOpComputeBWithF16<Shape_,
+                                 int8_t,
+                                 LayoutA_,
+                                 int8_t,
+                                 LayoutB_,
+                                 ElementC_,
+                                 LayoutC_,
+                                 Policy_,
+                                 SharedMemoryInstructionShape_,
+                                 PartitionsK_,
+                                 AccumulatorsInRowMajor,
+                                 Enable> {
+ public:
+  /// Shape of warp-level matrix operation (concept: GemmShape)
+  using Shape = Shape_;
 
-    /// Data type of multiplicand A
-    using ElementA = int8_t;
+  /// Data type of multiplicand A
+  using ElementA = int8_t;
 
-    /// Layout of multiplicand A
-    using LayoutA = LayoutA_;
+  /// Layout of multiplicand A
+  using LayoutA = LayoutA_;
 
-    /// Data type of multiplicand B
-    using ElementB = int8_t;
+  /// Data type of multiplicand B
+  using ElementB = int8_t;
 
-    /// Layout of multiplicand B
-    using LayoutB = LayoutB_;
+  /// Layout of multiplicand B
+  using LayoutB = LayoutB_;
 
-    /// Data type of accumulator matrix C
-    using ElementC = ElementC_;
+  /// Data type of accumulator matrix C
+  using ElementC = ElementC_;
 
-    /// Layout of accumulator matrix C
-    using LayoutC = LayoutC_;
+  /// Layout of accumulator matrix C
+  using LayoutC = LayoutC_;
 
-    /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
-    using Policy = Policy_;
+  /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
+  using Policy = Policy_;
 
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    using ArchMmaOperator = typename Policy::Operator;
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  using ArchMmaOperator = typename Policy::Operator;
 
-    /// Indicates math operator
-    using MathOperator = typename ArchMmaOperator::Operator;
+  /// Indicates math operator
+  using MathOperator = typename ArchMmaOperator::Operator;
 
-    /// Architecture tag from underlying instruction
-    using ArchTag = typename ArchMmaOperator::ArchTag;
-    static_assert((platform::is_same<typename ArchMmaOperator::ElementA, int8_t>::value
-                   && platform::is_same<typename ArchMmaOperator::ElementB, int8_t>::value),
-                  "MmaTensorOpCvtBToA only supports underlying iMMA");
+  /// Architecture tag from underlying instruction
+  using ArchTag = typename ArchMmaOperator::ArchTag;
+  static_assert(
+      (platform::is_same<typename ArchMmaOperator::ElementA, int8_t>::value &&
+       platform::is_same<typename ArchMmaOperator::ElementB, int8_t>::value),
+      "MmaTensorOpCvtBToA only supports underlying iMMA");
 
-    // static_assert(platform::is_same<ElementA, int8_t>::value,
-    //               "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on Ampere+");
+  // static_assert(platform::is_same<ElementA, int8_t>::value,
+  //               "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on
+  //               Ampere+");
 
-    /// Indicates class of matrix operator
-    using OperatorClass = arch::OpClassTensorOp;
+  /// Indicates class of matrix operator
+  using OperatorClass = arch::OpClassTensorOp;
 
-    /// Shape of underlying instruction
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  /// Shape of underlying instruction
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    /// Instruction shape to override shared memory iterators with
-    using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
+  /// Instruction shape to override shared memory iterators with
+  using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
 
-    static_assert(SharedMemoryInstructionShape::kM == InstructionShape::kM,
-                  "M dimension of compute instruction must match load");
-    static_assert(SharedMemoryInstructionShape::kN == InstructionShape::kN,
-                  "N dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kM == InstructionShape::kM,
+                "M dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kN == InstructionShape::kN,
+                "N dimension of compute instruction must match load");
 
-    static constexpr int kExpansionFactor = SharedMemoryInstructionShape::kK / InstructionShape::kK;
+  static constexpr int kExpansionFactor =
+      SharedMemoryInstructionShape::kK / InstructionShape::kK;
 
-    static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
+  static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = ComplexTransform::kNone;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = ComplexTransform::kNone;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
 
-    /// Number of threads participating in warp-level matrix product
-    static int const kThreadCount = 32;
+  /// Number of threads participating in warp-level matrix product
+  static int const kThreadCount = 32;
 
-    /// Number of partitions along K dimension
-    static int const kPartitionsK = PartitionsK_;
+  /// Number of partitions along K dimension
+  static int const kPartitionsK = PartitionsK_;
 
-public:
-    /// Iterates over the A operand in memory
-    using IteratorA = MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kM, Shape::kK>,
-                                                          Operand::kA,
-                                                          ElementA,
-                                                          LayoutA,
-                                                          MatrixShape<InstructionShape::kM, InstructionShape::kK>,
-                                                          Policy::OpDelta::kRow,
-                                                          kThreadCount,
-                                                          kPartitionsK>;
+ public:
+  /// Iterates over the A operand in memory
+  using IteratorA = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kM, Shape::kK>,
+      Operand::kA,
+      ElementA,
+      LayoutA,
+      MatrixShape<InstructionShape::kM, InstructionShape::kK>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
 
-    /// Storage for A tile
-    using FragmentA = typename IteratorA::Fragment;
+  /// Storage for A tile
+  using FragmentA = typename IteratorA::Fragment;
 
-    /// Storage for transformed A tile
-    using TransformedFragmentA = Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
+  /// Storage for transformed A tile
+  using TransformedFragmentA =
+      Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
 
-    /// Iterates over the B operand in memory
-    using IteratorB =
-        MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kK, Shape::kN>,
-                                            Operand::kB,
-                                            ElementB,
-                                            LayoutB,
-                                            MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>,
-                                            Policy::OpDelta::kRow,
-                                            kThreadCount,
-                                            kPartitionsK>;
+  /// Iterates over the B operand in memory
+  using IteratorB = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kK, Shape::kN>,
+      Operand::kB,
+      ElementB,
+      LayoutB,
+      MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
 
-    /// Storage for B tile
-    using FragmentB = typename IteratorB::Fragment;
+  /// Storage for B tile
+  using FragmentB = typename IteratorB::Fragment;
 
-    /// Storage for transformed B tile
-    using TransformedFragmentB = Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
+  /// Storage for transformed B tile
+  using TransformedFragmentB =
+      Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
 
-    /// Iterates over the C operand in memory
-    using IteratorC = MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>,
-                                                         ElementC,
-                                                         LayoutC,
-                                                         typename ArchMmaOperator::Shape,
-                                                         typename Policy::OpDelta>;
+  /// Iterates over the C operand in memory
+  using IteratorC =
+      MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>,
+                                         ElementC,
+                                         LayoutC,
+                                         typename ArchMmaOperator::Shape,
+                                         typename Policy::OpDelta>;
 
-    /// Storage for C tile
-    using FragmentC = typename IteratorC::Fragment;
+  /// Storage for C tile
+  using FragmentC = typename IteratorC::Fragment;
 
-    /// Number of mma operations performed
-    using MmaIterations = MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM,
-                                      (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN>;
+  /// Number of mma operations performed
+  using MmaIterations =
+      MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) /
+                      ArchMmaOperator::Shape::kM,
+                  (Shape::kN + ArchMmaOperator::Shape::kN - 1) /
+                      ArchMmaOperator::Shape::kN>;
 
-public:
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    ArchMmaOperator mma;
+ public:
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  ArchMmaOperator mma;
 
-public:
-    //
-    // Methods
-    //
+ public:
+  //
+  // Methods
+  //
 
-    /// Ctor
-    CUTLASS_DEVICE
-    MmaTensorOpComputeBWithF16() {}
+  /// Ctor
+  CUTLASS_DEVICE
+  MmaTensorOpComputeBWithF16() {}
 
-    /// Performs a warp-level matrix multiply-accumulate operation
-    CUTLASS_DEVICE
-    void operator()(FragmentC&                  D,
-                    TransformedFragmentA const& A,
-                    TransformedFragmentB const& B,
-                    FragmentC const&            C,
-                    const int                   warp_tileB_k_offset) const
-    {
+  /// Performs a warp-level matrix multiply-accumulate operation
+  CUTLASS_DEVICE
+  void operator()(FragmentC& D,
+                  TransformedFragmentA const& A,
+                  TransformedFragmentB const& B,
+                  FragmentC const& C,
+                  const int warp_tileB_k_offset) const {
+    using MmaOperandA = typename ArchMmaOperator::FragmentA;
+    using MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using MmaOperandC = typename ArchMmaOperator::FragmentC;
 
-        using MmaOperandA = typename ArchMmaOperator::FragmentA;
-        using MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using MmaOperandC = typename ArchMmaOperator::FragmentC;
+    static_assert(
+        TransformedFragmentB::kElements ==
+            MmaOperandB::kElements * kExpansionFactor * MmaIterations::kColumn,
+        "Each thread should have a pack of mma registers for each column "
+        "iteration AND for the expanded K dim of B");
 
-        static_assert(
-            TransformedFragmentB::kElements == MmaOperandB::kElements * kExpansionFactor * MmaIterations::kColumn,
-            "Each thread should have a pack of mma registers for each column iteration AND for the expanded K dim of B");
+    D = C;
 
-        D = C;
-
-        MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
-        MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
-        MmaOperandC*       ptr_D = reinterpret_cast<MmaOperandC*>(&D);
+    MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
+    MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
+    MmaOperandC* ptr_D = reinterpret_cast<MmaOperandC*>(&D);
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
-        // Serpentine visitation order maximizing reuse of Rb
-        CUTLASS_PRAGMA_UNROLL
-        for (int n = 0; n < MmaIterations::kColumn; ++n) {
+    // Serpentine visitation order maximizing reuse of Rb
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < MmaIterations::kColumn; ++n) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int m = 0; m < MmaIterations::kRow; ++m) {
+        int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
 
-            CUTLASS_PRAGMA_UNROLL
-            for (int m = 0; m < MmaIterations::kRow; ++m) {
-
-                int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
-
-                int n_offsetB = warp_tileB_k_offset + kExpansionFactor * n;
-                if (AccumulatorsInRowMajor) {  // matrix B is reordered
-                    mma(ptr_D[n + m_serpentine * MmaIterations::kColumn],
-                        ptr_A[m_serpentine],
-                        ptr_B[n_offsetB],
-                        ptr_D[n + m_serpentine * MmaIterations::kColumn]);
-                }
-                else {
-                    mma(ptr_D[m_serpentine + n * MmaIterations::kRow],
-                        ptr_A[m_serpentine],
-                        ptr_B[n_offsetB],
-                        ptr_D[m_serpentine + n * MmaIterations::kRow]);
-                }
-            }
+        int n_offsetB = warp_tileB_k_offset + kExpansionFactor * n;
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n + m_serpentine * MmaIterations::kColumn],
+              ptr_A[m_serpentine],
+              ptr_B[n_offsetB],
+              ptr_D[n + m_serpentine * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m_serpentine + n * MmaIterations::kRow],
+              ptr_A[m_serpentine],
+              ptr_B[n_offsetB],
+              ptr_D[m_serpentine + n * MmaIterations::kRow]);
         }
-#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
-        // Serpentine visitation order maximizing reuse of Ra
-        CUTLASS_PRAGMA_UNROLL
-        for (int m = 0; m < MmaIterations::kRow; ++m) {
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int n = 0; n < MmaIterations::kColumn; ++n) {
-
-                int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
-
-                int n_serpentine_offsetB = warp_tileB_k_offset + kExpansionFactor * n_serpentine;
-                if (AccumulatorsInRowMajor) {  // matrix B is reordered
-                    mma(ptr_D[n_serpentine + m * MmaIterations::kColumn],
-                        ptr_A[m],
-                        ptr_B[n_serpentine_offsetB],
-                        ptr_D[n_serpentine + m * MmaIterations::kColumn]);
-                }
-                else {
-                    mma(ptr_D[m + n_serpentine * MmaIterations::kRow],
-                        ptr_A[m],
-                        ptr_B[n_serpentine_offsetB],
-                        ptr_D[m + n_serpentine * MmaIterations::kRow]);
-                }
-            }
-        }
-#else
-        assert(0);
-#endif
+      }
     }
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+    // Serpentine visitation order maximizing reuse of Ra
+    CUTLASS_PRAGMA_UNROLL
+    for (int m = 0; m < MmaIterations::kRow; ++m) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < MmaIterations::kColumn; ++n) {
+        int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
+
+        int n_serpentine_offsetB =
+            warp_tileB_k_offset + kExpansionFactor * n_serpentine;
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n_serpentine + m * MmaIterations::kColumn],
+              ptr_A[m],
+              ptr_B[n_serpentine_offsetB],
+              ptr_D[n_serpentine + m * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m + n_serpentine * MmaIterations::kRow],
+              ptr_A[m],
+              ptr_B[n_serpentine_offsetB],
+              ptr_D[m + n_serpentine * MmaIterations::kRow]);
+        }
+      }
+    }
+#else
+    assert(0);
+#endif
+  }
 };
 
-
-
 // Specialization for int8 int8 int32. Author(zhengzekang)
-template<
+template <
     /// Size of the Gemm problem - concept: gemm::GemmShape<>
     typename Shape_,
     /// Layout of A matrix (concept: MatrixLayout)
@@ -576,235 +587,250 @@ template<
     bool AccumulatorsInRowMajor,
     /// Used for partial specialization
     typename Enable>
-class MmaTensorOpComputeBWithF16<
-    Shape_,
-    int8_t,
-    LayoutA_,
-    cutlass::uint4b_t,
-    LayoutB_,
-    ElementC_,
-    LayoutC_,
-    Policy_,
-    SharedMemoryInstructionShape_,
-    PartitionsK_,
-    AccumulatorsInRowMajor,
-    Enable> {
-public:
-    /// Shape of warp-level matrix operation (concept: GemmShape)
-    using Shape = Shape_;
+class MmaTensorOpComputeBWithF16<Shape_,
+                                 int8_t,
+                                 LayoutA_,
+                                 cutlass::uint4b_t,
+                                 LayoutB_,
+                                 ElementC_,
+                                 LayoutC_,
+                                 Policy_,
+                                 SharedMemoryInstructionShape_,
+                                 PartitionsK_,
+                                 AccumulatorsInRowMajor,
+                                 Enable> {
+ public:
+  /// Shape of warp-level matrix operation (concept: GemmShape)
+  using Shape = Shape_;
 
-    /// Data type of multiplicand A
-    using ElementA = int8_t;
+  /// Data type of multiplicand A
+  using ElementA = int8_t;
 
-    /// Layout of multiplicand A
-    using LayoutA = LayoutA_;
+  /// Layout of multiplicand A
+  using LayoutA = LayoutA_;
 
-    /// Data type of multiplicand B
-    using ElementB = cutlass::uint4b_t;
+  /// Data type of multiplicand B
+  using ElementB = cutlass::uint4b_t;
 
-    /// Layout of multiplicand B
-    using LayoutB = LayoutB_;
+  /// Layout of multiplicand B
+  using LayoutB = LayoutB_;
 
-    /// Data type of accumulator matrix C
-    using ElementC = ElementC_;
+  /// Data type of accumulator matrix C
+  using ElementC = ElementC_;
 
-    /// Layout of accumulator matrix C
-    using LayoutC = LayoutC_;
+  /// Layout of accumulator matrix C
+  using LayoutC = LayoutC_;
 
-    /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
-    using Policy = Policy_;
+  /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
+  using Policy = Policy_;
 
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    using ArchMmaOperator = typename Policy::Operator;
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  using ArchMmaOperator = typename Policy::Operator;
 
-    /// Indicates math operator
-    using MathOperator = typename ArchMmaOperator::Operator;
+  /// Indicates math operator
+  using MathOperator = typename ArchMmaOperator::Operator;
 
-    /// Architecture tag from underlying instruction
-    using ArchTag = typename ArchMmaOperator::ArchTag;
-    static_assert((platform::is_same<typename ArchMmaOperator::ElementA, int8_t>::value
-                   && platform::is_same<typename ArchMmaOperator::ElementB, int8_t>::value),
-                  "MmaTensorOpCvtBToA only supports underlying iMMA");
+  /// Architecture tag from underlying instruction
+  using ArchTag = typename ArchMmaOperator::ArchTag;
+  static_assert(
+      (platform::is_same<typename ArchMmaOperator::ElementA, int8_t>::value &&
+       platform::is_same<typename ArchMmaOperator::ElementB, int8_t>::value),
+      "MmaTensorOpCvtBToA only supports underlying iMMA");
 
-    // static_assert(platform::is_same<ElementA, int8_t>::value,
-    //               "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on Ampere+");
+  // static_assert(platform::is_same<ElementA, int8_t>::value,
+  //               "MmaTensorOpCvtBToA only supports Fp16 A or Bf16 A on
+  //               Ampere+");
 
-    /// Indicates class of matrix operator
-    using OperatorClass = arch::OpClassTensorOp;
+  /// Indicates class of matrix operator
+  using OperatorClass = arch::OpClassTensorOp;
 
-    /// Shape of underlying instruction
-    using InstructionShape = typename ArchMmaOperator::Shape;
+  /// Shape of underlying instruction
+  using InstructionShape = typename ArchMmaOperator::Shape;
 
-    /// Instruction shape to override shared memory iterators with
-    using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
+  /// Instruction shape to override shared memory iterators with
+  using SharedMemoryInstructionShape = SharedMemoryInstructionShape_;
 
-    static_assert(SharedMemoryInstructionShape::kM == InstructionShape::kM,
-                  "M dimension of compute instruction must match load");
-    static_assert(SharedMemoryInstructionShape::kN == InstructionShape::kN,
-                  "N dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kM == InstructionShape::kM,
+                "M dimension of compute instruction must match load");
+  static_assert(SharedMemoryInstructionShape::kN == InstructionShape::kN,
+                "N dimension of compute instruction must match load");
 
-    static constexpr int kExpansionFactor = SharedMemoryInstructionShape::kK / InstructionShape::kK;
+  static constexpr int kExpansionFactor =
+      SharedMemoryInstructionShape::kK / InstructionShape::kK;
 
-    static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
+  static_assert(!(Shape::kK % SharedMemoryInstructionShape::kK), "");
 
-    /// Complex transform on A operand
-    static ComplexTransform const kTransformA = ComplexTransform::kNone;
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
 
-    /// Complex transform on B operand
-    static ComplexTransform const kTransformB = ComplexTransform::kNone;
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
 
-    /// Number of threads participating in warp-level matrix product
-    static int const kThreadCount = 32;
+  /// Number of threads participating in warp-level matrix product
+  static int const kThreadCount = 32;
 
-    /// Number of partitions along K dimension
-    static int const kPartitionsK = PartitionsK_;
+  /// Number of partitions along K dimension
+  static int const kPartitionsK = PartitionsK_;
 
-public:
-    /// Iterates over the A operand in memory
+ public:
+  /// Iterates over the A operand in memory
 
+  // ::cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator<
+  // ::cutlass::MatrixShape<(int)32, (int)64> , ( ::cutlass::gemm::Operand)0,
+  // signed char,
+  // ::cutlass::layout::RowMajorTensorOpMultiplicandCrosswise<(int)8, (int)64> ,
+  // ::cutlass::MatrixShape<(int)16, (int)32> , (int)1, (int)32, (int)1> > ()")
+  using IteratorA = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kM, Shape::kK>,
+      Operand::kA,
+      ElementA,
+      LayoutA,
+      MatrixShape<InstructionShape::kM, InstructionShape::kK>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
 
-    // ::cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator< ::cutlass::MatrixShape<(int)32, (int)64> , ( ::cutlass::gemm::Operand)0, signed char,  ::cutlass::layout::RowMajorTensorOpMultiplicandCrosswise<(int)8, (int)64> ,  ::cutlass::MatrixShape<(int)16, (int)32> , (int)1, (int)32, (int)1> > ()")
-    using IteratorA = MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kM, Shape::kK>,
-                                                          Operand::kA,
-                                                          ElementA,
-                                                          LayoutA,
-                                                          MatrixShape<InstructionShape::kM, InstructionShape::kK>,
-                                                          Policy::OpDelta::kRow,
-                                                          kThreadCount,
-                                                          kPartitionsK>;
+  /// Storage for A tile
+  using FragmentA = typename IteratorA::Fragment;
 
-    /// Storage for A tile
-    using FragmentA = typename IteratorA::Fragment;
+  /// Storage for transformed A tile
+  using TransformedFragmentA =
+      Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
+  // bool  ::print_type<
+  // ::cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator<
+  // ::cutlass::MatrixShape<(int)64, (int)32> , ( ::cutlass::gemm::Operand)1,
+  // ::cutlass::integer_subbyte<(int)4, (bool)0> ,
+  // ::cutlass::layout::ColumnMajorTensorOpMultiplicandCrosswise<(int)4,
+  // (int)64> ,  ::cutlass::MatrixShape<(int)64, (int)8> , (int)1, (int)32,
+  // (int)1> > ()")
+  /// Iterates over the B operand in memory
+  using IteratorB = MmaTensorOpMultiplicandTileIterator<
+      MatrixShape<Shape::kK, Shape::kN>,
+      Operand::kB,
+      ElementB,
+      LayoutB,
+      MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>,
+      Policy::OpDelta::kRow,
+      kThreadCount,
+      kPartitionsK>;
+  // ::cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator<
+  // ::cutlass::MatrixShape<(int)64, (int)32> , ( ::cutlass::gemm::Operand)1,
+  // ::cutlass::integer_subbyte<(int)4, (bool)0> ,
+  // ::cutlass::layout::ColumnMajorTensorOpMultiplicandCrosswise<(int)4,
+  // (int)64> ,  ::cutlass::MatrixShape<(int)64, (int)8> , (int)1, (int)32,
+  // (int)1> > print_type<typename IteratorB>();
+  /// Storage for B tile
+  using FragmentB = typename IteratorB::Fragment;
 
-    /// Storage for transformed A tile
-    using TransformedFragmentA = Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
-// bool  ::print_type< ::cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator< ::cutlass::MatrixShape<(int)64, (int)32> , ( ::cutlass::gemm::Operand)1,  ::cutlass::integer_subbyte<(int)4, (bool)0> ,  ::cutlass::layout::ColumnMajorTensorOpMultiplicandCrosswise<(int)4, (int)64> ,  ::cutlass::MatrixShape<(int)64, (int)8> , (int)1, (int)32, (int)1> > ()")
-    /// Iterates over the B operand in memory
-    using IteratorB =
-        MmaTensorOpMultiplicandTileIterator<MatrixShape<Shape::kK, Shape::kN>,
-                                            Operand::kB,
-                                            ElementB,
-                                            LayoutB,
-                                            MatrixShape<SharedMemoryInstructionShape::kK, InstructionShape::kN>,
-                                            Policy::OpDelta::kRow,
-                                            kThreadCount,
-                                            kPartitionsK>;
-// ::cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator< ::cutlass::MatrixShape<(int)64, (int)32> , ( ::cutlass::gemm::Operand)1,  ::cutlass::integer_subbyte<(int)4, (bool)0> ,  ::cutlass::layout::ColumnMajorTensorOpMultiplicandCrosswise<(int)4, (int)64> ,  ::cutlass::MatrixShape<(int)64, (int)8> , (int)1, (int)32, (int)1> >
-    // print_type<typename IteratorB>();
-    /// Storage for B tile
-    using FragmentB = typename IteratorB::Fragment;
+  /// Storage for transformed B tile
+  using TransformedFragmentB =
+      Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
 
-    /// Storage for transformed B tile
-    using TransformedFragmentB = Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
+  /// Iterates over the C operand in memory
+  using IteratorC =
+      MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>,
+                                         ElementC,
+                                         LayoutC,
+                                         typename ArchMmaOperator::Shape,
+                                         typename Policy::OpDelta>;
 
-    /// Iterates over the C operand in memory
-    using IteratorC = MmaTensorOpAccumulatorTileIterator<MatrixShape<Shape::kM, Shape::kN>,
-                                                         ElementC,
-                                                         LayoutC,
-                                                         typename ArchMmaOperator::Shape,
-                                                         typename Policy::OpDelta>;
+  /// Storage for C tile
+  using FragmentC = typename IteratorC::Fragment;
 
-    /// Storage for C tile
-    using FragmentC = typename IteratorC::Fragment;
+  /// Number of mma operations performed
+  using MmaIterations =
+      MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) /
+                      ArchMmaOperator::Shape::kM,
+                  (Shape::kN + ArchMmaOperator::Shape::kN - 1) /
+                      ArchMmaOperator::Shape::kN>;
 
-    /// Number of mma operations performed
-    using MmaIterations = MatrixShape<(Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM,
-                                      (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN>;
+ public:
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  ArchMmaOperator mma;
 
-public:
-    /// Underlying matrix multiply operator (concept: arch::Mma)
-    ArchMmaOperator mma;
+ public:
+  //
+  // Methods
+  //
 
-public:
-    //
-    // Methods
-    //
+  /// Ctor
+  CUTLASS_DEVICE
+  MmaTensorOpComputeBWithF16() {}
 
-    /// Ctor
-    CUTLASS_DEVICE
-    MmaTensorOpComputeBWithF16() {}
+  /// Performs a warp-level matrix multiply-accumulate operation
+  CUTLASS_DEVICE
+  void operator()(FragmentC& D,
+                  TransformedFragmentA const& A,
+                  TransformedFragmentB const& B,
+                  FragmentC const& C,
+                  const int warp_tileB_k_offset) const {
+    using MmaOperandA = typename ArchMmaOperator::FragmentA;
+    using MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using MmaOperandC = typename ArchMmaOperator::FragmentC;
 
-    /// Performs a warp-level matrix multiply-accumulate operation
-    CUTLASS_DEVICE
-    void operator()(FragmentC&                  D,
-                    TransformedFragmentA const& A,
-                    TransformedFragmentB const& B,
-                    FragmentC const&            C,
-                    const int                   warp_tileB_k_offset) const
-    {
+    static_assert(
+        TransformedFragmentB::kElements ==
+            MmaOperandB::kElements * kExpansionFactor * MmaIterations::kColumn,
+        "Each thread should have a pack of mma registers for each column "
+        "iteration AND for the expanded K dim of B");
 
-        using MmaOperandA = typename ArchMmaOperator::FragmentA;
-        using MmaOperandB = typename ArchMmaOperator::FragmentB;
-        using MmaOperandC = typename ArchMmaOperator::FragmentC;
+    D = C;
 
-        static_assert(
-            TransformedFragmentB::kElements == MmaOperandB::kElements * kExpansionFactor * MmaIterations::kColumn,
-            "Each thread should have a pack of mma registers for each column iteration AND for the expanded K dim of B");
-
-        D = C;
-
-        MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
-        MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
-        MmaOperandC*       ptr_D = reinterpret_cast<MmaOperandC*>(&D);
+    MmaOperandA const* ptr_A = reinterpret_cast<MmaOperandA const*>(&A);
+    MmaOperandB const* ptr_B = reinterpret_cast<MmaOperandB const*>(&B);
+    MmaOperandC* ptr_D = reinterpret_cast<MmaOperandC*>(&D);
 
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
-        // Serpentine visitation order maximizing reuse of Rb
-        CUTLASS_PRAGMA_UNROLL
-        for (int n = 0; n < MmaIterations::kColumn; ++n) {
+    // Serpentine visitation order maximizing reuse of Rb
+    CUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < MmaIterations::kColumn; ++n) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int m = 0; m < MmaIterations::kRow; ++m) {
+        int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
 
-            CUTLASS_PRAGMA_UNROLL
-            for (int m = 0; m < MmaIterations::kRow; ++m) {
-
-                int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
-
-                int n_offsetB = warp_tileB_k_offset + kExpansionFactor * n;
-                if (AccumulatorsInRowMajor) {  // matrix B is reordered
-                    mma(ptr_D[n + m_serpentine * MmaIterations::kColumn],
-                        ptr_A[m_serpentine],
-                        ptr_B[n_offsetB],
-                        ptr_D[n + m_serpentine * MmaIterations::kColumn]);
-                }
-                else {
-                    mma(ptr_D[m_serpentine + n * MmaIterations::kRow],
-                        ptr_A[m_serpentine],
-                        ptr_B[n_offsetB],
-                        ptr_D[m_serpentine + n * MmaIterations::kRow]);
-                }
-            }
+        int n_offsetB = warp_tileB_k_offset + kExpansionFactor * n;
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n + m_serpentine * MmaIterations::kColumn],
+              ptr_A[m_serpentine],
+              ptr_B[n_offsetB],
+              ptr_D[n + m_serpentine * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m_serpentine + n * MmaIterations::kRow],
+              ptr_A[m_serpentine],
+              ptr_B[n_offsetB],
+              ptr_D[m_serpentine + n * MmaIterations::kRow]);
         }
-#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
-        // Serpentine visitation order maximizing reuse of Ra
-        CUTLASS_PRAGMA_UNROLL
-        for (int m = 0; m < MmaIterations::kRow; ++m) {
-
-            CUTLASS_PRAGMA_UNROLL
-            for (int n = 0; n < MmaIterations::kColumn; ++n) {
-
-                int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
-
-                int n_serpentine_offsetB = warp_tileB_k_offset + kExpansionFactor * n_serpentine;
-                if (AccumulatorsInRowMajor) {  // matrix B is reordered
-                    mma(ptr_D[n_serpentine + m * MmaIterations::kColumn],
-                        ptr_A[m],
-                        ptr_B[n_serpentine_offsetB],
-                        ptr_D[n_serpentine + m * MmaIterations::kColumn]);
-                }
-                else {
-                    mma(ptr_D[m + n_serpentine * MmaIterations::kRow],
-                        ptr_A[m],
-                        ptr_B[n_serpentine_offsetB],
-                        ptr_D[m + n_serpentine * MmaIterations::kRow]);
-                }
-            }
-        }
-#else
-        assert(0);
-#endif
+      }
     }
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+    // Serpentine visitation order maximizing reuse of Ra
+    CUTLASS_PRAGMA_UNROLL
+    for (int m = 0; m < MmaIterations::kRow; ++m) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < MmaIterations::kColumn; ++n) {
+        int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);
+
+        int n_serpentine_offsetB =
+            warp_tileB_k_offset + kExpansionFactor * n_serpentine;
+        if (AccumulatorsInRowMajor) {  // matrix B is reordered
+          mma(ptr_D[n_serpentine + m * MmaIterations::kColumn],
+              ptr_A[m],
+              ptr_B[n_serpentine_offsetB],
+              ptr_D[n_serpentine + m * MmaIterations::kColumn]);
+        } else {
+          mma(ptr_D[m + n_serpentine * MmaIterations::kRow],
+              ptr_A[m],
+              ptr_B[n_serpentine_offsetB],
+              ptr_D[m + n_serpentine * MmaIterations::kRow]);
+        }
+      }
+    }
+#else
+    assert(0);
+#endif
+  }
 };
 
-
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 }  // namespace warp
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h
index 73a1122433..73c8fbb457 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h
@@ -187,10 +187,9 @@ class MmaTensorOpDequantizer<
 
   // Adds a pointer offset in units of elements.
   CUTLASS_DEVICE
-  void add_pointer_offset(int64_t const& offset)
-  {
-      static_assert(sizeof(ElementScale) > 1, "");
-      pointer_ += offset;
+  void add_pointer_offset(int64_t const& offset) {
+    static_assert(sizeof(ElementScale) > 1, "");
+    pointer_ += offset;
   }
 
  private:
@@ -297,10 +296,9 @@ class MmaTensorOpDequantizer<
   }
   // Adds a pointer offset in units of elements.
   CUTLASS_DEVICE
-  void add_pointer_offset(int64_t const& offset)
-  {
-      static_assert(sizeof(ElementScale) > 1, "");
-      pointer_ += offset;
+  void add_pointer_offset(int64_t const& offset) {
+    static_assert(sizeof(ElementScale) > 1, "");
+    pointer_ += offset;
   }
 
  private:
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/interleaved_numeric_conversion.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/interleaved_numeric_conversion.h
index 6f09be73e3..d68a683f4a 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/interleaved_numeric_conversion.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/interleaved_numeric_conversion.h
@@ -435,180 +435,197 @@ struct FastInterleavedAndBiasedNumericArrayConverter<bfloat16_t, uint4b_t, N> {
   result_type operator()(source_type const& s) { return convert(s); }
 };
 
-template<>
+template <>
 struct FastInterleavedAndBiasedNumericArrayConverter<int8_t, int8_t, 4> {
-    using result_type = Array<int8_t, 4>;
-    using source_type = Array<int8_t, 4>;
+  using result_type = Array<int8_t, 4>;
+  using source_type = Array<int8_t, 4>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 
-        uint32_t*      h   = reinterpret_cast<uint32_t*>(&result);
-        uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
 
-        // 3 2 1 0 -> 3 1 2 0
-        static constexpr uint32_t mask_for_elt     = 0x3120;
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[0]) : "r"(i8s), "r"(i8s), "n"(mask_for_elt));
+    // 3 2 1 0 -> 3 1 2 0
+    static constexpr uint32_t mask_for_elt = 0x3120;
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[0])
+                 : "r"(i8s), "r"(i8s), "n"(mask_for_elt));
 
-        // Author zhengzekang
-        uint8_t*  tmp  = reinterpret_cast<uint8_t*>(&result);
-        #pragma unroll
-        for(int i = 0; i < 4; i++){
-            result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
-        }
-        return result;
+    // Author zhengzekang
+    uint8_t* tmp = reinterpret_cast<uint8_t*>(&result);
+#pragma unroll
+    for (int i = 0; i < 4; i++) {
+      result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
     }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
-
-template<>
+template <>
 struct FastInterleavedAndBiasedNumericArrayConverter<int8_t, int8_t, 8> {
-    using result_type = Array<int8_t, 8>;
-    using source_type = Array<int8_t, 8>;
+  using result_type = Array<int8_t, 8>;
+  using source_type = Array<int8_t, 8>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 
-        uint32_t*      h   = reinterpret_cast<uint32_t*>(&result);
-        const uint32_t * i8s = reinterpret_cast<const uint32_t*>(&source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    const uint32_t* i8s = reinterpret_cast<const uint32_t*>(&source);
 
-        // 3 2 1 0 -> 3 1 2 0
-        static constexpr uint32_t mask_for_elt_1     = 0x7120;
-        static constexpr uint32_t mask_for_elt_2     = 0x3654;
+    // 3 2 1 0 -> 3 1 2 0
+    static constexpr uint32_t mask_for_elt_1 = 0x7120;
+    static constexpr uint32_t mask_for_elt_2 = 0x3654;
 
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[0]) : "r"(i8s[0]), "r"(i8s[1]), "n"(mask_for_elt_1));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[1]) : "r"(i8s[0]), "r"(i8s[1]), "n"(mask_for_elt_2));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[0])
+                 : "r"(i8s[0]), "r"(i8s[1]), "n"(mask_for_elt_1));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[1])
+                 : "r"(i8s[0]), "r"(i8s[1]), "n"(mask_for_elt_2));
 
-        // Author zhengzekang
-        uint8_t*  tmp  = reinterpret_cast<uint8_t*>(&result);
-        #pragma unroll
-        for(int i = 0; i < 8; i++){
-            result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
-        }
-        return result;
+    // Author zhengzekang
+    uint8_t* tmp = reinterpret_cast<uint8_t*>(&result);
+#pragma unroll
+    for (int i = 0; i < 8; i++) {
+      result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
     }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
-
-template<>
+template <>
 struct FastInterleavedAndBiasedNumericArrayConverter<int8_t, int8_t, 16> {
-    using result_type = Array<int8_t, 16>;
-    using source_type = Array<int8_t, 16>;
+  using result_type = Array<int8_t, 16>;
+  using source_type = Array<int8_t, 16>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 
-        uint32_t*      h   = reinterpret_cast<uint32_t*>(&result);
-        const uint32_t * i8s = reinterpret_cast<const uint32_t*>(&source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    const uint32_t* i8s = reinterpret_cast<const uint32_t*>(&source);
 
-        // 3 2 1 0 -> 3 1 2 0
-        static constexpr uint32_t mask_for_elt_1     = 0x3120;
-        static constexpr uint32_t mask_for_elt_2     = 0x3120;
-        static constexpr uint32_t mask_for_elt_3     = 0x3120;
-        static constexpr uint32_t mask_for_elt_4     = 0x3120;
+    // 3 2 1 0 -> 3 1 2 0
+    static constexpr uint32_t mask_for_elt_1 = 0x3120;
+    static constexpr uint32_t mask_for_elt_2 = 0x3120;
+    static constexpr uint32_t mask_for_elt_3 = 0x3120;
+    static constexpr uint32_t mask_for_elt_4 = 0x3120;
 
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[0]) : "r"(i8s[0]), "r"(i8s[2]), "n"(mask_for_elt_1));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[1]) : "r"(i8s[1]), "r"(i8s[3]), "n"(mask_for_elt_2));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[2]) : "r"(i8s[2]), "r"(i8s[0]), "n"(mask_for_elt_3));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[3]) : "r"(i8s[3]), "r"(i8s[1]), "n"(mask_for_elt_4));
-
-
-        // Author zhengzekang
-        uint8_t*  tmp  = reinterpret_cast<uint8_t*>(&result);
-        #pragma unroll
-        for(int i = 0; i < 16; i++){
-            result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
-        }
-
-        return result;
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[0])
+                 : "r"(i8s[0]), "r"(i8s[2]), "n"(mask_for_elt_1));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[1])
+                 : "r"(i8s[1]), "r"(i8s[3]), "n"(mask_for_elt_2));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[2])
+                 : "r"(i8s[2]), "r"(i8s[0]), "n"(mask_for_elt_3));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[3])
+                 : "r"(i8s[3]), "r"(i8s[1]), "n"(mask_for_elt_4));
 
+    // Author zhengzekang
+    uint8_t* tmp = reinterpret_cast<uint8_t*>(&result);
+#pragma unroll
+    for (int i = 0; i < 16; i++) {
+      result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
     }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
-template<int N>
+template <int N>
 struct FastInterleavedAndBiasedNumericArrayConverter<int8_t, int8_t, N> {
-    static constexpr int VEC_WIDTH = 4;
-    static_assert(N == 32,"N must be 32");
-    static_assert(!(N % VEC_WIDTH), "N must be multiple of 16.");
+  static constexpr int VEC_WIDTH = 4;
+  static_assert(N == 32, "N must be 32");
+  static_assert(!(N % VEC_WIDTH), "N must be multiple of 16.");
 
-    using result_type = Array<int8_t, N>;
-    using source_type = Array<int8_t, N>;
+  using result_type = Array<int8_t, N>;
+  using source_type = Array<int8_t, N>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type, scalar_source_type, VEC_WIDTH>
-            convert_vector_;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    FastInterleavedAndBiasedNumericArrayConverter<scalar_result_type,
+                                                  scalar_source_type,
+                                                  VEC_WIDTH>
+        convert_vector_;
 
-        result_type result;
-        using vec_result = Array<scalar_result_type, VEC_WIDTH>;
-        using vec_source = Array<scalar_source_type, VEC_WIDTH>;
+    result_type result;
+    using vec_result = Array<scalar_result_type, VEC_WIDTH>;
+    using vec_source = Array<scalar_source_type, VEC_WIDTH>;
 
-        vec_result*       result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int32_t i = 0; i < N / VEC_WIDTH; ++i) {
-            result_ptr[i] = convert_vector_((source_ptr)[i]);
-        }
-        Array<int8_t, 8> temp_rearrange_array;
-        auto lane_idx_div4 = threadIdx.x%4;
+    CUTLASS_PRAGMA_UNROLL
+    for (int32_t i = 0; i < N / VEC_WIDTH; ++i) {
+      result_ptr[i] = convert_vector_((source_ptr)[i]);
+    }
+    Array<int8_t, 8> temp_rearrange_array;
+    auto lane_idx_div4 = threadIdx.x % 4;
 
-        CUTLASS_PRAGMA_UNROLL
-        for (int32_t i=0;i< N/8;++i){
-            uint32_t*       temp_rearrange_array_ptr = reinterpret_cast<uint32_t*>(&temp_rearrange_array);
-            uint32_t*       result_reg_ptr = reinterpret_cast<uint32_t*>(result_ptr)+i * 2;
-            temp_rearrange_array_ptr[0] = __shfl_xor_sync(0xFFFFFFFF,reinterpret_cast<uint32_t*>(result_reg_ptr)[0],3);
-            temp_rearrange_array_ptr[1] = __shfl_xor_sync(0xFFFFFFFF,reinterpret_cast<uint32_t*>(result_reg_ptr)[1],3);
-            if( lane_idx_div4==1 || lane_idx_div4==2 ){
-                result_reg_ptr[0]=temp_rearrange_array_ptr[0];
-                result_reg_ptr[1]=temp_rearrange_array_ptr[1];
-            }
-            temp_rearrange_array_ptr[0] = __shfl_xor_sync(0xFFFFFFFF,reinterpret_cast<uint32_t*>(result_reg_ptr)[0],2);
-            temp_rearrange_array_ptr[1] = __shfl_xor_sync(0xFFFFFFFF,reinterpret_cast<uint32_t*>(result_reg_ptr)[1],2);
-            if(lane_idx_div4<2){
-                asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(result_reg_ptr[0]) : "r"(result_reg_ptr[0]), "r"(temp_rearrange_array_ptr[0]), "n"(0x5410));
-                asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(result_reg_ptr[1]) : "r"(result_reg_ptr[1]), "r"(temp_rearrange_array_ptr[1]), "n"(0x5410));
-            }
-            else{
-                asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(result_reg_ptr[0]) : "r"(result_reg_ptr[0]), "r"(temp_rearrange_array_ptr[0]), "n"(0x3276));
-                asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(result_reg_ptr[1]) : "r"(result_reg_ptr[1]), "r"(temp_rearrange_array_ptr[1]), "n"(0x3276));
-            }
-        }
-
-        return result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int32_t i = 0; i < N / 8; ++i) {
+      uint32_t* temp_rearrange_array_ptr =
+          reinterpret_cast<uint32_t*>(&temp_rearrange_array);
+      uint32_t* result_reg_ptr =
+          reinterpret_cast<uint32_t*>(result_ptr) + i * 2;
+      temp_rearrange_array_ptr[0] = __shfl_xor_sync(
+          0xFFFFFFFF, reinterpret_cast<uint32_t*>(result_reg_ptr)[0], 3);
+      temp_rearrange_array_ptr[1] = __shfl_xor_sync(
+          0xFFFFFFFF, reinterpret_cast<uint32_t*>(result_reg_ptr)[1], 3);
+      if (lane_idx_div4 == 1 || lane_idx_div4 == 2) {
+        result_reg_ptr[0] = temp_rearrange_array_ptr[0];
+        result_reg_ptr[1] = temp_rearrange_array_ptr[1];
+      }
+      temp_rearrange_array_ptr[0] = __shfl_xor_sync(
+          0xFFFFFFFF, reinterpret_cast<uint32_t*>(result_reg_ptr)[0], 2);
+      temp_rearrange_array_ptr[1] = __shfl_xor_sync(
+          0xFFFFFFFF, reinterpret_cast<uint32_t*>(result_reg_ptr)[1], 2);
+      if (lane_idx_div4 < 2) {
+        asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                     : "=r"(result_reg_ptr[0])
+                     : "r"(result_reg_ptr[0]),
+                       "r"(temp_rearrange_array_ptr[0]),
+                       "n"(0x5410));
+        asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                     : "=r"(result_reg_ptr[1])
+                     : "r"(result_reg_ptr[1]),
+                       "r"(temp_rearrange_array_ptr[1]),
+                       "n"(0x5410));
+      } else {
+        asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                     : "=r"(result_reg_ptr[0])
+                     : "r"(result_reg_ptr[0]),
+                       "r"(temp_rearrange_array_ptr[0]),
+                       "n"(0x3276));
+        asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                     : "=r"(result_reg_ptr[1])
+                     : "r"(result_reg_ptr[1]),
+                       "r"(temp_rearrange_array_ptr[1]),
+                       "n"(0x3276));
+      }
     }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/interleaved_numeric_conversion_nf4.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/interleaved_numeric_conversion_nf4.h
index b21b204c85..c4f03583ce 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/interleaved_numeric_conversion_nf4.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/interleaved_numeric_conversion_nf4.h
@@ -104,8 +104,8 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<half_t, uint8_t, N> {
     using scalar_result_type = typename result_type::Element;
     using scalar_source_type = typename source_type::Element;
     FastInterleavedAndBiasedNumericArrayConverterNf4<scalar_result_type,
-                                                  scalar_source_type,
-                                                  VEC_WIDTH>
+                                                     scalar_source_type,
+                                                     VEC_WIDTH>
         convert_vector_;
 
     result_type result;
@@ -128,7 +128,9 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<half_t, uint8_t, N> {
 };
 
 template <>
-struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint8_t, 4> {
+struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t,
+                                                        uint8_t,
+                                                        4> {
   using result_type = Array<bfloat16_t, 4>;
   using source_type = Array<uint8_t, 4>;
 
@@ -179,7 +181,9 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint8_t, 4>
 };
 
 template <int N>
-struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint8_t, N> {
+struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t,
+                                                        uint8_t,
+                                                        N> {
   static constexpr int VEC_WIDTH = 4;
   static_assert(!(N % VEC_WIDTH), "N must be multiple of 4.");
 
@@ -191,8 +195,8 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint8_t, N>
     using scalar_result_type = typename result_type::Element;
     using scalar_source_type = typename source_type::Element;
     FastInterleavedAndBiasedNumericArrayConverterNf4<scalar_result_type,
-                                                  scalar_source_type,
-                                                  VEC_WIDTH>
+                                                     scalar_source_type,
+                                                     VEC_WIDTH>
         convert_vector_;
 
     result_type result;
@@ -314,8 +318,8 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<half_t, uint4b_t, N> {
     using scalar_result_type = typename result_type::Element;
     using scalar_source_type = typename source_type::Element;
     FastInterleavedAndBiasedNumericArrayConverterNf4<scalar_result_type,
-                                                  scalar_source_type,
-                                                  VEC_WIDTH>
+                                                     scalar_source_type,
+                                                     VEC_WIDTH>
         convert_vector_;
 
     result_type result;
@@ -338,7 +342,9 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<half_t, uint4b_t, N> {
 };
 
 template <>
-struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint4b_t, 8> {
+struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t,
+                                                        uint4b_t,
+                                                        8> {
   using result_type = Array<bfloat16_t, 8>;
   using source_type = Array<uint4b_t, 8>;
 
@@ -400,7 +406,9 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint4b_t, 8>
 };
 
 template <int N>
-struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint4b_t, N> {
+struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t,
+                                                        uint4b_t,
+                                                        N> {
   static constexpr int VEC_WIDTH = 8;
   static_assert(!(N % VEC_WIDTH), "N must be multiple of 8.");
 
@@ -412,8 +420,8 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint4b_t, N>
     using scalar_result_type = typename result_type::Element;
     using scalar_source_type = typename source_type::Element;
     FastInterleavedAndBiasedNumericArrayConverterNf4<scalar_result_type,
-                                                  scalar_source_type,
-                                                  VEC_WIDTH>
+                                                     scalar_source_type,
+                                                     VEC_WIDTH>
         convert_vector_;
 
     result_type result;
@@ -435,388 +443,405 @@ struct FastInterleavedAndBiasedNumericArrayConverterNf4<bfloat16_t, uint4b_t, N>
   result_type operator()(source_type const& s) { return convert(s); }
 };
 
-template<>
+template <>
 struct FastInterleavedAndBiasedNumericArrayConverterNf4<int8_t, int8_t, 4> {
-    using result_type = Array<int8_t, 4>;
-    using source_type = Array<int8_t, 4>;
+  using result_type = Array<int8_t, 4>;
+  using source_type = Array<int8_t, 4>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 
-        uint32_t*      h   = reinterpret_cast<uint32_t*>(&result);
-        uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    uint32_t const i8s = reinterpret_cast<uint32_t const&>(source);
 
-        // 3 2 1 0 -> 3 1 2 0
-        static constexpr uint32_t mask_for_elt     = 0x3120;
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[0]) : "r"(i8s), "r"(i8s), "n"(mask_for_elt));
+    // 3 2 1 0 -> 3 1 2 0
+    static constexpr uint32_t mask_for_elt = 0x3120;
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[0])
+                 : "r"(i8s), "r"(i8s), "n"(mask_for_elt));
 
-        // Author zhengzekang
-        uint8_t*  tmp  = reinterpret_cast<uint8_t*>(&result);
-        #pragma unroll
-        for(int i = 0; i < 4; i++){
-            result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
-        }
-        return result;
+    // Author zhengzekang
+    uint8_t* tmp = reinterpret_cast<uint8_t*>(&result);
+#pragma unroll
+    for (int i = 0; i < 4; i++) {
+      result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
     }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
-
-template<>
+template <>
 struct FastInterleavedAndBiasedNumericArrayConverterNf4<int8_t, int8_t, 8> {
-    using result_type = Array<int8_t, 8>;
-    using source_type = Array<int8_t, 8>;
+  using result_type = Array<int8_t, 8>;
+  using source_type = Array<int8_t, 8>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 
-        uint32_t*      h   = reinterpret_cast<uint32_t*>(&result);
-        const uint32_t * i8s = reinterpret_cast<const uint32_t*>(&source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    const uint32_t* i8s = reinterpret_cast<const uint32_t*>(&source);
 
-        // 3 2 1 0 -> 3 1 2 0
-        static constexpr uint32_t mask_for_elt_1     = 0x7120;
-        static constexpr uint32_t mask_for_elt_2     = 0x3654;
+    // 3 2 1 0 -> 3 1 2 0
+    static constexpr uint32_t mask_for_elt_1 = 0x7120;
+    static constexpr uint32_t mask_for_elt_2 = 0x3654;
 
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[0]) : "r"(i8s[0]), "r"(i8s[1]), "n"(mask_for_elt_1));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[1]) : "r"(i8s[0]), "r"(i8s[1]), "n"(mask_for_elt_2));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[0])
+                 : "r"(i8s[0]), "r"(i8s[1]), "n"(mask_for_elt_1));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[1])
+                 : "r"(i8s[0]), "r"(i8s[1]), "n"(mask_for_elt_2));
 
-        // Author zhengzekang
-        uint8_t*  tmp  = reinterpret_cast<uint8_t*>(&result);
-        #pragma unroll
-        for(int i = 0; i < 8; i++){
-            result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
-        }
-        return result;
+    // Author zhengzekang
+    uint8_t* tmp = reinterpret_cast<uint8_t*>(&result);
+#pragma unroll
+    for (int i = 0; i < 8; i++) {
+      result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
     }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
-
-template<>
+template <>
 struct FastInterleavedAndBiasedNumericArrayConverterNf4<int8_t, int8_t, 16> {
-    using result_type = Array<int8_t, 16>;
-    using source_type = Array<int8_t, 16>;
+  using result_type = Array<int8_t, 16>;
+  using source_type = Array<int8_t, 16>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        result_type result;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    result_type result;
 
-        uint32_t*      h   = reinterpret_cast<uint32_t*>(&result);
-        const uint32_t * i8s = reinterpret_cast<const uint32_t*>(&source);
+    uint32_t* h = reinterpret_cast<uint32_t*>(&result);
+    const uint32_t* i8s = reinterpret_cast<const uint32_t*>(&source);
 
-        // 3 2 1 0 -> 3 1 2 0
-        static constexpr uint32_t mask_for_elt_1     = 0x3120;
-        static constexpr uint32_t mask_for_elt_2     = 0x3120;
-        static constexpr uint32_t mask_for_elt_3     = 0x3120;
-        static constexpr uint32_t mask_for_elt_4     = 0x3120;
+    // 3 2 1 0 -> 3 1 2 0
+    static constexpr uint32_t mask_for_elt_1 = 0x3120;
+    static constexpr uint32_t mask_for_elt_2 = 0x3120;
+    static constexpr uint32_t mask_for_elt_3 = 0x3120;
+    static constexpr uint32_t mask_for_elt_4 = 0x3120;
 
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[0]) : "r"(i8s[0]), "r"(i8s[2]), "n"(mask_for_elt_1));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[1]) : "r"(i8s[1]), "r"(i8s[3]), "n"(mask_for_elt_2));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[2]) : "r"(i8s[2]), "r"(i8s[0]), "n"(mask_for_elt_3));
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(h[3]) : "r"(i8s[3]), "r"(i8s[1]), "n"(mask_for_elt_4));
-
-
-        // Author zhengzekang
-        uint8_t*  tmp  = reinterpret_cast<uint8_t*>(&result);
-        #pragma unroll
-        for(int i = 0; i < 16; i++){
-            result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
-        }
-
-        return result;
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[0])
+                 : "r"(i8s[0]), "r"(i8s[2]), "n"(mask_for_elt_1));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[1])
+                 : "r"(i8s[1]), "r"(i8s[3]), "n"(mask_for_elt_2));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[2])
+                 : "r"(i8s[2]), "r"(i8s[0]), "n"(mask_for_elt_3));
+    asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                 : "=r"(h[3])
+                 : "r"(i8s[3]), "r"(i8s[1]), "n"(mask_for_elt_4));
 
+    // Author zhengzekang
+    uint8_t* tmp = reinterpret_cast<uint8_t*>(&result);
+#pragma unroll
+    for (int i = 0; i < 16; i++) {
+      result[i] = static_cast<int8_t>(static_cast<int32_t>(tmp[i]) - 128);
     }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
-template<int N>
+template <int N>
 struct FastInterleavedAndBiasedNumericArrayConverterNf4<int8_t, int8_t, N> {
-    static constexpr int VEC_WIDTH = 4;
-    static_assert(N == 32,"N must be 32");
-    static_assert(!(N % VEC_WIDTH), "N must be multiple of 16.");
+  static constexpr int VEC_WIDTH = 4;
+  static_assert(N == 32, "N must be 32");
+  static_assert(!(N % VEC_WIDTH), "N must be multiple of 16.");
 
-    using result_type = Array<int8_t, N>;
-    using source_type = Array<int8_t, N>;
+  using result_type = Array<int8_t, N>;
+  using source_type = Array<int8_t, N>;
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        return source;
-    }
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) { return source; }
 
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
 };
 
-template<int N>
-struct FastInterleavedAndBiasedNumericArrayConverterNf4<int8_t, cutlass::uint4b_t, N> {
-    static constexpr int VEC_WIDTH = 8;
-    // static_assert(N == 64,"N must be 64");
-    static_assert(!(N % VEC_WIDTH), "N must be multiple of VEC_WIDTH.");
+template <int N>
+struct FastInterleavedAndBiasedNumericArrayConverterNf4<int8_t,
+                                                        cutlass::uint4b_t,
+                                                        N> {
+  static constexpr int VEC_WIDTH = 8;
+  // static_assert(N == 64,"N must be 64");
+  static_assert(!(N % VEC_WIDTH), "N must be multiple of VEC_WIDTH.");
 
-    using result_type = Array<int8_t, N>;
-    using source_type = Array<cutlass::uint4b_t, N>;
-    using vec_source = Array<cutlass::uint4b_t, VEC_WIDTH>;
-    using vec_result = Array<int8_t, VEC_WIDTH>;
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source)
-    {
-        //nf4
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        result_type result;
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / VEC_WIDTH; ++i) {
-          vec_result index;
-          uint32_t* result_i_ptr = reinterpret_cast<uint32_t*>(&(result_ptr[i]));
-          uint32_t const i4s = reinterpret_cast<uint32_t const&>(source_ptr[i]);
-          static constexpr uint32_t up_int4_mask = 0xf0f0f0f0;
-          static constexpr uint32_t immLut_0 = 0x40;
-          static constexpr uint32_t immLut_1 = 0x80;
-          asm volatile(
-              "lop3.b32 %0, %1, %2, %3, %4;\n"
-              : "=r"(result_i_ptr[1])
-              : "r"(i4s), "r"(i4s), "n"(up_int4_mask), "n"(immLut_1));
-          asm volatile(
-              "lop3.b32 %0, %1, %2, %3, %4;\n"
-              : "=r"(result_i_ptr[0])
-              : "r"(i4s), "r"(i4s), "n"(up_int4_mask), "n"(immLut_0));
-          result_i_ptr[0]=(result_i_ptr[0]<<4);
-        }
-        return result;
+  using result_type = Array<int8_t, N>;
+  using source_type = Array<cutlass::uint4b_t, N>;
+  using vec_source = Array<cutlass::uint4b_t, VEC_WIDTH>;
+  using vec_result = Array<int8_t, VEC_WIDTH>;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source) {
+    // nf4
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    result_type result;
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / VEC_WIDTH; ++i) {
+      vec_result index;
+      uint32_t* result_i_ptr = reinterpret_cast<uint32_t*>(&(result_ptr[i]));
+      uint32_t const i4s = reinterpret_cast<uint32_t const&>(source_ptr[i]);
+      static constexpr uint32_t up_int4_mask = 0xf0f0f0f0;
+      static constexpr uint32_t immLut_0 = 0x40;
+      static constexpr uint32_t immLut_1 = 0x80;
+      asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+                   : "=r"(result_i_ptr[1])
+                   : "r"(i4s), "r"(i4s), "n"(up_int4_mask), "n"(immLut_1));
+      asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+                   : "=r"(result_i_ptr[0])
+                   : "r"(i4s), "r"(i4s), "n"(up_int4_mask), "n"(immLut_0));
+      result_i_ptr[0] = (result_i_ptr[0] << 4);
     }
+    return result;
+  }
 
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source, int32_t* shared_look_up_table)
-    {
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        result_type result;
-        // static constexpr uint32_t loop_up_table[4]{0x03020100,0x07060504,0x0B0A0908,0x0F0E0D0C};
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
-        // FastInterleavedAndBiasedNumericArrayConverterNf4<scalar_result_type, scalar_source_type, VEC_WIDTH>
-        //     convert_vector_;
-        // static constexpr Array<uint32_t, 4> loop_up_table{0x00010203,0x04050607,0x08090A0B,0x0C0D0E0F};
-        for (int i = 0; i < N / VEC_WIDTH; ++i) {
-          vec_result index;
-          uint32_t* h = reinterpret_cast<uint32_t*>(&index);
-          // const int8_t* loop_up_table_int8 = reinterpret_cast<const int8_t*>(&loop_up_table);
-          uint32_t const i4s = reinterpret_cast<uint32_t const&>(source_ptr[i]);
-          static constexpr uint32_t down_int4_mask = 0x0f0f0f0f;
-          static constexpr uint32_t up_int4_mask = 0xf0f0f0f0;
-          h[0]=i4s&down_int4_mask;
-          h[1]=i4s&up_int4_mask;
-          h[1]=h[1]>>4;
+  CUTLASS_DEVICE
+  static result_type convert(source_type const& source,
+                             int32_t* shared_look_up_table) {
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    result_type result;
+    // static constexpr uint32_t
+    // loop_up_table[4]{0x03020100,0x07060504,0x0B0A0908,0x0F0E0D0C};
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    // FastInterleavedAndBiasedNumericArrayConverterNf4<scalar_result_type,
+    // scalar_source_type, VEC_WIDTH>
+    //     convert_vector_;
+    // static constexpr Array<uint32_t, 4>
+    // loop_up_table{0x00010203,0x04050607,0x08090A0B,0x0C0D0E0F};
+    for (int i = 0; i < N / VEC_WIDTH; ++i) {
+      vec_result index;
+      uint32_t* h = reinterpret_cast<uint32_t*>(&index);
+      // const int8_t* loop_up_table_int8 = reinterpret_cast<const
+      // int8_t*>(&loop_up_table);
+      uint32_t const i4s = reinterpret_cast<uint32_t const&>(source_ptr[i]);
+      static constexpr uint32_t down_int4_mask = 0x0f0f0f0f;
+      static constexpr uint32_t up_int4_mask = 0xf0f0f0f0;
+      h[0] = i4s & down_int4_mask;
+      h[1] = i4s & up_int4_mask;
+      h[1] = h[1] >> 4;
 
-          //TODO(wangbojun)!!!! do nf4 lookup table
-          CUTLASS_PRAGMA_UNROLL
-          for(int ii=0; ii<VEC_WIDTH;++ii){
-            result_ptr[i][ii]=shared_look_up_table[index[ii]];
-          }
-          // CUTLASS_PRAGMA_UNROLL
-          // for(int ii=0; ii<VEC_WIDTH;++ii){
-          //   result_ptr[i][ii]=static_cast<int8_t>(index[ii]);
-          // }
+      // TODO(wangbojun)!!!! do nf4 lookup table
+      CUTLASS_PRAGMA_UNROLL
+      for (int ii = 0; ii < VEC_WIDTH; ++ii) {
+        result_ptr[i][ii] = shared_look_up_table[index[ii]];
+      }
+      // CUTLASS_PRAGMA_UNROLL
+      // for(int ii=0; ii<VEC_WIDTH;++ii){
+      //   result_ptr[i][ii]=static_cast<int8_t>(index[ii]);
+      // }
 
-          // CUTLASS_PRAGMA_UNROLL
-          // for(int ii=0; ii<VEC_WIDTH;++ii){
-          //   result_ptr[i][ii]=0;
-          // }
+      // CUTLASS_PRAGMA_UNROLL
+      // for(int ii=0; ii<VEC_WIDTH;++ii){
+      //   result_ptr[i][ii]=0;
+      // }
 
-          // printf("### i:%d index:%d-%d-%d-%d-%d-%d-%d-%d, result:%d-%d-%d-%d-%d-%d-%d-%d\n",
-          //       i,
-          //       static_cast<int32_t>(index[0]),
-          //       static_cast<int32_t>(index[1]),
-          //       static_cast<int32_t>(index[2]),
-          //       static_cast<int32_t>(index[3]),
-          //       static_cast<int32_t>(index[4]),
-          //       static_cast<int32_t>(index[5]),
-          //       static_cast<int32_t>(index[6]),
-          //       static_cast<int32_t>(index[7]),
-          //       static_cast<int32_t>(result_ptr[i][0]),
-          //       static_cast<int32_t>(result_ptr[i][1]),
-          //       static_cast<int32_t>(result_ptr[i][2]),
-          //       static_cast<int32_t>(result_ptr[i][3]),
-          //       static_cast<int32_t>(result_ptr[i][4]),
-          //       static_cast<int32_t>(result_ptr[i][5]),
-          //       static_cast<int32_t>(result_ptr[i][6]),
-          //       static_cast<int32_t>(result_ptr[i][7])
-          //       );
-        }
-        return result;
+      // printf("### i:%d index:%d-%d-%d-%d-%d-%d-%d-%d,
+      // result:%d-%d-%d-%d-%d-%d-%d-%d\n",
+      //       i,
+      //       static_cast<int32_t>(index[0]),
+      //       static_cast<int32_t>(index[1]),
+      //       static_cast<int32_t>(index[2]),
+      //       static_cast<int32_t>(index[3]),
+      //       static_cast<int32_t>(index[4]),
+      //       static_cast<int32_t>(index[5]),
+      //       static_cast<int32_t>(index[6]),
+      //       static_cast<int32_t>(index[7]),
+      //       static_cast<int32_t>(result_ptr[i][0]),
+      //       static_cast<int32_t>(result_ptr[i][1]),
+      //       static_cast<int32_t>(result_ptr[i][2]),
+      //       static_cast<int32_t>(result_ptr[i][3]),
+      //       static_cast<int32_t>(result_ptr[i][4]),
+      //       static_cast<int32_t>(result_ptr[i][5]),
+      //       static_cast<int32_t>(result_ptr[i][6]),
+      //       static_cast<int32_t>(result_ptr[i][7])
+      //       );
     }
+    return result;
+  }
 
+  // #define NF4_LUT_DEBUG
+  CUTLASS_DEVICE
+  static result_type convert(
+      source_type const& source,
+      cutlass::Array<uint32_t, 4, true> const& reg_look_up_table) {
+    // static_assert(VEC_WIDTH==16, "VEC_WIDTH == 16 for int8 int8 int32")
+    using scalar_result_type = typename result_type::Element;
+    using scalar_source_type = typename source_type::Element;
+    result_type result;
+    // static constexpr uint32_t
+    // loop_up_table[4]{0x03020100,0x07060504,0x0B0A0908,0x0F0E0D0C};
+    vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
+    vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
+    // FastInterleavedAndBiasedNumericArrayConverterNf4<scalar_result_type,
+    // scalar_source_type, VEC_WIDTH>
+    //     convert_vector_;
+    // static constexpr Array<uint32_t, 4>
+    // loop_up_table{0x00010203,0x04050607,0x08090A0B,0x0C0D0E0F};
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / VEC_WIDTH; ++i) {
+      uint32_t bitwise_workspace;
+      uint32_t* result_reg = reinterpret_cast<uint32_t*>(&(result_ptr[i]));
+      uint32_t const i4s = reinterpret_cast<uint32_t const&>(source_ptr[i]);
+      static constexpr uint32_t down_int4_mask = 0x0f0f0f0f;
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### lup 0:%08x, lup 1:%08x, lup2:%08x, lup3:%08x \n",
+               reg_look_up_table[0],
+               reg_look_up_table[1],
+               reg_look_up_table[2],
+               reg_look_up_table[3]);
+        printf("#### i4s:%08x \n", i4s);
+      }
+#endif
+      bitwise_workspace = i4s & down_int4_mask;  // 0x0v0v0v0v
 
-// #define NF4_LUT_DEBUG
-    CUTLASS_DEVICE
-    static result_type convert(source_type const& source, cutlass::Array<uint32_t, 4, true> const& reg_look_up_table)
-    {
-        // static_assert(VEC_WIDTH==16, "VEC_WIDTH == 16 for int8 int8 int32")
-        using scalar_result_type = typename result_type::Element;
-        using scalar_source_type = typename source_type::Element;
-        result_type result;
-        // static constexpr uint32_t loop_up_table[4]{0x03020100,0x07060504,0x0B0A0908,0x0F0E0D0C};
-        vec_result* result_ptr = reinterpret_cast<vec_result*>(&result);
-        vec_source const* source_ptr = reinterpret_cast<vec_source const*>(&source);
-        // FastInterleavedAndBiasedNumericArrayConverterNf4<scalar_result_type, scalar_source_type, VEC_WIDTH>
-        //     convert_vector_;
-        // static constexpr Array<uint32_t, 4> loop_up_table{0x00010203,0x04050607,0x08090A0B,0x0C0D0E0F};
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < N / VEC_WIDTH; ++i) {
-          uint32_t bitwise_workspace;
-          uint32_t* result_reg = reinterpret_cast<uint32_t*>(&(result_ptr[i]));
-          uint32_t const i4s = reinterpret_cast<uint32_t const&>(source_ptr[i]);
-          static constexpr uint32_t down_int4_mask = 0x0f0f0f0f;
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### lup 0:%08x, lup 1:%08x, lup2:%08x, lup3:%08x \n",
-                                                                              reg_look_up_table[0],
-                                                                              reg_look_up_table[1],
-                                                                              reg_look_up_table[2],
-                                                                              reg_look_up_table[3]);
-                                                                        printf("#### i4s:%08x \n", i4s);
-                                                                      }
-                                                                    #endif
-          bitwise_workspace = i4s & down_int4_mask; // 0x0v0v0v0v
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### h0 0x0v0v0v0v:%08x \n", bitwise_workspace);
+      }
+#endif
 
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### h0 0x0v0v0v0v:%08x \n", bitwise_workspace);
-                                                                      }
-                                                                    #endif
+      uint32_t ge_7_mask_h_0 =
+          (bitwise_workspace | bitwise_workspace << 4) & 0x88888888;
+      ge_7_mask_h_0 |= ge_7_mask_h_0 >> 1;
+      ge_7_mask_h_0 |= ge_7_mask_h_0 >> 2;
 
-          uint32_t ge_7_mask_h_0 = (bitwise_workspace |bitwise_workspace << 4)  & 0x88888888;
-          ge_7_mask_h_0 |= ge_7_mask_h_0 >> 1;
-          ge_7_mask_h_0 |= ge_7_mask_h_0 >> 2;
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### ge_7_mask_h_0:%08x \n", ge_7_mask_h_0);
+      }
+#endif
 
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### ge_7_mask_h_0:%08x \n", ge_7_mask_h_0);
-                                                                      }
-                                                                    #endif
+      // uint32_t look_up_h_0 = h[0];
+      bitwise_workspace =
+          (bitwise_workspace | (bitwise_workspace >> 4)) & 0x00FF00FF;
+      bitwise_workspace = (bitwise_workspace | (bitwise_workspace >> 8)) &
+                          0x00007777;  // make h[i] into 0x0000vvvv
 
-          // uint32_t look_up_h_0 = h[0];
-          bitwise_workspace = (bitwise_workspace | (bitwise_workspace >> 4)) & 0x00FF00FF;
-          bitwise_workspace = (bitwise_workspace | (bitwise_workspace >> 8)) & 0x00007777; // make h[i] into 0x0000vvvv
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### h0 0x0000vvvv:%08x \n", bitwise_workspace);
+      }
+#endif
 
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### h0 0x0000vvvv:%08x \n", bitwise_workspace);
-                                                                      }
-                                                                    #endif
-
-          uint32_t result_1=0;
-          asm volatile("prmt.b32 %0,%1,%2,%3;\n"
-                 : "=r"(result_1)
-                 : "r"(reg_look_up_table[0]), "r"(reg_look_up_table[1]), "r"(bitwise_workspace));
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### result_1:%08x \n", result_1);
-                                                                      }
-                                                                    #endif
-          result_reg[0] = result_1 & (~ge_7_mask_h_0);
-          result_1 = ((~result_1)) & (ge_7_mask_h_0);
-          // result_1 = ((~result_1) | 0x01010101) & (ge_7_mask_h_0); // half mirror
-          result_reg[0] = result_reg[0] | result_1;
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### result[0]:%08x \n", result_reg[0]);
-                                                                      }
-                                                                    #endif
-          bitwise_workspace = i4s >> 4;
-          bitwise_workspace = bitwise_workspace & down_int4_mask;
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### h1 0x0v0v0v0v:%08x \n", bitwise_workspace);
-                                                                      }
-                                                                    #endif
-          ge_7_mask_h_0 = (bitwise_workspace | bitwise_workspace << 4)  & 0x88888888;
-          ge_7_mask_h_0 |= ge_7_mask_h_0 >> 1;
-          ge_7_mask_h_0 |= ge_7_mask_h_0 >> 2;
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### ge_7_mask_h_0:%08x \n", ge_7_mask_h_0);
-                                                                      }
-                                                                    #endif
-          bitwise_workspace = (bitwise_workspace | (bitwise_workspace >> 4)) & 0x00FF00FF;
-          bitwise_workspace = (bitwise_workspace | (bitwise_workspace >> 8)) & 0x00007777; // make look_up_h_0 into 0x0000vvvv
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### h1 0x0000vvvv:%08x \n", bitwise_workspace);
-                                                                      }
-                                                                    #endif
-          asm volatile("prmt.b32 %0,%1,%2,%3;\n"
-                 : "=r"(result_1)
-                 : "r"(reg_look_up_table[0]), "r"(reg_look_up_table[1]), "r"(bitwise_workspace));
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### result_1:%08x \n", result_1);
-                                                                      }
-                                                                    #endif
-          result_reg[1] = result_1 & (~ge_7_mask_h_0);
-          result_1 = ((~result_1)) & (ge_7_mask_h_0);
-          // result_1 = ((~result_1) | 0x01010101) & (ge_7_mask_h_0); // half mirror
-          result_reg[1] = result_reg[1] | result_1;
-                                                                    #ifdef NF4_LUT_DEBUG
-                                                                      if(threadIdx.x==0 && blockIdx.x==0 && blockIdx.y==0 && blockIdx.z==0){
-                                                                        printf("#### result[1]:%08x \n", result_reg[1]);
-                                                                      }
-                                                                    #endif
-        }
-        return result;
+      uint32_t result_1 = 0;
+      asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                   : "=r"(result_1)
+                   : "r"(reg_look_up_table[0]),
+                     "r"(reg_look_up_table[1]),
+                     "r"(bitwise_workspace));
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### result_1:%08x \n", result_1);
+      }
+#endif
+      result_reg[0] = result_1 & (~ge_7_mask_h_0);
+      result_1 = ((~result_1)) & (ge_7_mask_h_0);
+      // result_1 = ((~result_1) | 0x01010101) & (ge_7_mask_h_0); // half mirror
+      result_reg[0] = result_reg[0] | result_1;
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### result[0]:%08x \n", result_reg[0]);
+      }
+#endif
+      bitwise_workspace = i4s >> 4;
+      bitwise_workspace = bitwise_workspace & down_int4_mask;
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### h1 0x0v0v0v0v:%08x \n", bitwise_workspace);
+      }
+#endif
+      ge_7_mask_h_0 = (bitwise_workspace | bitwise_workspace << 4) & 0x88888888;
+      ge_7_mask_h_0 |= ge_7_mask_h_0 >> 1;
+      ge_7_mask_h_0 |= ge_7_mask_h_0 >> 2;
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### ge_7_mask_h_0:%08x \n", ge_7_mask_h_0);
+      }
+#endif
+      bitwise_workspace =
+          (bitwise_workspace | (bitwise_workspace >> 4)) & 0x00FF00FF;
+      bitwise_workspace = (bitwise_workspace | (bitwise_workspace >> 8)) &
+                          0x00007777;  // make look_up_h_0 into 0x0000vvvv
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### h1 0x0000vvvv:%08x \n", bitwise_workspace);
+      }
+#endif
+      asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+                   : "=r"(result_1)
+                   : "r"(reg_look_up_table[0]),
+                     "r"(reg_look_up_table[1]),
+                     "r"(bitwise_workspace));
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### result_1:%08x \n", result_1);
+      }
+#endif
+      result_reg[1] = result_1 & (~ge_7_mask_h_0);
+      result_1 = ((~result_1)) & (ge_7_mask_h_0);
+      // result_1 = ((~result_1) | 0x01010101) & (ge_7_mask_h_0); // half mirror
+      result_reg[1] = result_reg[1] | result_1;
+#ifdef NF4_LUT_DEBUG
+      if (threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 &&
+          blockIdx.z == 0) {
+        printf("#### result[1]:%08x \n", result_reg[1]);
+      }
+#endif
     }
+    return result;
+  }
 
-
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s)
-    {
-        return convert(s);
-    }
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s, uint32_t shared_look_up_table[32][32], int32_t warp_idx, int32_t lane_idx)
-    {
-        return convert(s,shared_look_up_table,warp_idx, lane_idx);
-    }
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s, uint32_t shared_look_up_table[16])
-    {
-        return convert(s,shared_look_up_table);
-    }
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s, int32_t* shared_look_up_table)
-    {
-        return convert(s,shared_look_up_table);
-    }
-    CUTLASS_DEVICE
-    result_type operator()(source_type const& s, cutlass::Array<uint32_t, 4, true> const& reg_look_up_table)
-    {
-        return convert(s,reg_look_up_table);
-    }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s) { return convert(s); }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s,
+                         uint32_t shared_look_up_table[32][32],
+                         int32_t warp_idx,
+                         int32_t lane_idx) {
+    return convert(s, shared_look_up_table, warp_idx, lane_idx);
+  }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s,
+                         uint32_t shared_look_up_table[16]) {
+    return convert(s, shared_look_up_table);
+  }
+  CUTLASS_DEVICE
+  result_type operator()(source_type const& s, int32_t* shared_look_up_table) {
+    return convert(s, shared_look_up_table);
+  }
+  CUTLASS_DEVICE
+  result_type operator()(
+      source_type const& s,
+      cutlass::Array<uint32_t, 4, true> const& reg_look_up_table) {
+    return convert(s, reg_look_up_table);
+  }
 };
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/tile_interleaved_layout.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/tile_interleaved_layout.h
index 89152360f1..40265273e0 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/tile_interleaved_layout.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_extensions/tile_interleaved_layout.h
@@ -40,20 +40,20 @@ limitations under the License. */
 namespace cutlass {
 namespace layout {
 
-template<int RowsPerTile, int ColumnsInterleaved>
+template <int RowsPerTile, int ColumnsInterleaved>
 class ColumnMajorTileInterleave {
-    static constexpr int kRowsPerTile        = RowsPerTile;
-    static constexpr int kColumnsInterleaved = ColumnsInterleaved;
+  static constexpr int kRowsPerTile = RowsPerTile;
+  static constexpr int kColumnsInterleaved = ColumnsInterleaved;
 };
 
-template<class T>
+template <class T>
 struct IsColumnMajorTileInterleave {
-    static constexpr bool value = false;
+  static constexpr bool value = false;
 };
 
-template<int U, int V>
+template <int U, int V>
 struct IsColumnMajorTileInterleave<ColumnMajorTileInterleave<U, V>> {
-    static constexpr bool value = true;
+  static constexpr bool value = true;
 };
 
 }  // namespace layout
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_heuristic_w4a4.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_heuristic_w4a4.h
index 096a857de6..5fa4d7c98d 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_heuristic_w4a4.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/cutlass_heuristic_w4a4.h
@@ -33,7 +33,6 @@ limitations under the License. */
 #include "glog/logging.h"
 #include "w4a4_gemm_configs.h"
 
-
 static TileShape get_cta_shape_for_config_w4a4(CutlassTileConfig tile_config) {
   switch (tile_config) {
     case CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64:
@@ -58,12 +57,12 @@ static TileShape get_cta_shape_for_config_w4a4(CutlassTileConfig tile_config) {
 }
 
 static bool is_valid_split_k_factor_w4a4(const int64_t m,
-                                    const int64_t n,
-                                    const int64_t k,
-                                    const TileShape tile_shape,
-                                    const int split_k_factor,
-                                    const size_t workspace_bytes,
-                                    const bool is_weight_only) {
+                                         const int64_t n,
+                                         const int64_t k,
+                                         const TileShape tile_shape,
+                                         const int split_k_factor,
+                                         const size_t workspace_bytes,
+                                         const bool is_weight_only) {
   // All tile sizes have a k_tile of 64.
   static constexpr int k_tile = 64;
 
@@ -125,7 +124,7 @@ static std::vector<CutlassTileConfig> get_candidate_tiles_w4a4(
       CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64
       // CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64,
       // CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64,
-      };
+  };
   std::vector<CutlassTileConfig> quant_B_configs;
   switch (sm) {
     case 90:
@@ -148,7 +147,6 @@ static std::vector<CutlassTileConfig> get_candidate_tiles_w4a4(
   return simt_configs_only ? simt_configs : allowed_configs;
 }
 
-
 static std::vector<CutlassGemmConfig> get_candidate_configs_nf4(
     int sm,
     const bool is_weight_only,
@@ -171,7 +169,6 @@ static std::vector<CutlassGemmConfig> get_candidate_configs_nf4(
   return candidate_configs;
 }
 
-
 static CutlassGemmConfig estimate_best_config_from_occupancies_w4a4(
     const std::vector<CutlassGemmConfig>& candidate_configs,
     const std::vector<int>& occupancies,
@@ -189,7 +186,7 @@ static CutlassGemmConfig estimate_best_config_from_occupancies_w4a4(
         "candidate configs vectors must have equal length.");
   }
 
-  VLOG(1)<<"estimate_best_config_from_occupancies_w4a4";
+  VLOG(1) << "estimate_best_config_from_occupancies_w4a4";
   CutlassGemmConfig best_config;
   // Score will be [0, 1]. The objective is to minimize this score.
   // It represents the fraction of SM resources unused in the last wave.
@@ -198,25 +195,25 @@ static CutlassGemmConfig estimate_best_config_from_occupancies_w4a4(
   int current_m_tile = 0;
 
   {
-    VLOG(1)<<"######## begin of cutlass gemm search";
-  if (m >= 256 &&
-      std::find_if(
-          candidate_configs.begin(),
-          candidate_configs.end(),
-          [](const CutlassGemmConfig& gemm_config) {
-            return gemm_config.tile_config ==
-                   CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64;
-          }) != candidate_configs.end()) {
-    VLOG(1) << "m >= 256, encoder config";
-    best_config = CutlassGemmConfig{
-        CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64,
-        SplitKStyle::SPLIT_K_STREAM,
-        // SplitKStyle::NO_SPLIT_K,
-        1,
-        3};
+    VLOG(1) << "######## begin of cutlass gemm search";
+    if (m >= 256 &&
+        std::find_if(
+            candidate_configs.begin(),
+            candidate_configs.end(),
+            [](const CutlassGemmConfig& gemm_config) {
+              return gemm_config.tile_config ==
+                     CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64;
+            }) != candidate_configs.end()) {
+      VLOG(1) << "m >= 256, encoder config";
+      best_config = CutlassGemmConfig{
+          CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64,
+          SplitKStyle::SPLIT_K_STREAM,
+          // SplitKStyle::NO_SPLIT_K,
+          1,
+          3};
 
-  } else {
-    VLOG(1) << "m <= 64 , decoder config";
+    } else {
+      VLOG(1) << "m <= 64 , decoder config";
       const int max_split_k =
           n >= multi_processor_count * 256 ? 1 : split_k_limit;
       for (int ii = 0; ii < candidate_configs.size(); ++ii) {
@@ -240,14 +237,14 @@ static CutlassGemmConfig estimate_best_config_from_occupancies_w4a4(
         const int ctas_in_n_dim = (n + tile_shape.n - 1) / tile_shape.n;
 
         for (int split_k_factor = 1; split_k_factor <= max_split_k;
-            ++split_k_factor) {
+             ++split_k_factor) {
           if (is_valid_split_k_factor_w4a4(m,
-                                      n,
-                                      k,
-                                      tile_shape,
-                                      split_k_factor,
-                                      workspace_bytes,
-                                      is_weight_only)) {
+                                           n,
+                                           k,
+                                           tile_shape,
+                                           split_k_factor,
+                                           workspace_bytes,
+                                           is_weight_only)) {
             const int ctas_per_wave = occupancy * multi_processor_count;
             const int ctas_for_problem =
                 ctas_in_m_dim * ctas_in_n_dim * split_k_factor;
@@ -262,7 +259,7 @@ static CutlassGemmConfig estimate_best_config_from_occupancies_w4a4(
             const float score_slack = 0.1f;
             if (current_score < config_score ||
                 ((config_waves > num_waves_total) &&
-                (current_score < config_score + score_slack))) {
+                 (current_score < config_score + score_slack))) {
               config_score = current_score;
               config_waves = num_waves_total;
               SplitKStyle split_style = split_k_factor > 1
@@ -273,9 +270,10 @@ static CutlassGemmConfig estimate_best_config_from_occupancies_w4a4(
                                               split_k_factor,
                                               candidate_config.stages};
               current_m_tile = tile_shape.m;
-              // std::cout<<"#### split-k factor: "<<split_k_factor<<" config: "<<static_cast<std::underlying_type<CutlassTileConfig>::type>(candidate_config.tile_config)<<std::endl;
+              // std::cout<<"#### split-k factor: "<<split_k_factor<<" config:
+              // "<<static_cast<std::underlying_type<CutlassTileConfig>::type>(candidate_config.tile_config)<<std::endl;
             } else if (current_score == config_score &&
-                      (best_config.stages < candidate_config.stages ||
+                       (best_config.stages < candidate_config.stages ||
                         split_k_factor < best_config.split_k_factor ||
                         current_m_tile < tile_shape.m)) {
               // Prefer deeper pipeline or smaller split-k
@@ -298,6 +296,9 @@ static CutlassGemmConfig estimate_best_config_from_occupancies_w4a4(
     throw std::runtime_error(
         "[FT Error] Heurisitc failed to find a valid config.");
   }
-  VLOG(1)<<"#### best split-k factor: "<<best_config.split_k_factor<<" config: "<<static_cast<std::underlying_type<CutlassTileConfig>::type>(best_config.tile_config);
+  VLOG(1) << "#### best split-k factor: " << best_config.split_k_factor
+          << " config: "
+          << static_cast<std::underlying_type<CutlassTileConfig>::type>(
+                 best_config.tile_config);
   return best_config;
 }
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a4_gemm_configs.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a4_gemm_configs.h
index 790ec3b17b..00c338ba97 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a4_gemm_configs.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a4_gemm_configs.h
@@ -33,44 +33,44 @@ limitations under the License. */
 // Note: The shapes are in the format MxNxK. The K shape of the runtime config
 enum class CutlassTileConfig {
   // Signals that we should run heuristics do choose a config
-  Undefined, // 0
+  Undefined,  // 0
 
   // Signals that we should run heuristics do choose a config
-  ChooseWithHeuristic, // 1
+  ChooseWithHeuristic,  // 1
 
   // SiMT config
-  CtaShape128x128x8_WarpShape64x64x8, // 2
+  CtaShape128x128x8_WarpShape64x64x8,  // 2
 
   // TensorCore configs CTA_N = 128, CTA_K = 64
   // Warp configs for M=16
-  CtaShape16x128x64_WarpShape16x32x64, // 3
-  CtaShape16x256x64_WarpShape16x64x64, // 4
+  CtaShape16x128x64_WarpShape16x32x64,  // 3
+  CtaShape16x256x64_WarpShape16x64x64,  // 4
 
   // Warp configs for M=32
-  CtaShape32x128x64_WarpShape32x32x64, // 5
+  CtaShape32x128x64_WarpShape32x32x64,  // 5
 
   // Warp configs for M=64
-  CtaShape64x128x64_WarpShape32x64x64, // 6
-  CtaShape64x128x64_WarpShape64x32x64, // 7
+  CtaShape64x128x64_WarpShape32x64x64,  // 6
+  CtaShape64x128x64_WarpShape64x32x64,  // 7
 
   // Warp configs for M=128
-  CtaShape128x128x64_WarpShape64x32x64, // 8
-  CtaShape128x128x64_WarpShape128x32x64, // 9
+  CtaShape128x128x64_WarpShape64x32x64,   // 8
+  CtaShape128x128x64_WarpShape128x32x64,  // 9
 
   // configs for large M in encoder
-  CtaShape128x256x64_WarpShape64x64x64, // 10
-  CtaShape256x128x64_WarpShape64x64x64, // 11
+  CtaShape128x256x64_WarpShape64x64x64,  // 10
+  CtaShape256x128x64_WarpShape64x64x64,  // 11
 
-  CtaShape32x256x64_WarpShape32x64x64,  // 12
-  CtaShape64x256x64_WarpShape64x64x64, // 13
-  CtaShape128x256x64_WarpShape128x64x64, // 14
-  CtaShape32x512x64_WarpShape32x128x64, // 15
+  CtaShape32x256x64_WarpShape32x64x64,    // 12
+  CtaShape64x256x64_WarpShape64x64x64,    // 13
+  CtaShape128x256x64_WarpShape128x64x64,  // 14
+  CtaShape32x512x64_WarpShape32x128x64,   // 15
 };
 
 enum class SplitKStyle {
-  NO_SPLIT_K, //0
-  SPLIT_K_SERIAL, //1
-  SPLIT_K_STREAM, //2
+  NO_SPLIT_K,      // 0
+  SPLIT_K_SERIAL,  // 1
+  SPLIT_K_STREAM,  // 2
   // SPLIT_K_PARALLEL // Not supported yet
 };
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_gemm_grouped.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_gemm_grouped.h
index 1301cc351c..ae90627fc7 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_gemm_grouped.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_gemm_grouped.h
@@ -1,12 +1,12 @@
 /***************************************************************************************************
- * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
+ * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights
+ *reserved. SPDX-License-Identifier: BSD-3-Clause
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
  *
- * 1. Redistributions of source code must retain the above copyright notice, this
- * list of conditions and the following disclaimer.
+ * 1. Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
  *
  * 2. Redistributions in binary form must reproduce the above copyright notice,
  * this list of conditions and the following disclaimer in the documentation
@@ -18,22 +18,23 @@
  *
  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ *POSSIBILITY OF SUCH DAMAGE.
  *
  **************************************************************************************************/
 /*!
   \file
-  \brief The universal GEMM accommodates streamk, batched strided, and batched array variants.
+  \brief The universal GEMM accommodates streamk, batched strided, and batched
+  array variants.
 */
 
-
 #pragma once
 
 #include <limits>
@@ -59,11 +60,9 @@ namespace device {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-
 template <typename GemmKernel_>
 class W4A8MoeGemmUniversalBase {
-public:
-
+ public:
   using GemmKernel = GemmKernel_;
   using ThreadblockShape = typename GemmKernel::Mma::Shape;
 
@@ -92,8 +91,7 @@ public:
   /// Argument structure
   using Arguments = typename GemmKernel::Arguments;
 
-protected:
-
+ protected:
   //
   // Device properties (uniform across all instances of the current thread)
   //
@@ -112,8 +110,7 @@ protected:
 
   /// Initialize static thread-local members for the thread's current device,
   /// if necessary.
-  static Status init_device_props()
-  {
+  static Status init_device_props() {
     CUTLASS_TRACE_HOST("W4A8MoeGemmUniversalBase::init_device_props()");
 
     cudaError_t cudart_result;
@@ -122,7 +119,8 @@ protected:
     int current_ordinal;
     cudart_result = cudaGetDevice(&current_ordinal);
     if (cudart_result != cudaSuccess) {
-      CUTLASS_TRACE_HOST("  cudaGetDevice() returned error " << cudaGetErrorString(cudart_result));
+      CUTLASS_TRACE_HOST("  cudaGetDevice() returned error "
+                         << cudaGetErrorString(cudart_result));
       return Status::kErrorInternal;
     }
 
@@ -133,64 +131,80 @@ protected:
     }
 
     // Update SM count member
-    cudart_result = cudaDeviceGetAttribute (&device_sms_, cudaDevAttrMultiProcessorCount, current_ordinal);
+    cudart_result = cudaDeviceGetAttribute(
+        &device_sms_, cudaDevAttrMultiProcessorCount, current_ordinal);
     if (cudart_result != cudaSuccess) {
-      CUTLASS_TRACE_HOST("  cudaDeviceGetAttribute() returned error " << cudaGetErrorString(cudart_result));
+      CUTLASS_TRACE_HOST("  cudaDeviceGetAttribute() returned error "
+                         << cudaGetErrorString(cudart_result));
       return Status::kErrorInternal;
     }
 
-    // Update the kernel function's shared memory configuration for the current device
+    // Update the kernel function's shared memory configuration for the current
+    // device
     smem_size_ = int(sizeof(typename GemmKernel::SharedStorage));
 
     // If requires more than 48KB: configure for extended, dynamic shared memory
-    if (smem_size_ >= (48 << 10))
-    {
-      cudart_result = cudaFuncSetAttribute(
-        Kernel2<GemmKernel>,
-        cudaFuncAttributeMaxDynamicSharedMemorySize,
-        smem_size_);
+    if (smem_size_ >= (48 << 10)) {
+      cudart_result =
+          cudaFuncSetAttribute(Kernel2<GemmKernel>,
+                               cudaFuncAttributeMaxDynamicSharedMemorySize,
+                               smem_size_);
       if (cudart_result != cudaSuccess) {
-        CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error " << cudaGetErrorString(cudart_result));
+        CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error "
+                           << cudaGetErrorString(cudart_result));
         return Status::kErrorInternal;
       }
 
-      cudart_result = cudaFuncSetAttribute(
-          Kernel2<GemmKernel>,
-          cudaFuncAttributePreferredSharedMemoryCarveout, 100); // 100% shared memory
+      cudart_result =
+          cudaFuncSetAttribute(Kernel2<GemmKernel>,
+                               cudaFuncAttributePreferredSharedMemoryCarveout,
+                               100);  // 100% shared memory
       if (cudart_result != cudaSuccess) {
-        CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error " << cudaGetErrorString(cudart_result));
+        CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error "
+                           << cudaGetErrorString(cudart_result));
         return Status::kErrorInternal;
       }
     }
 
     // Update SM occupancy member
     cudart_result = cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
-      &sm_occupancy_,
-      Kernel2<GemmKernel>,
-      GemmKernel::kThreadCount,
-      smem_size_,
-      cudaOccupancyDisableCachingOverride);
+        &sm_occupancy_,
+        Kernel2<GemmKernel>,
+        GemmKernel::kThreadCount,
+        smem_size_,
+        cudaOccupancyDisableCachingOverride);
     if (cudart_result != cudaSuccess) {
-      CUTLASS_TRACE_HOST("  cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned error " << cudaGetErrorString(cudart_result));
+      CUTLASS_TRACE_HOST(
+          "  cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned "
+          "error "
+          << cudaGetErrorString(cudart_result));
       return Status::kErrorInternal;
     }
 
     // Update device ordinal member on success
     device_ordinal_ = current_ordinal;
 
-    CUTLASS_TRACE_HOST("  "
-      "device_ordinal: (" << device_ordinal_ << "), "
-      "device_sms: (" << device_sms_ << "), "
-      "sm_occupancy: (" << sm_occupancy_ << ") "
-      "smem_size: (" << smem_size_ << ") "
-      "GemmKernel::kThreadCount: (" << GemmKernel::kThreadCount << ")");
+    CUTLASS_TRACE_HOST(
+        "  "
+        "device_ordinal: ("
+        << device_ordinal_
+        << "), "
+           "device_sms: ("
+        << device_sms_
+        << "), "
+           "sm_occupancy: ("
+        << sm_occupancy_
+        << ") "
+           "smem_size: ("
+        << smem_size_
+        << ") "
+           "GemmKernel::kThreadCount: ("
+        << GemmKernel::kThreadCount << ")");
 
     return Status::kSuccess;
   }
 
-
-protected:
-
+ protected:
   //
   // Instance data members
   //
@@ -198,10 +212,8 @@ protected:
   /// Kernel parameters
   typename GemmKernel::Params params_;
 
-
   /// Initialize params member
-  Status init_params(Arguments const &args)
-  {
+  Status init_params(Arguments const &args) {
     // Initialize static device properties, if necessary
     Status result = init_device_props();
     if (result != Status::kSuccess) {
@@ -213,15 +225,13 @@ protected:
     return Status::kSuccess;
   }
 
-public:
-
+ public:
   //---------------------------------------------------------------------------------------------
   // Stateless API
   //---------------------------------------------------------------------------------------------
 
   /// Determines whether the GEMM can execute the given problem.
-  static Status can_implement(Arguments const &args)
-  {
+  static Status can_implement(Arguments const &args) {
     CUTLASS_TRACE_HOST("W4A8MoeGemmUniversalBase::can_implement()");
     // Initialize static kernel and device properties, if necessary.
     Status result = init_device_props();
@@ -231,18 +241,15 @@ public:
     dim3 grid = get_grid_shape(args);
     // printf("--grid:%d, %d, %d\n", grid.x, grid.y, grid.z);
     if (!(grid.y <= std::numeric_limits<uint16_t>::max() &&
-          grid.z <= std::numeric_limits<uint16_t>::max()))
-    {
+          grid.z <= std::numeric_limits<uint16_t>::max())) {
       return Status::kErrorInvalidProblem;
     }
     return GemmKernel::can_implement(args);
   }
 
-
   /// Returns the workspace size (in bytes) needed for the problem
   /// geometry expressed by these arguments
-  static size_t get_workspace_size(Arguments const &args)
-  {
+  static size_t get_workspace_size(Arguments const &args) {
     CUTLASS_TRACE_HOST("W4A8MoeGemmUniversalBase::get_workspace_size()");
 
     // Initialize parameters from args
@@ -258,33 +265,28 @@ public:
     return workspace_bytes;
   }
 
-
   /// Returns the grid extents in thread blocks to launch
-  static dim3 get_grid_shape(Arguments const &args)
-  {
+  static dim3 get_grid_shape(Arguments const &args) {
     CUTLASS_TRACE_HOST("W4A8MoeGemmUniversalBase::get_grid_shape()");
 
     // Initialize parameters from args
     W4A8MoeGemmUniversalBase base;
     if (base.init_params(args) != Status::kSuccess) {
-      return dim3(0,0,0);
+      return dim3(0, 0, 0);
     }
 
     // Get dims from parameters
     dim3 grid_dims = base.params_.get_grid_dims();
 
-    CUTLASS_TRACE_HOST(
-         "  tiled_shape: " << base.params_.get_tiled_shape()  << "\n"
-      << "  grid_dims: {" << grid_dims << "}");
+    CUTLASS_TRACE_HOST("  tiled_shape: "
+                       << base.params_.get_tiled_shape() << "\n"
+                       << "  grid_dims: {" << grid_dims << "}");
 
     return grid_dims;
   }
 
-
-
   /// Returns the maximum number of active thread blocks per multiprocessor
-  static int maximum_active_blocks(int smem_capacity = -1)
-  {
+  static int maximum_active_blocks(int smem_capacity = -1) {
     CUTLASS_TRACE_HOST("W4A8MoeGemmUniversalBase::maximum_active_blocks()");
 
     int smem_size = int(sizeof(typename GemmKernel_::SharedStorage));
@@ -300,26 +302,25 @@ public:
       if (result != cudaSuccess) {
         // Call cudaGetLastError() to clear the error bit
         result = cudaGetLastError();
-        CUTLASS_TRACE_HOST(
-          "  cudaFuncSetAttribute() returned error "
-          << cudaGetErrorString(result));
+        CUTLASS_TRACE_HOST("  cudaFuncSetAttribute() returned error "
+                           << cudaGetErrorString(result));
         return -1;
       }
     }
 
     int max_active_blocks = -1;
-    result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-        &max_active_blocks,
-        Kernel2<GemmKernel_>,
-        GemmKernel_::kThreadCount,
-        smem_size);
+    result =
+        cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_active_blocks,
+                                                      Kernel2<GemmKernel_>,
+                                                      GemmKernel_::kThreadCount,
+                                                      smem_size);
 
     if (result != cudaSuccess) {
       // Call cudaGetLastError() to clear the error bit
       result = cudaGetLastError();
       CUTLASS_TRACE_HOST(
-        "  cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error "
-        << cudaGetErrorString(result));
+          "  cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error "
+          << cudaGetErrorString(result));
       return -1;
     }
 
@@ -327,19 +328,17 @@ public:
     return max_active_blocks;
   }
 
-
   //---------------------------------------------------------------------------------------------
   // Stateful API
   //---------------------------------------------------------------------------------------------
 
   /// Initializes GEMM state from arguments and workspace memory
-  Status initialize(
-    Arguments const &args,
-    void *workspace,
-    cudaStream_t stream = nullptr)
-  {
+  Status initialize(Arguments const &args,
+                    void *workspace,
+                    cudaStream_t stream = nullptr) {
     CUTLASS_TRACE_HOST("W4A8MoeGemmUniversalBase::initialize() - workspace "
-      << workspace << ", stream: " << (stream ? "non-null" : "null"));
+                       << workspace
+                       << ", stream: " << (stream ? "non-null" : "null"));
 
     // Initialize parameters from args
     Status result = init_params(args);
@@ -351,20 +350,16 @@ public:
     return params_.init_workspace(workspace, stream);
   }
 
-
-  /// Lightweight update given a subset of arguments.  Problem geometry is assumed to
-  /// remain the same.
-  Status update(Arguments const &args)
-  {
+  /// Lightweight update given a subset of arguments.  Problem geometry is
+  /// assumed to remain the same.
+  Status update(Arguments const &args) {
     CUTLASS_TRACE_HOST("W4A8MoeGemmUniversalBase()::update()");
     params_.update(args);
     return Status::kSuccess;
   }
 
-
   /// Runs the kernel using initialized state.
-  Status run(cudaStream_t stream = nullptr)
-  {
+  Status run(cudaStream_t stream = nullptr) {
     CUTLASS_TRACE_HOST("W4A8MoeGemmUniversalBase::run()");
 
     // Configure grid and block dimensions
@@ -372,37 +367,37 @@ public:
     dim3 grid(params_.threadblock_count, 1, 1);
 
     // Launch kernel
-    CUTLASS_TRACE_HOST("  "
-      "grid: (" << grid << "), "
-      "block: (" << block << "), "
-      "SMEM: (" << smem_size_ << ")");
+    CUTLASS_TRACE_HOST(
+        "  "
+        "grid: ("
+        << grid
+        << "), "
+           "block: ("
+        << block
+        << "), "
+           "SMEM: ("
+        << smem_size_ << ")");
 
     Kernel2<GemmKernel><<<grid, block, smem_size_, stream>>>(params_);
 
     // Query for errors
     cudaError_t result = cudaGetLastError();
     if (result != cudaSuccess) {
-      CUTLASS_TRACE_HOST("  grid launch failed with error " << cudaGetErrorString(result));
+      CUTLASS_TRACE_HOST("  grid launch failed with error "
+                         << cudaGetErrorString(result));
       return Status::kErrorInternal;
     }
 
     return Status::kSuccess;
   }
 
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
 
   /// Runs the kernel using initialized state.
-  Status operator()(cudaStream_t stream = nullptr)
-  {
-    return run(stream);
-  }
-
-
-  /// Runs the kernel using initialized state.
-  Status operator()(
-    Arguments const &args,
-    void *workspace = nullptr,
-    cudaStream_t stream = nullptr)
-  {
+  Status operator()(Arguments const &args,
+                    void *workspace = nullptr,
+                    cudaStream_t stream = nullptr) {
     Status status = initialize(args, workspace, stream);
 
     if (status == Status::kSuccess) {
@@ -413,7 +408,6 @@ public:
   }
 };
 
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 /// Static initializers
 /////////////////////////////////////////////////////////////////////////////////////////////////
@@ -434,12 +428,10 @@ thread_local int W4A8MoeGemmUniversalBase<GemmKernel_>::sm_occupancy_ = -1;
 template <typename GemmKernel_>
 thread_local int W4A8MoeGemmUniversalBase<GemmKernel_>::smem_size_ = -1;
 
-
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-} // namespace device
-} // namespace gemm
-} // namespace cutlass
+}  // namespace device
+}  // namespace gemm
+}  // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_cutlass_kernel.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_cutlass_kernel.h
index dbcc8912f2..18cdeca89c 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_cutlass_kernel.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_cutlass_kernel.h
@@ -37,60 +37,65 @@ namespace kernel {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-template<typename Mma_,                 ///! Threadblock-scoped matrix multiply-accumulate
-         typename Epilogue_,            ///! Epilogue
-         typename ThreadblockSwizzle_,  ///! Threadblock swizzling function
-         typename KernelArch,  ///! The Architecture this kernel is compiled for. Used since SIMT kernels lose top-level
-                               /// arch.
-         bool SplitKSerial,     ///! If true, code supporting split-K via serial reduction is enabled.
-    GroupScheduleMode GroupScheduleMode_  ///! Type of scheduling to //
-         >
+template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
+          typename Epilogue_,            ///! Epilogue
+          typename ThreadblockSwizzle_,  ///! Threadblock swizzling function
+          typename KernelArch,  ///! The Architecture this kernel is compiled
+                                /// for. Used since SIMT kernels lose top-level
+                                /// arch.
+          bool SplitKSerial,    ///! If true, code supporting split-K via serial
+                                /// reduction is enabled.
+          GroupScheduleMode GroupScheduleMode_  ///! Type of scheduling to //
+          >
 struct MoeW4A8Gemm {
-  using Mma                       = Mma_;
-  using Epilogue                  = Epilogue_;
-  using EpilogueOutputOp          = typename Epilogue::OutputOp;
-  using ThreadblockSwizzle        = ThreadblockSwizzle_;
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using EpilogueOutputOp = typename Epilogue::OutputOp;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
   static bool const kSplitKSerial = SplitKSerial;
   static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_;
   static bool const kTransposed = false;
 
-  using ElementA     = typename Mma::IteratorA::Element;
-  using LayoutA      = typename Mma::IteratorA::Layout;
+  using ElementA = typename Mma::IteratorA::Element;
+  using LayoutA = typename Mma::IteratorA::Layout;
   using TensorRefA = TensorRef<ElementA, LayoutA>;
-  using ElementB     = typename Mma::IteratorB::Element;
-  using LayoutB      = typename Mma::IteratorB::Layout;
+  using ElementB = typename Mma::IteratorB::Element;
+  using LayoutB = typename Mma::IteratorB::Layout;
   using TensorRefB = TensorRef<ElementB, LayoutB>;
 
   using LayoutAlphaCol = cutlass::layout::RowMajor;
   using LayoutAlphaRow = cutlass::layout::ColumnMajor;
   using TensorRefNf4LookUpTable = TensorRef<int32_t, cutlass::layout::RowMajor>;
 
-  using ElementC     = typename Epilogue::OutputTileIterator::Element;
-  using LayoutC      = typename Mma::LayoutC;
+  using ElementC = typename Epilogue::OutputTileIterator::Element;
+  using LayoutC = typename Mma::LayoutC;
   using TensorRefC = TensorRef<ElementC, LayoutC>;
 
   static ComplexTransform const kTransformA = Mma::kTransformA;
   static ComplexTransform const kTransformB = Mma::kTransformA;
 
   // Type definitions about the mainloop.
-  using Operator         = typename Mma::Operator;
-  using OperatorClass    = typename Mma::Operator::OperatorClass;
+  using Operator = typename Mma::Operator;
+  using OperatorClass = typename Mma::Operator::OperatorClass;
   using ThreadblockShape = typename Mma::Shape;
-  using WarpShape        = typename Mma::Operator::Shape;
+  using WarpShape = typename Mma::Operator::Shape;
   using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
-  using ArchTag          = typename Mma::ArchTag;
+  using ArchTag = typename Mma::ArchTag;
 
-  static int const kStages     = Mma::kStages;
+  static int const kStages = Mma::kStages;
   static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
   static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
-  static int const kAlignmentD = Epilogue::OutputTileIterator::kElementsPerAccess;
+  static int const kAlignmentD =
+      Epilogue::OutputTileIterator::kElementsPerAccess;
 
   /// Warp count (concept: GemmShape)
-  using WarpCount               = typename Mma::WarpCount;
+  using WarpCount = typename Mma::WarpCount;
   static int const kThreadCount = 32 * WarpCount::kCount;
-  static int const kSplitKAlignment = const_max(128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value);
+  static int const kSplitKAlignment = const_max(
+      128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value);
 
-  static constexpr int kInterleave = Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
+  static constexpr int kInterleave =
+      Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
 
   using ProblemVisitor = GemmMoeProblemVisitor<ThreadblockShape,
                                                kGroupScheduleMode,
@@ -98,7 +103,6 @@ struct MoeW4A8Gemm {
                                                kThreadCount,
                                                kTransposed>;
 
-
   /// Argument structure
   struct Arguments : UniversalArgumentsBase {
     //
@@ -133,8 +137,7 @@ struct MoeW4A8Gemm {
     Arguments() {}
 
     /// constructs an arguments structure
-    Arguments(
-        cutlass::gemm::GemmUniversalMode mode_,
+    Arguments(cutlass::gemm::GemmUniversalMode mode_,
               GemmCoord problem_size_,
               int problem_count,
               int batch_count_,
@@ -170,7 +173,7 @@ struct MoeW4A8Gemm {
           host_problem_sizes(nullptr) {}
   };
 
-    /// Parameters structure
+  /// Parameters structure
   struct Params : UniversalParamsBase<ThreadblockSwizzle,
                                       ThreadblockShape,
                                       ElementA,
@@ -219,120 +222,111 @@ struct MoeW4A8Gemm {
           ptr_C(args.ref_C.data()),
           ptr_D(args.ref_D.data()),
           batch_stride_A(args.batch_stride_A),
-          batch_stride_B(args.batch_stride_B)
-          {
-    }
+          batch_stride_B(args.batch_stride_B) {}
   };
 
-    /// Shared memory storage structure
-    union SharedStorage {
-        typename ProblemVisitor::SharedStorage problem_visitor;
-        typename Mma::SharedStorage      main_loop;
-        struct {
-        typename Epilogue::SharedStorage epilogue;
-        } epilogue;
-    };
+  /// Shared memory storage structure
+  union SharedStorage {
+    typename ProblemVisitor::SharedStorage problem_visitor;
+    typename Mma::SharedStorage main_loop;
+    struct {
+      typename Epilogue::SharedStorage epilogue;
+    } epilogue;
+  };
 
-    //
-    // Methods
-    //
+  //
+  // Methods
+  //
 
-    CUTLASS_HOST_DEVICE
-    MoeW4A8Gemm() {}
+  CUTLASS_HOST_DEVICE
+  MoeW4A8Gemm() {}
 
-    /// Determines whether kernel satisfies alignment
-    CUTLASS_HOST_DEVICE
-    static Status can_implement(cutlass::gemm::GemmCoord const& problem_size)
-    {
+  /// Determines whether kernel satisfies alignment
+  CUTLASS_HOST_DEVICE
+  static Status can_implement(cutlass::gemm::GemmCoord const& problem_size) {
+    CUTLASS_TRACE_HOST(
+        "GemmWithEpilogueVisitorInterleavedNf4::can_implement()");
 
-        CUTLASS_TRACE_HOST("GemmWithEpilogueVisitorInterleavedNf4::can_implement()");
+    static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+    static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
 
-        static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
-        static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+    bool isAMisaligned = false;
+    bool isBMisaligned = false;
+    bool isCMisaligned = false;
 
-        bool isAMisaligned = false;
-        bool isBMisaligned = false;
-        bool isCMisaligned = false;
-
-        if (platform::is_same<LayoutA, layout::RowMajor>::value) {
-        isAMisaligned = problem_size.k() % kAlignmentA;
-        } else if (platform::is_same<LayoutA, layout::ColumnMajor>::value) {
-        isAMisaligned = problem_size.m() % kAlignmentA;
-        } else if (platform::is_same<LayoutA,
-                                    layout::ColumnMajorInterleaved<32>>::value ||
-                platform::is_same<LayoutA,
-                                    layout::ColumnMajorInterleaved<64>>::value) {
-        isAMisaligned = problem_size.k() % kAlignmentA;
-        }
-
-        if (platform::is_same<LayoutB, layout::RowMajor>::value) {
-        isBMisaligned = problem_size.n() % kAlignmentB;
-        } else if (platform::is_same<LayoutB, layout::ColumnMajor>::value) {
-        isBMisaligned = problem_size.k() % kAlignmentB;
-        } else if (platform::is_same<LayoutB,
-                                    layout::RowMajorInterleaved<32>>::value ||
-                platform::is_same<LayoutB,
-                                    layout::RowMajorInterleaved<64>>::value) {
-        isBMisaligned = problem_size.k() % kAlignmentB;
-        }
-
-        if (platform::is_same<LayoutC, layout::RowMajor>::value) {
-        // isCMisaligned = problem_size.n() % kAlignmentC;
-        } else if (platform::is_same<LayoutC, layout::ColumnMajor>::value) {
-        // isCMisaligned = problem_size.m() % kAlignmentC;
-        } else if (platform::is_same<LayoutC,
-                                    layout::ColumnMajorInterleaved<32>>::value ||
-                platform::is_same<LayoutC,
-                                    layout::ColumnMajorInterleaved<64>>::value) {
-        // isCMisaligned = problem_size.n() % kAlignmentC;
-        }
-
-        if (isAMisaligned) {
-        CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for A operand");
-        return Status::kErrorMisalignedOperand;
-        }
-
-        if (isBMisaligned) {
-        CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for B operand");
-        return Status::kErrorMisalignedOperand;
-        }
-
-        if (isCMisaligned) {
-        CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for C operand");
-        return Status::kErrorMisalignedOperand;
-        }
-
-        CUTLASS_TRACE_HOST("  returning kSuccess");
-
-        return Status::kSuccess;
+    if (platform::is_same<LayoutA, layout::RowMajor>::value) {
+      isAMisaligned = problem_size.k() % kAlignmentA;
+    } else if (platform::is_same<LayoutA, layout::ColumnMajor>::value) {
+      isAMisaligned = problem_size.m() % kAlignmentA;
+    } else if (platform::is_same<LayoutA,
+                                 layout::ColumnMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutA,
+                                 layout::ColumnMajorInterleaved<64>>::value) {
+      isAMisaligned = problem_size.k() % kAlignmentA;
     }
 
+    if (platform::is_same<LayoutB, layout::RowMajor>::value) {
+      isBMisaligned = problem_size.n() % kAlignmentB;
+    } else if (platform::is_same<LayoutB, layout::ColumnMajor>::value) {
+      isBMisaligned = problem_size.k() % kAlignmentB;
+    } else if (platform::is_same<LayoutB,
+                                 layout::RowMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutB,
+                                 layout::RowMajorInterleaved<64>>::value) {
+      isBMisaligned = problem_size.k() % kAlignmentB;
+    }
+
+    if (platform::is_same<LayoutC, layout::RowMajor>::value) {
+      // isCMisaligned = problem_size.n() % kAlignmentC;
+    } else if (platform::is_same<LayoutC, layout::ColumnMajor>::value) {
+      // isCMisaligned = problem_size.m() % kAlignmentC;
+    } else if (platform::is_same<LayoutC,
+                                 layout::ColumnMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutC,
+                                 layout::ColumnMajorInterleaved<64>>::value) {
+      // isCMisaligned = problem_size.n() % kAlignmentC;
+    }
+
+    if (isAMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for A operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (isBMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for B operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (isCMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for C operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    CUTLASS_TRACE_HOST("  returning kSuccess");
+
+    return Status::kSuccess;
+  }
+
   static Status can_implement(Arguments const& args) {
     return can_implement(args.problem_size);
   }
 
-    static size_t get_extra_workspace_size(Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape)
-    {
+  static size_t get_extra_workspace_size(
+      Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape) {
+    return 0;
+  }
 
-        return 0;
-    }
-
-
-    CUTLASS_DEVICE
-    static void invoke(
-        Params const &params,
-        SharedStorage &shared_storage)
-    {
-        MoeW4A8Gemm op;
-        op(params, shared_storage);
-    }
+  CUTLASS_DEVICE
+  static void invoke(Params const& params, SharedStorage& shared_storage) {
+    MoeW4A8Gemm op;
+    op(params, shared_storage);
+  }
 
 #define SPLIT_K_ENABLED 1
 
   /// Executes one GEMM
   CUTLASS_DEVICE
   void operator()(Params const& params, SharedStorage& shared_storage) {
-
     using ElementA = typename Mma::IteratorA::Element;
     using LayoutA = typename Mma::IteratorA::Layout;
     using ElementB = typename Mma::IteratorB::Element;
@@ -340,12 +334,11 @@ struct MoeW4A8Gemm {
 
     static constexpr int kInterleave =
         Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
-    static_assert(
-        platform::is_same<LayoutB, layout::RowMajor>::value &&
-                kInterleave == 1 ||
-            platform::is_same<LayoutB, layout::ColumnMajor>::value &&
-                kInterleave >= 1,
-        "B must be row major/col major OR col major interleaved.");
+    static_assert(platform::is_same<LayoutB, layout::RowMajor>::value &&
+                          kInterleave == 1 ||
+                      platform::is_same<LayoutB, layout::ColumnMajor>::value &&
+                          kInterleave >= 1,
+                  "B must be row major/col major OR col major interleaved.");
 
     //
     // Problem visitor.
@@ -357,190 +350,187 @@ struct MoeW4A8Gemm {
     int64_t bytes_per_expert_matrix =
         (gemm_k * gemm_n / 8) * cutlass::sizeof_bits<ElementB>::value;
 
-      // Outer 'persistent' loop to iterate over tiles
-      while (problem_visitor.next_tile()) {
+    // Outer 'persistent' loop to iterate over tiles
+    while (problem_visitor.next_tile()) {
       // // Compute threadblock location
       ThreadblockSwizzle threadblock_swizzle;
 
       cutlass::gemm::GemmCoord threadblock_tile_offset =
           threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-        GemmCoord problem_size = problem_visitor.problem_size();
-        int32_t problem_idx = problem_visitor.problem_index();
-        int32_t cta_idx = int32_t(problem_visitor.threadblock_idx());
+      GemmCoord problem_size = problem_visitor.problem_size();
+      int32_t problem_idx = problem_visitor.problem_index();
+      int32_t cta_idx = int32_t(problem_visitor.threadblock_idx());
 
-        GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
+      GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
 
-        cutlass::gemm::GemmCoord threadblock_offset(
-            int(cta_idx / grid_shape.n()) * Mma::Shape::kM,  // NOLINT
-            int(cta_idx % grid_shape.n()) * Mma::Shape::kN,  // NOLINT
-            0);
+      cutlass::gemm::GemmCoord threadblock_offset(
+          int(cta_idx / grid_shape.n()) * Mma::Shape::kM,  // NOLINT
+          int(cta_idx % grid_shape.n()) * Mma::Shape::kN,  // NOLINT
+          0);
 
-        // Load element pointers. Exchange pointers and strides if working on
-        // the transpose
-        const int64_t rows_to_jump =
-            problem_idx == 0
-                ? 0
-                : params.problem_visitor.last_row_for_problem[problem_idx - 1];
-        ElementA* ptr_A =
-            reinterpret_cast<ElementA*>(params.ptr_A) + rows_to_jump * gemm_k;
-        typename LayoutA::LongIndex ldm_A = gemm_k;
+      // Load element pointers. Exchange pointers and strides if working on
+      // the transpose
+      const int64_t rows_to_jump =
+          problem_idx == 0
+              ? 0
+              : params.problem_visitor.last_row_for_problem[problem_idx - 1];
+      ElementA* ptr_A =
+          reinterpret_cast<ElementA*>(params.ptr_A) + rows_to_jump * gemm_k;
+      typename LayoutA::LongIndex ldm_A = gemm_k;
 
-        char* byte_ptr_B = ((char*)params.ptr_B) +                 // NOLINT
-                           problem_idx * bytes_per_expert_matrix;  // NOLINT
-        ElementB* ptr_B = reinterpret_cast<ElementB*>(byte_ptr_B);
-        typename LayoutB::LongIndex ldm_B =
-            platform::is_same<layout::RowMajor, LayoutB>::value
-                ? gemm_n
-                : gemm_k * kInterleave;
+      char* byte_ptr_B = ((char*)params.ptr_B) +                 // NOLINT
+                         problem_idx * bytes_per_expert_matrix;  // NOLINT
+      ElementB* ptr_B = reinterpret_cast<ElementB*>(byte_ptr_B);
+      typename LayoutB::LongIndex ldm_B =
+          platform::is_same<layout::RowMajor, LayoutB>::value
+              ? gemm_n
+              : gemm_k * kInterleave;
 
       int offset_k = 0;
       int problem_size_k = params.problem_size.k();
 
-        // Compute initial location in logical coordinates
+      // Compute initial location in logical coordinates
       cutlass::MatrixCoord tb_offset_A{
           threadblock_offset.m(),
           0,
       };
 
-      cutlass::MatrixCoord tb_offset_B{
-          0,
-          threadblock_offset.n() / kInterleave};
+      cutlass::MatrixCoord tb_offset_B{0, threadblock_offset.n() / kInterleave};
 
-        // Compute position within threadblock
-        int thread_idx = threadIdx.x;
+      // Compute position within threadblock
+      int thread_idx = threadIdx.x;
 
-        // Construct iterators to A and B operands
-        typename Mma::IteratorA iterator_A(
-            params.params_A,
-            ptr_A,
-            {params.problem_size.m(), problem_size_k},
-            thread_idx,
-            tb_offset_A);
+      // Construct iterators to A and B operands
+      typename Mma::IteratorA iterator_A(
+          params.params_A,
+          ptr_A,
+          {params.problem_size.m(), problem_size_k},
+          thread_idx,
+          tb_offset_A);
 
+      typename Mma::IteratorB iterator_B(
+          params.params_B,
+          ptr_B,
+          {problem_size_k * kInterleave, params.problem_size.n() / kInterleave},
+          thread_idx,
+          tb_offset_B);
+      typename Mma::IteratorNF4LookUpTable iterator_nf4_look_up_table =
+          Mma::IteratorNF4LookUpTable(params.params_nf4_look_up_table,
+                                      params.ref_nf4_look_up_table.data(),
+                                      {0, 16},
+                                      threadIdx.x,
+                                      {0, 0});
 
-        typename Mma::IteratorB iterator_B(
-            params.params_B,
-            ptr_B,
-            {problem_size_k * kInterleave, params.problem_size.n() / kInterleave},
-            thread_idx,
-            tb_offset_B);
-        typename Mma::IteratorNF4LookUpTable iterator_nf4_look_up_table =
-          Mma::IteratorNF4LookUpTable(
-            params.params_nf4_look_up_table,
-          params.ref_nf4_look_up_table.data(),
-            {0,16},
-            threadIdx.x,
-            {0,0}
-          );
+      // Broadcast the warp_id computed by lane 0 to ensure dependent code
+      // is compiled as warp-uniform.
+      int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
 
-        // Broadcast the warp_id computed by lane 0 to ensure dependent code
-        // is compiled as warp-uniform.
-        int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+      int lane_idx = threadIdx.x % 32;
 
-        int lane_idx = threadIdx.x % 32;
+      //
+      // Main loop
+      //
 
-        //
-        // Main loop
-        //
+      // Construct thread-scoped matrix multiply
+      Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
 
-        // Construct thread-scoped matrix multiply
-        Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
+      typename Mma::FragmentC accumulators;
 
-        typename Mma::FragmentC accumulators;
+      accumulators.clear();
 
-        accumulators.clear();
+      // Compute threadblock-scoped matrix multiply-add
+      int gemm_k_iterations =
+          (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
+      // printf("#### gemm_k_iterations: %d \n", gemm_k_iterations);
+      // Compute threadblock-scoped matrix multiply-add
+      mma(gemm_k_iterations,
+          accumulators,
+          iterator_A,
+          iterator_B,
+          iterator_nf4_look_up_table,
+          accumulators);
+      // if(threadIdx.x==0){
+      //   printf("##### block: %d-%d-%d, offset-m-n-k:%d-%d-%d \n",
+      //           blockIdx.x, blockIdx.y, blockIdx.z,
+      //           threadblock_tile_offset.m(),
+      //           threadblock_tile_offset.n(),
+      //           threadblock_tile_offset.k()
+      //           );
+      // }
+      //
+      // Masked tile iterators constructed from members
+      //
+      EpilogueOutputOp output_op(params.output_op);
 
-        // Compute threadblock-scoped matrix multiply-add
-        int gemm_k_iterations =
-            (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
-        // printf("#### gemm_k_iterations: %d \n", gemm_k_iterations);
-        // Compute threadblock-scoped matrix multiply-add
-        mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, iterator_nf4_look_up_table, accumulators);
-        // if(threadIdx.x==0){
-        //   printf("##### block: %d-%d-%d, offset-m-n-k:%d-%d-%d \n",
-        //           blockIdx.x, blockIdx.y, blockIdx.z,
-        //           threadblock_tile_offset.m(),
-        //           threadblock_tile_offset.n(),
-        //           threadblock_tile_offset.k()
-        //           );
-        // }
-        //
-        // Masked tile iterators constructed from members
-        //
-            EpilogueOutputOp output_op(params.output_op);
+      threadblock_tile_offset =
+          threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-        threadblock_tile_offset =
-            threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+      ElementC* ptr_D =
+          reinterpret_cast<ElementC*>(params.ptr_D) + rows_to_jump * gemm_n;
 
-        ElementC* ptr_D =
-            reinterpret_cast<ElementC*>(params.ptr_D) + rows_to_jump * gemm_n;
+      int block_idx = threadblock_tile_offset.m() +
+                      threadblock_tile_offset.n() * params.grid_tiled_shape.m();
 
+      // Construct the semaphore.
+      Semaphore semaphore(params.semaphore + block_idx, thread_idx);
 
-        int block_idx = threadblock_tile_offset.m() +
-                        threadblock_tile_offset.n() * params.grid_tiled_shape.m();
+      // If performing a reduction via split-K, fetch the initial
+      // synchronization
+      if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+        // Fetch the synchronization lock initially but do not block.
+        semaphore.fetch();
 
-
-
-        // Construct the semaphore.
-        Semaphore semaphore(params.semaphore + block_idx, thread_idx);
-
-        // If performing a reduction via split-K, fetch the initial synchronization
-        if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
-
-            // Fetch the synchronization lock initially but do not block.
-            semaphore.fetch();
-
-            // Indicate which position in a serial reduction the output operator is currently updating
-            output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k());
-        }
-
-            // Tile iterator writing to destination tensor.
-            typename Epilogue::OutputTileIterator iterator_D(params.params_D,
-                                                             ptr_D,
-                                                             params.problem_size.mn(),
-                                                             thread_idx,
-                                                             threadblock_offset,
-                                                             params.scatter_D_indices);
-
-
-            Epilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx);
-
-            // Wait on the semaphore - this latency may have been covered by iterator construction
-            if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
-
-                // For subsequent threadblocks, the source matrix is held in the 'D' tensor.
-
-                semaphore.wait(threadblock_tile_offset.k());
-            }
-
-            // Execute the epilogue operator to update the destination tensor.
-            epilogue(output_op, iterator_D, accumulators, iterator_D);
-
-            //
-            // Release the semaphore
-            //
-
-            if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
-
-                int lock = 0;
-                if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) {
-
-                    // The final threadblock resets the semaphore for subsequent grids.
-                    lock = 0;
-                }
-                else {
-                    // Otherwise, the semaphore is incremented
-                    lock = threadblock_tile_offset.k() + 1;
-                }
-
-                semaphore.release(lock);
-            }
-
-            // Next tile
-            shared_storage.problem_visitor.advance(gridDim.x);
+        // Indicate which position in a serial reduction the output operator is
+        // currently updating
+        output_op.set_k_partition(threadblock_tile_offset.k(),
+                                  params.grid_tiled_shape.k());
       }
 
+      // Tile iterator writing to destination tensor.
+      typename Epilogue::OutputTileIterator iterator_D(
+          params.params_D,
+          ptr_D,
+          params.problem_size.mn(),
+          thread_idx,
+          threadblock_offset,
+          params.scatter_D_indices);
+
+      Epilogue epilogue(
+          shared_storage.epilogue, thread_idx, warp_idx, lane_idx);
+
+      // Wait on the semaphore - this latency may have been covered by iterator
+      // construction
+      if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+        // For subsequent threadblocks, the source matrix is held in the 'D'
+        // tensor.
+
+        semaphore.wait(threadblock_tile_offset.k());
+      }
+
+      // Execute the epilogue operator to update the destination tensor.
+      epilogue(output_op, iterator_D, accumulators, iterator_D);
+
+      //
+      // Release the semaphore
+      //
+
+      if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+        int lock = 0;
+        if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) {
+          // The final threadblock resets the semaphore for subsequent grids.
+          lock = 0;
+        } else {
+          // Otherwise, the semaphore is incremented
+          lock = threadblock_tile_offset.k() + 1;
+        }
+
+        semaphore.release(lock);
+      }
+
+      // Next tile
+      shared_storage.problem_visitor.advance(gridDim.x);
+    }
   }
 };
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_cutlass_kernel_template.cu b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_cutlass_kernel_template.cu
index ede2178d27..cf02b86af1 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_cutlass_kernel_template.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_cutlass_kernel_template.cu
@@ -38,10 +38,9 @@
 #include "cutlass_kernels/w4a8_moe/cutlass_extensions/gemm/kernel/gemm_with_epilogue_visitor_interleaved_nf4.h"
 #include "w4a8_moe_gemm_with_epilogue_visitor.h"
 
-
 template <int val>
 class IntegerType {
-  public:
+ public:
   static constexpr int value = val;
 };
 
@@ -76,17 +75,21 @@ void generic_w4a8_moe_gemm_kernelLauncher(
     int multi_processor_count,
     cudaStream_t stream,
     int* occupancy) {
-  if (gemm_config.split_k_style == SplitKStyle::NO_SPLIT_K){
+  if (gemm_config.split_k_style == SplitKStyle::NO_SPLIT_K) {
     static_assert(cutlass::platform::is_same<IntAType, int8_t>::value,
                   "input type must be int8_t");
 
-    // The cutlass type for the input elements. This is needed to convert to cutlass::half_t if necessary.
-    // using OutputElementType_      = OutputType;
-    using OutputElementType_ = typename cutlass::platform::conditional<cutlass::platform::is_same<OutputType, __nv_bfloat16>::value,
-                                                cutlass::bfloat16_t, OutputType>::type;
+    // The cutlass type for the input elements. This is needed to convert to
+    // cutlass::half_t if necessary. using OutputElementType_      = OutputType;
+    using OutputElementType_ = typename cutlass::platform::conditional<
+        cutlass::platform::is_same<OutputType, __nv_bfloat16>::value,
+        cutlass::bfloat16_t,
+        OutputType>::type;
 
-    using OutputElementType       = typename cutlass::platform::conditional<cutlass::platform::is_same<OutputElementType_, half>::value,
-                                                cutlass::half_t, OutputElementType_>::type;
+    using OutputElementType = typename cutlass::platform::conditional<
+        cutlass::platform::is_same<OutputElementType_, half>::value,
+        cutlass::half_t,
+        OutputElementType_>::type;
 
     using CutlassIntAType_ = IntAType;
     using CutlassIntAType = CutlassIntAType_;
@@ -94,47 +97,55 @@ void generic_w4a8_moe_gemm_kernelLauncher(
     using CutlassIntBType_ = IntBType;
     using CutlassIntBType = CutlassIntBType_;
 
-    // We need separate config for each architecture since we will target different tensorcore instructions. For float,
-    // we do not target TCs.
+    // We need separate config for each architecture since we will target
+    // different tensorcore instructions. For float, we do not target TCs.
 
-    using MixedGemmArchTraits = cutlass::gemm::kernel::
-        Int8Nf4GemmArchTraits<CutlassIntAType, CutlassIntBType, OutputElementType, arch>;
-
-    using ElementAccumulator  = typename MixedGemmArchTraits::AccType;
-    using ElementCompute  = float;
+    using MixedGemmArchTraits =
+        cutlass::gemm::kernel::Int8Nf4GemmArchTraits<CutlassIntAType,
+                                                     CutlassIntBType,
+                                                     OutputElementType,
+                                                     arch>;
 
+    using ElementAccumulator = typename MixedGemmArchTraits::AccType;
+    using ElementCompute = float;
 
     // ==============
     using EpilogueOp =
-        typename Epilogue<OutputElementType, MixedGemmArchTraits::ElementsPerAccessC, ElementCompute, EpilogueTag>::Op;
+        typename Epilogue<OutputElementType,
+                          MixedGemmArchTraits::ElementsPerAccessC,
+                          ElementCompute,
+                          EpilogueTag>::Op;
 
-    using ThreadBlockSwizzle = typename cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
-    using GemmKernel_ = typename cutlass::gemm::kernel::DefaultInt8InterleavedGemm<
-        CutlassIntAType,
-        cutlass::layout::RowMajor,
-        MixedGemmArchTraits::ElementsPerAccessA,
-        CutlassIntBType,
-        typename MixedGemmArchTraits::LayoutB,
-        MixedGemmArchTraits::ElementsPerAccessB,
-        OutputElementType,
-        cutlass::layout::RowMajor,
-        ElementAccumulator,
-        cutlass::arch::OpClassTensorOp,
-        arch,
-        ThreadblockShape,
-        WarpShape,
-        typename MixedGemmArchTraits::InstructionShape,
-        EpilogueOp,
-        ThreadBlockSwizzle,
-        Stages,
-        true,
-        typename MixedGemmArchTraits::Operator>::GemmKernel;
-    using GemmKernel = cutlass::gemm::kernel::MoeW4A8Gemm<typename GemmKernel_::Mma,
-                                                          typename GemmKernel_::Epilogue,
-                                                          typename GemmKernel_::ThreadblockSwizzle,
-                                                          arch,  // Ensure top level arch is used for dispatch
-                                                          GemmKernel_::kSplitKSerial,
-                                                          cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>;
+    using ThreadBlockSwizzle = typename cutlass::gemm::threadblock::
+        GemmBatchedIdentityThreadblockSwizzle;
+    using GemmKernel_ =
+        typename cutlass::gemm::kernel::DefaultInt8InterleavedGemm<
+            CutlassIntAType,
+            cutlass::layout::RowMajor,
+            MixedGemmArchTraits::ElementsPerAccessA,
+            CutlassIntBType,
+            typename MixedGemmArchTraits::LayoutB,
+            MixedGemmArchTraits::ElementsPerAccessB,
+            OutputElementType,
+            cutlass::layout::RowMajor,
+            ElementAccumulator,
+            cutlass::arch::OpClassTensorOp,
+            arch,
+            ThreadblockShape,
+            WarpShape,
+            typename MixedGemmArchTraits::InstructionShape,
+            EpilogueOp,
+            ThreadBlockSwizzle,
+            Stages,
+            true,
+            typename MixedGemmArchTraits::Operator>::GemmKernel;
+    using GemmKernel = cutlass::gemm::kernel::MoeW4A8Gemm<
+        typename GemmKernel_::Mma,
+        typename GemmKernel_::Epilogue,
+        typename GemmKernel_::ThreadblockSwizzle,
+        arch,  // Ensure top level arch is used for dispatch
+        GemmKernel_::kSplitKSerial,
+        cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>;
     using AlphaColTileIterator =
         cutlass::epilogue::threadblock::PredicatedTileIterator<
             cutlass::epilogue::threadblock::OutputTileOptimalThreadMap<
@@ -147,38 +158,44 @@ void generic_w4a8_moe_gemm_kernelLauncher(
                 cutlass::sizeof_bits<OutputElementType>::value>,
             OutputElementType>;
 
-    using EpilogueVisitor = typename cutlass::epilogue::threadblock::EpilogueVisitorPerRowPerColNf4<
-        ThreadblockShape,
-        GemmKernel::kThreadCount,
-        AlphaColTileIterator,
-        typename GemmKernel::Epilogue::OutputTileIterator,
-        ElementAccumulator,
-        ElementCompute,
-        EpilogueOp>;
+    using EpilogueVisitor =
+        typename cutlass::epilogue::threadblock::EpilogueVisitorPerRowPerColNf4<
+            ThreadblockShape,
+            GemmKernel::kThreadCount,
+            AlphaColTileIterator,
+            typename GemmKernel::Epilogue::OutputTileIterator,
+            ElementAccumulator,
+            ElementCompute,
+            EpilogueOp>;
 
     /// Epilogue
     using Epilogue = typename cutlass::epilogue::threadblock::
-        EpilogueWithVisitorFromExistingEpilogue<EpilogueVisitor, typename GemmKernel::Epilogue>::Epilogue;
+        EpilogueWithVisitorFromExistingEpilogue<
+            EpilogueVisitor,
+            typename GemmKernel::Epilogue>::Epilogue;
 
     // GEMM
     using GemmWithEpilogueVisitorKernel =
-        cutlass::gemm::kernel::MoeW4A8GemmWithEpilogueVisitorInterleavedNf4<typename GemmKernel::Mma,
-                                                                            Epilogue,
-                                                                            ThreadBlockSwizzle,
-                                                                            cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>;
-
+        cutlass::gemm::kernel::MoeW4A8GemmWithEpilogueVisitorInterleavedNf4<
+            typename GemmKernel::Mma,
+            Epilogue,
+            ThreadBlockSwizzle,
+            cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>;
 
     if (occupancy != nullptr) {
-        *occupancy = compute_occupancy_for_kernel<GemmWithEpilogueVisitorKernel>();
-        return;
+      *occupancy =
+          compute_occupancy_for_kernel<GemmWithEpilogueVisitorKernel>();
+      return;
     }
 
-    using Gemm = cutlass::gemm::device::W4A8MoeGemmUniversalBase<GemmWithEpilogueVisitorKernel>;
+    using Gemm = cutlass::gemm::device::W4A8MoeGemmUniversalBase<
+        GemmWithEpilogueVisitorKernel>;
 
     const int ldb =
-        cutlass::platform::is_same<cutlass::layout::RowMajor, typename MixedGemmArchTraits::LayoutB>::value ?
-            n :
-            k * GemmKernel::kInterleave;
+        cutlass::platform::is_same<cutlass::layout::RowMajor,
+                                   typename MixedGemmArchTraits::LayoutB>::value
+            ? n
+            : k * GemmKernel::kInterleave;
 
     typename EpilogueOp::Params linear_scaling_params;
 
@@ -192,64 +209,79 @@ void generic_w4a8_moe_gemm_kernelLauncher(
 
     const int threadblock_count = multi_processor_count * occupancy_;
 
-    typename Gemm::Arguments args{cutlass::gemm::GemmUniversalMode::kBatched,
-                                  num_experts,
-                                  threadblock_count,
-                                  {total_rows, n, k},
-                                  1,
-                                  {reinterpret_cast<CutlassIntAType*>(const_cast<IntAType*>(A)), k},
-                                  {reinterpret_cast<CutlassIntBType*>(const_cast<IntBType*>(B)), ldb},
-                                  quant_mode,
-                                  {reinterpret_cast<OutputElementType*>(const_cast<OutputType*>(col_scale)), 0},
-                                  {reinterpret_cast<OutputElementType*>(const_cast<OutputType*>(row_scale)), 0},
-                                  {const_cast<int32_t*>(nf4_look_up_table), 0},
-                                  {reinterpret_cast<OutputElementType*>(C), n},
-                                  {reinterpret_cast<OutputElementType*>(C), n},
-                                  total_rows_before_expert,
-                                  total_rows_in_ll_else_minus1,
-                                  n,
-                                  k,
-                                  (int64_t)0,
-                                  (int64_t)0,
-                                  typename EpilogueVisitor::Arguments(linear_scaling_params, 0, 0, 0)};
+    typename Gemm::Arguments args{
+        cutlass::gemm::GemmUniversalMode::kBatched,
+        num_experts,
+        threadblock_count,
+        {total_rows, n, k},
+        1,
+        {reinterpret_cast<CutlassIntAType*>(const_cast<IntAType*>(A)), k},
+        {reinterpret_cast<CutlassIntBType*>(const_cast<IntBType*>(B)), ldb},
+        quant_mode,
+        {reinterpret_cast<OutputElementType*>(
+             const_cast<OutputType*>(col_scale)),
+         0},
+        {reinterpret_cast<OutputElementType*>(
+             const_cast<OutputType*>(row_scale)),
+         0},
+        {const_cast<int32_t*>(nf4_look_up_table), 0},
+        {reinterpret_cast<OutputElementType*>(C), n},
+        {reinterpret_cast<OutputElementType*>(C), n},
+        total_rows_before_expert,
+        total_rows_in_ll_else_minus1,
+        n,
+        k,
+        (int64_t)0,
+        (int64_t)0,
+        typename EpilogueVisitor::Arguments(linear_scaling_params, 0, 0, 0)};
 
-    // This assertion is enabled because because for the column interleaved layout, K MUST be a multiple of
-    // threadblockK. The reason for this is that the default pitchlinear iterators are used to handle walking over the
-    // interleaved matrix. The way masking in handled in these do not map to the interleaved layout. We need to write
-    // our own predicated iterator in order to relax this limitation.
-    if (GemmKernel::kInterleave > 1
-        && ((k % MixedGemmArchTraits::ThreadblockK)
-            || ((k / gemm_config.split_k_factor) % MixedGemmArchTraits::ThreadblockK))) {
-        throw std::runtime_error("Temp assertion: k must be multiple of threadblockK");
+    // This assertion is enabled because because for the column interleaved
+    // layout, K MUST be a multiple of threadblockK. The reason for this is that
+    // the default pitchlinear iterators are used to handle walking over the
+    // interleaved matrix. The way masking in handled in these do not map to the
+    // interleaved layout. We need to write our own predicated iterator in order
+    // to relax this limitation.
+    if (GemmKernel::kInterleave > 1 &&
+        ((k % MixedGemmArchTraits::ThreadblockK) ||
+         ((k / gemm_config.split_k_factor) %
+          MixedGemmArchTraits::ThreadblockK))) {
+      throw std::runtime_error(
+          "Temp assertion: k must be multiple of threadblockK");
     }
 
     Gemm gemm;
     if (gemm.get_workspace_size(args) > workspace_bytes) {
-        std::cout<<
-            "Requested split-k but workspace size insufficient. Falling back to non-split-k implementation."<<std::endl;
-        // If requested split-k factor will require more workspace bytes, revert to standard gemm.
-        args.batch_count = 1;
+      std::cout << "Requested split-k but workspace size insufficient. Falling "
+                   "back to non-split-k implementation."
+                << std::endl;
+      // If requested split-k factor will require more workspace bytes, revert
+      // to standard gemm.
+      args.batch_count = 1;
     }
     auto can_implement = gemm.can_implement(args);
     if (can_implement != cutlass::Status::kSuccess) {
-        std::string err_msg = "intA_intB cutlass kernel will fail for params. Error: "
-                              + std::string(cutlassGetStatusString(can_implement));
-        throw std::runtime_error("[FT Error][intA_intB Runner] " + err_msg);
+      std::string err_msg =
+          "intA_intB cutlass kernel will fail for params. Error: " +
+          std::string(cutlassGetStatusString(can_implement));
+      throw std::runtime_error("[FT Error][intA_intB Runner] " + err_msg);
     }
     auto init_status = gemm.initialize(args, workspace, stream);
     if (init_status != cutlass::Status::kSuccess) {
-        std::string err_msg =
-            "Failed to initialize cutlass intA_intB gemm. Error: " + std::string(cutlassGetStatusString(init_status));
-        throw std::runtime_error("[FT Error][intA_intB Runner] " + err_msg);
+      std::string err_msg =
+          "Failed to initialize cutlass intA_intB gemm. Error: " +
+          std::string(cutlassGetStatusString(init_status));
+      throw std::runtime_error("[FT Error][intA_intB Runner] " + err_msg);
     }
     auto run_status = gemm.run(stream);
     if (run_status != cutlass::Status::kSuccess) {
-        std::string err_msg =
-            "Failed to run cutlass intA_intB gemm. Error: " + std::string(cutlassGetStatusString(run_status));
-        CUTLASS_TRACE_HOST("  [FT Error][intA_intB Runner] " << cutlassGetStatusString(run_status));
-        // throw std::runtime_error("[FT Error][intA_intB Runner] " + err_msg);
+      std::string err_msg = "Failed to run cutlass intA_intB gemm. Error: " +
+                            std::string(cutlassGetStatusString(run_status));
+      CUTLASS_TRACE_HOST("  [FT Error][intA_intB Runner] "
+                         << cutlassGetStatusString(run_status));
+      // throw std::runtime_error("[FT Error][intA_intB Runner] " + err_msg);
     }
-    CUTLASS_TRACE_HOST("  finish run kernel " << cutlassGetStatusString(run_status));
+    CUTLASS_TRACE_HOST("  finish run kernel "
+                       << cutlassGetStatusString(run_status));
   }
 }
 
@@ -279,7 +311,6 @@ void dispatch_gemm_config(const IntAType* A,
                           int multi_processor_count,
                           cudaStream_t stream,
                           int* occupancy = nullptr) {
-
   auto dispatch_by_stage = [&](auto temp_args) {
     using DispatcherStages = dispatch_stages<OutputType,
                                              IntAType,
@@ -331,8 +362,7 @@ void dispatch_gemm_config(const IntAType* A,
     default:
       std::string err_msg = "dispatch_gemm_config does not support stages " +
                             std::to_string(gemm_config.stages);
-      throw std::runtime_error("[W4A8MoE][dispatch_gemm_config] " +
-                               err_msg);
+      throw std::runtime_error("[W4A8MoE][dispatch_gemm_config] " + err_msg);
       break;
   }
 }
@@ -342,27 +372,26 @@ template <typename OutputType,
           typename IntBType,
           typename arch,
           typename EpilogueTag>
-void dispatch_moe_gemm_to_cutlass(
-    const IntAType* A,
-    const IntBType* B,
-    cutlass::epilogue::QuantMode quant_mode,
-    const OutputType* col_scale,
-    const OutputType* row_scale,
-    const int32_t* nf4_look_up_table,
-    OutputType* C,
-    int64_t* total_rows_before_expert,
-    int64_t total_rows_in_ll_else_minus1,
-    int64_t total_rows,
-    int64_t gemm_n,
-    int64_t gemm_k,
-    int num_experts,
-    CutlassGemmConfig gemm_config,
-    char* workspace_ptr,
-    const size_t workspace_bytes,
-    // int sm_version,
-    int multi_processor_count,
-    cudaStream_t stream,
-    int* occupancy = nullptr) {
+void dispatch_moe_gemm_to_cutlass(const IntAType* A,
+                                  const IntBType* B,
+                                  cutlass::epilogue::QuantMode quant_mode,
+                                  const OutputType* col_scale,
+                                  const OutputType* row_scale,
+                                  const int32_t* nf4_look_up_table,
+                                  OutputType* C,
+                                  int64_t* total_rows_before_expert,
+                                  int64_t total_rows_in_ll_else_minus1,
+                                  int64_t total_rows,
+                                  int64_t gemm_n,
+                                  int64_t gemm_k,
+                                  int num_experts,
+                                  CutlassGemmConfig gemm_config,
+                                  char* workspace_ptr,
+                                  const size_t workspace_bytes,
+                                  // int sm_version,
+                                  int multi_processor_count,
+                                  cudaStream_t stream,
+                                  int* occupancy = nullptr) {
   // VLOG(1)<<__PRETTY_FUNCTION__;
 
   auto dispatch_by_tile = [&](auto ThreadblockShapeM,
@@ -371,67 +400,67 @@ void dispatch_moe_gemm_to_cutlass(
                               auto WarpShapeM,
                               auto WarpShapeN,
                               auto WarpShapeK) {
-      dispatch_gemm_config<
-          OutputType,
-          IntAType,
-          IntBType,
-          arch,
-          EpilogueTag,
-          cutlass::gemm::GemmShape<decltype(ThreadblockShapeM)::value,
-                                   decltype(ThreadblockShapeN)::value,
-                                   decltype(ThreadblockShapeK)::value>,
-          cutlass::gemm::GemmShape<decltype(WarpShapeM)::value,
-                                   decltype(WarpShapeN)::value,
-                                   decltype(WarpShapeK)::value>>
-                                   (A,
-                                    B,
-                                    quant_mode,
-                                    col_scale,
-                                    row_scale,
-                                    nf4_look_up_table,
-                                    C,
-                                    total_rows_before_expert,
-                                    total_rows_in_ll_else_minus1,
-                                    total_rows,
-                                    gemm_n,
-                                    gemm_k,
-                                    num_experts,
-                                    gemm_config,
-                                    workspace_ptr,
-                                    workspace_bytes,
-                                    multi_processor_count,
-                                    stream,
-                                    occupancy);
+    dispatch_gemm_config<
+        OutputType,
+        IntAType,
+        IntBType,
+        arch,
+        EpilogueTag,
+        cutlass::gemm::GemmShape<decltype(ThreadblockShapeM)::value,
+                                 decltype(ThreadblockShapeN)::value,
+                                 decltype(ThreadblockShapeK)::value>,
+        cutlass::gemm::GemmShape<decltype(WarpShapeM)::value,
+                                 decltype(WarpShapeN)::value,
+                                 decltype(WarpShapeK)::value>>(
+        A,
+        B,
+        quant_mode,
+        col_scale,
+        row_scale,
+        nf4_look_up_table,
+        C,
+        total_rows_before_expert,
+        total_rows_in_ll_else_minus1,
+        total_rows,
+        gemm_n,
+        gemm_k,
+        num_experts,
+        gemm_config,
+        workspace_ptr,
+        workspace_bytes,
+        multi_processor_count,
+        stream,
+        occupancy);
   };
 
   switch (gemm_config.tile_config) {
     case CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64:
-      dispatch_by_tile(Int<16>(), Int<64>(), Int<64>(),
-                       Int<16>(), Int<32>(), Int<64>());
+      dispatch_by_tile(
+          Int<16>(), Int<64>(), Int<64>(), Int<16>(), Int<32>(), Int<64>());
       break;
     case CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64:
-      dispatch_by_tile(Int<32>(), Int<128>(), Int<64>(),
-                       Int<32>(), Int<32>(), Int<64>());
+      dispatch_by_tile(
+          Int<32>(), Int<128>(), Int<64>(), Int<32>(), Int<32>(), Int<64>());
       break;
     case CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64:
-      dispatch_by_tile(Int<64>(), Int<128>(), Int<64>(),
-                       Int<64>(), Int<32>(), Int<64>());
+      dispatch_by_tile(
+          Int<64>(), Int<128>(), Int<64>(), Int<64>(), Int<32>(), Int<64>());
       break;
     case CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64:
-      dispatch_by_tile(Int<128>(), Int<128>(), Int<64>(),
-                       Int<128>(), Int<32>(), Int<64>());
+      dispatch_by_tile(
+          Int<128>(), Int<128>(), Int<64>(), Int<128>(), Int<32>(), Int<64>());
       break;
     case CutlassTileConfig::CtaShape32x512x64_WarpShape32x128x64:
-      dispatch_by_tile(Int<32>(), Int<512>(), Int<64>(),
-                       Int<32>(), Int<128>(), Int<64>());
+      dispatch_by_tile(
+          Int<32>(), Int<512>(), Int<64>(), Int<32>(), Int<128>(), Int<64>());
       break;
     case CutlassTileConfig::CtaShape32x256x64_WarpShape32x64x64:
-      dispatch_by_tile(Int<32>(), Int<256>(), Int<64>(),
-                       Int<32>(), Int<64>(), Int<64>());
+      dispatch_by_tile(
+          Int<32>(), Int<256>(), Int<64>(), Int<32>(), Int<64>(), Int<64>());
       break;
     case CutlassTileConfig::CtaShape64x256x64_WarpShape64x64x64:
-      dispatch_by_tile(Int<64>(), Int<256>(), Int<64>(),
-                       Int<64>(), Int<64>(), Int<64>());
+      dispatch_by_tile(
+          Int<64>(), Int<256>(), Int<64>(), Int<64>(), Int<64>(), Int<64>());
       break;
     // case CutlassTileConfig::CtaShape128x256x64_WarpShape128x64x64:
     //   dispatch_by_tile(Int<128>(), Int<256>(), Int<64>(),
@@ -463,7 +492,6 @@ void dispatch_moe_gemm_to_cutlass(
   }
 }
 
-
 template <typename OutputType, typename IntAType, typename IntBType>
 W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::W4A8MoeGemmRunner() {
   int device{-1};
@@ -472,137 +500,155 @@ W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::W4A8MoeGemmRunner() {
   // sm_ = 80;
   check_cuda_error(cudaDeviceGetAttribute(
       &multi_processor_count_, cudaDevAttrMultiProcessorCount, device));
-  std::string FLAGS_cutlass_w4a8_moe_best_config="";
+  std::string FLAGS_cutlass_w4a8_moe_best_config = "";
   if (getenv("FLAGS_cutlass_w4a8_moe_best_config")) {
-        FLAGS_cutlass_w4a8_moe_best_config = getenv("FLAGS_cutlass_w4a8_moe_best_config");
+    FLAGS_cutlass_w4a8_moe_best_config =
+        getenv("FLAGS_cutlass_w4a8_moe_best_config");
   }
-  if(tuned_configs_from_file.empty() && FLAGS_cutlass_w4a8_moe_best_config!="") {
+  if (tuned_configs_from_file.empty() &&
+      FLAGS_cutlass_w4a8_moe_best_config != "") {
     std::string config_file_path = FLAGS_cutlass_w4a8_moe_best_config;
-    if (config_file_path.find(".config")!=std::string::npos) {
+    if (config_file_path.find(".config") != std::string::npos) {
       std::ifstream config_file(FLAGS_cutlass_w4a8_moe_best_config);
-        if (config_file.is_open()) {
-          VLOG(1)<<"Get tuned w4a8 moe gemm config from: "<<config_file_path;
-          std::string config_string;
-          while(std::getline(config_file, config_string)) {
-            // decode one line of base64 string
-            config_string = base64_decode(config_string);
-            VLOG(1)<<"decode config_string: " << config_string;
-            std::stringstream ss(config_string);
-            std::string item;
-            std::vector<int> vec_configs;
-            while(std::getline(ss, item, ',')) {
-              try {
-                  int value = std::stoi(item);
-                  vec_configs.push_back(value);
-              } catch (const std::invalid_argument& e) {
-                  std::cerr << "Invalid argument: " << item << " is not an integer." << std::endl;
-                  return;
-              } catch (const std::out_of_range& e) {
-                  std::cerr << "Out of range: " << item << " is out of the range of representable values." << std::endl;
-                  return;
-              }
+      if (config_file.is_open()) {
+        VLOG(1) << "Get tuned w4a8 moe gemm config from: " << config_file_path;
+        std::string config_string;
+        while (std::getline(config_file, config_string)) {
+          // decode one line of base64 string
+          config_string = base64_decode(config_string);
+          VLOG(1) << "decode config_string: " << config_string;
+          std::stringstream ss(config_string);
+          std::string item;
+          std::vector<int> vec_configs;
+          while (std::getline(ss, item, ',')) {
+            try {
+              int value = std::stoi(item);
+              vec_configs.push_back(value);
+            } catch (const std::invalid_argument& e) {
+              std::cerr << "Invalid argument: " << item << " is not an integer."
+                        << std::endl;
+              return;
+            } catch (const std::out_of_range& e) {
+              std::cerr << "Out of range: " << item
+                        << " is out of the range of representable values."
+                        << std::endl;
+              return;
             }
-            W4A8MoeGEMMConfig search_config;
-            search_config.total_rows = vec_configs[0];
-            search_config.n = vec_configs[1];
-            search_config.k = vec_configs[2];
-            search_config.num_experts = vec_configs[3];
-            search_config.tile_config = static_cast<CutlassTileConfig>(vec_configs[4]);
-            search_config.split_k_style = static_cast<SplitKStyle>(vec_configs[5]);
-            search_config.split_k_factor = vec_configs[6];
-            search_config.stages = vec_configs[7];
-            tuned_configs_from_file.push_back(search_config);
-            VLOG(1)<<"tuned_configs_from_file: "<<search_config.total_rows<<","<<search_config.n<<","<<search_config.k<<","<<search_config.num_experts<<","<<static_cast<int>(search_config.tile_config)<<"," << static_cast<int>(search_config.split_k_style)<<","<<search_config.split_k_factor<<","<<search_config.stages;
-
           }
-        } else {
-          VLOG(1)<<"No tuned w4a8 gemm config.";
+          W4A8MoeGEMMConfig search_config;
+          search_config.total_rows = vec_configs[0];
+          search_config.n = vec_configs[1];
+          search_config.k = vec_configs[2];
+          search_config.num_experts = vec_configs[3];
+          search_config.tile_config =
+              static_cast<CutlassTileConfig>(vec_configs[4]);
+          search_config.split_k_style =
+              static_cast<SplitKStyle>(vec_configs[5]);
+          search_config.split_k_factor = vec_configs[6];
+          search_config.stages = vec_configs[7];
+          tuned_configs_from_file.push_back(search_config);
+          VLOG(1) << "tuned_configs_from_file: " << search_config.total_rows
+                  << "," << search_config.n << "," << search_config.k << ","
+                  << search_config.num_experts << ","
+                  << static_cast<int>(search_config.tile_config) << ","
+                  << static_cast<int>(search_config.split_k_style) << ","
+                  << search_config.split_k_factor << ","
+                  << search_config.stages;
         }
+      } else {
+        VLOG(1) << "No tuned w4a8 gemm config.";
+      }
 
     } else {
-      FILE * fp;
+      FILE* fp;
       fp = fopen(config_file_path.c_str(), "r");
-      if(fp) {
-        VLOG(1)<<"Get tuned w4a8 moe gemm config from: "<<config_file_path;
+      if (fp) {
+        VLOG(1) << "Get tuned w4a8 moe gemm config from: " << config_file_path;
         int tile_config, split_k_style, split_k_factor, stages;
         int total_rows_tmp, k_tmp, n_tmp, num_experts_tmp;
-        while(1) {
-            fscanf(fp, "%d, %d, %d, %d, %d, %d, %d, %d", &total_rows_tmp, &n_tmp, &k_tmp, &num_experts_tmp, &tile_config, &split_k_style, &split_k_factor, &stages);
-            W4A8MoeGEMMConfig search_config;
-            search_config.total_rows = total_rows_tmp;
-            search_config.n = n_tmp;
-            search_config.k = k_tmp;
-            search_config.num_experts = num_experts_tmp;
-            search_config.tile_config = static_cast<CutlassTileConfig>(tile_config);
-            search_config.split_k_style = static_cast<SplitKStyle>(split_k_style);
-            search_config.split_k_factor = split_k_factor;
-            search_config.stages = stages;
-            tuned_configs_from_file.push_back(search_config);
-            VLOG(1)<<"tuned_configs_from_file: "<<total_rows_tmp<<","<<n_tmp<<","<<k_tmp<<","<<num_experts_tmp<<","<<tile_config<<"," << split_k_style<<","<<split_k_factor<<","<<stages;
-            if (feof(fp))
-            break;
+        while (1) {
+          fscanf(fp,
+                 "%d, %d, %d, %d, %d, %d, %d, %d",
+                 &total_rows_tmp,
+                 &n_tmp,
+                 &k_tmp,
+                 &num_experts_tmp,
+                 &tile_config,
+                 &split_k_style,
+                 &split_k_factor,
+                 &stages);
+          W4A8MoeGEMMConfig search_config;
+          search_config.total_rows = total_rows_tmp;
+          search_config.n = n_tmp;
+          search_config.k = k_tmp;
+          search_config.num_experts = num_experts_tmp;
+          search_config.tile_config =
+              static_cast<CutlassTileConfig>(tile_config);
+          search_config.split_k_style = static_cast<SplitKStyle>(split_k_style);
+          search_config.split_k_factor = split_k_factor;
+          search_config.stages = stages;
+          tuned_configs_from_file.push_back(search_config);
+          VLOG(1) << "tuned_configs_from_file: " << total_rows_tmp << ","
+                  << n_tmp << "," << k_tmp << "," << num_experts_tmp << ","
+                  << tile_config << "," << split_k_style << ","
+                  << split_k_factor << "," << stages;
+          if (feof(fp)) break;
         }
-      } else if(FLAGS_cutlass_w4a8_moe_best_config=="") {
-          VLOG(1)<<"No tuned w4a8 gemm config.";
+      } else if (FLAGS_cutlass_w4a8_moe_best_config == "") {
+        VLOG(1) << "No tuned w4a8 gemm config.";
       }
     }
   }
 }
 
 template <typename OutputType, typename IntAType, typename IntBType>
-W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::~W4A8MoeGemmRunner() {
-}
-
-
+W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::~W4A8MoeGemmRunner() {}
 
 template <typename OutputType, typename IntAType, typename IntBType>
 template <typename EpilogueTag>
-void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::dispatch_to_arch<EpilogueTag>(
-    const IntAType* A,
-    const IntBType* B,
-    cutlass::epilogue::QuantMode quant_mode,
-    const OutputType* col_scale,
-    const OutputType* row_scale,
-    const int32_t* nf4_look_up_table,
-    OutputType* C,
-    int64_t* total_rows_before_expert,
-    int64_t total_rows_in_ll_else_minus1,
-    int64_t total_rows,
-    int64_t gemm_n,
-    int64_t gemm_k,
-    int num_experts,
-    CutlassGemmConfig gemm_config,
-    char* workspace_ptr,
-    const size_t workspace_bytes,
-    cudaStream_t stream,
-    int* occupancy) {
-
+void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::dispatch_to_arch<
+    EpilogueTag>(const IntAType* A,
+                 const IntBType* B,
+                 cutlass::epilogue::QuantMode quant_mode,
+                 const OutputType* col_scale,
+                 const OutputType* row_scale,
+                 const int32_t* nf4_look_up_table,
+                 OutputType* C,
+                 int64_t* total_rows_before_expert,
+                 int64_t total_rows_in_ll_else_minus1,
+                 int64_t total_rows,
+                 int64_t gemm_n,
+                 int64_t gemm_k,
+                 int num_experts,
+                 CutlassGemmConfig gemm_config,
+                 char* workspace_ptr,
+                 const size_t workspace_bytes,
+                 cudaStream_t stream,
+                 int* occupancy) {
   // only sm80 here
   dispatch_moe_gemm_to_cutlass<OutputType,
-                              IntAType,
-                              IntBType,
-                              cutlass::arch::Sm80,
-                              EpilogueTag>(A,
-                                          B,
-                                          quant_mode,
-                                          col_scale,
-                                          row_scale,
-                                          nf4_look_up_table,
-                                          C,
-                                          total_rows_before_expert,
-                                          total_rows_in_ll_else_minus1,
-                                          total_rows,
-                                          gemm_n,
-                                          gemm_k,
-                                          num_experts,
-                                          gemm_config,
-                                          workspace_ptr,
-                                          workspace_bytes,
-                                          multi_processor_count_,
-                                          stream,
-                                          occupancy);
-
-
+                               IntAType,
+                               IntBType,
+                               cutlass::arch::Sm80,
+                               EpilogueTag>(A,
+                                            B,
+                                            quant_mode,
+                                            col_scale,
+                                            row_scale,
+                                            nf4_look_up_table,
+                                            C,
+                                            total_rows_before_expert,
+                                            total_rows_in_ll_else_minus1,
+                                            total_rows,
+                                            gemm_n,
+                                            gemm_k,
+                                            num_experts,
+                                            gemm_config,
+                                            workspace_ptr,
+                                            workspace_bytes,
+                                            multi_processor_count_,
+                                            stream,
+                                            occupancy);
 }
 
 template <typename OutputType, typename IntAType, typename IntBType>
@@ -625,8 +671,8 @@ void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::run_gemm<EpilogueTag>(
     int num_experts,
     cudaStream_t stream,
     CutlassGemmConfig gemm_config) {
-  VLOG(1)<<__PRETTY_FUNCTION__;
-  static constexpr bool is_weight_only = true; //todo(yuanxiaolan)
+  VLOG(1) << __PRETTY_FUNCTION__;
+  static constexpr bool is_weight_only = true;  // todo(yuanxiaolan)
   bool is_weight_only_encoder = total_rows >= 512 ? true : false;
 
   VLOG(1) << "gemm_config tile_config"
@@ -636,29 +682,29 @@ void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::run_gemm<EpilogueTag>(
   VLOG(1) << "gemm_config split_k_factor " << gemm_config.split_k_factor;
   VLOG(1) << "gemm_config stages " << gemm_config.stages;
 
-  if(gemm_config.tile_config != CutlassTileConfig::Undefined) {
+  if (gemm_config.tile_config != CutlassTileConfig::Undefined) {
     dispatch_to_arch<EpilogueTag>(A,
-                                    B,
-                                    quant_mode,
-                                    col_scale,
-                                    row_scale,
-                                    nf4_look_up_table,
-                                    C,
-                                    total_rows_before_expert,
-                                    total_rows_in_ll_else_minus1,
-                                    total_rows,
-                                    gemm_n,
-                                    gemm_k,
-                                    num_experts,
-                                    gemm_config,
-                                    workspace_ptr,
-                                    workspace_bytes,
-                                    stream);
+                                  B,
+                                  quant_mode,
+                                  col_scale,
+                                  row_scale,
+                                  nf4_look_up_table,
+                                  C,
+                                  total_rows_before_expert,
+                                  total_rows_in_ll_else_minus1,
+                                  total_rows,
+                                  gemm_n,
+                                  gemm_k,
+                                  num_experts,
+                                  gemm_config,
+                                  workspace_ptr,
+                                  workspace_bytes,
+                                  stream);
     return;
   }
 
-  std::vector<CutlassGemmConfig> candidate_configs =
-      get_candidate_configs_nf4(80, is_weight_only, is_weight_only_encoder, false);
+  std::vector<CutlassGemmConfig> candidate_configs = get_candidate_configs_nf4(
+      80, is_weight_only, is_weight_only_encoder, false);
   std::vector<int> occupancies(candidate_configs.size());
 
   for (size_t ii = 0; ii < candidate_configs.size(); ++ii) {
@@ -686,20 +732,21 @@ void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::run_gemm<EpilogueTag>(
   int local_multi_processor_count{0};
   check_cuda_error(cudaGetDevice(&local_device));
   // sm_ = getSMVersion();
-  check_cuda_error(cudaDeviceGetAttribute(
-      &local_multi_processor_count, cudaDevAttrMultiProcessorCount, local_device));
+  check_cuda_error(cudaDeviceGetAttribute(&local_multi_processor_count,
+                                          cudaDevAttrMultiProcessorCount,
+                                          local_device));
 
   CutlassGemmConfig chosen_config =
       estimate_best_config_from_occupancies_w4a4(candidate_configs,
-                                            occupancies,
-                                            total_rows,
-                                            gemm_n,
-                                            gemm_k,
-                                            num_experts,
-                                            split_k_limit,
-                                            workspace_bytes,
-                                            local_multi_processor_count,
-                                            is_weight_only);
+                                                 occupancies,
+                                                 total_rows,
+                                                 gemm_n,
+                                                 gemm_k,
+                                                 num_experts,
+                                                 split_k_limit,
+                                                 workspace_bytes,
+                                                 local_multi_processor_count,
+                                                 is_weight_only);
 
   VLOG(1) << "chosen_config tile_config "
           << static_cast<int>(chosen_config.tile_config);
@@ -711,7 +758,6 @@ void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::run_gemm<EpilogueTag>(
   VLOG(1) << "total_rows  " << total_rows << "gemm_n  " << gemm_n << "gemm_k  "
           << gemm_k;
 
-
   dispatch_to_arch<EpilogueTag>(A,
                                 B,
                                 quant_mode,
@@ -732,7 +778,8 @@ void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::run_gemm<EpilogueTag>(
 }
 
 // template <typename OutputType, typename IntAType, typename IntBType>
-// void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::moe_gemm_bias_act(    const IntAType* A,
+// void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::moe_gemm_bias_act(
+// const IntAType* A,
 //     const IntBType* B,
 //     QuantMode quant_mode,
 //     const OutputType* col_scale,
@@ -769,65 +816,82 @@ void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::run_gemm<EpilogueTag>(
 
 template <typename OutputType, typename IntAType, typename IntBType>
 void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::moe_gemm(
-  const IntAType* A,
-  const IntBType* B,
-  cutlass::epilogue::QuantMode quant_mode,
-  const OutputType* col_scale,
-  const OutputType* row_scale,
-  const int32_t* nf4_look_up_table,
-  OutputType* C,
-  int64_t* total_rows_before_expert,
-  int64_t total_rows_in_ll_else_minus1,
-  int64_t total_rows,
-  int64_t gemm_n,
-  int64_t gemm_k,
-  char* workspace_ptr,
-  const size_t workspace_bytes,
-  int num_experts,
-  cudaStream_t stream,
-  CutlassGemmConfig gemm_config) {
+    const IntAType* A,
+    const IntBType* B,
+    cutlass::epilogue::QuantMode quant_mode,
+    const OutputType* col_scale,
+    const OutputType* row_scale,
+    const int32_t* nf4_look_up_table,
+    OutputType* C,
+    int64_t* total_rows_before_expert,
+    int64_t total_rows_in_ll_else_minus1,
+    int64_t total_rows,
+    int64_t gemm_n,
+    int64_t gemm_k,
+    char* workspace_ptr,
+    const size_t workspace_bytes,
+    int num_experts,
+    cudaStream_t stream,
+    CutlassGemmConfig gemm_config) {
   CutlassGemmConfig gemm_config_from_file_and_param = gemm_config;
-  if(!tuned_configs_from_file.empty()){
-    bool match=false;
+  if (!tuned_configs_from_file.empty()) {
+    bool match = false;
     int best_total_rows, best_n, best_k, best_num_experts;
-    int max_config_total_rows_in_file=0;
+    int max_config_total_rows_in_file = 0;
     W4A8MoeGEMMConfig max_total_rows_config;
-    for(const auto& tuned_config:tuned_configs_from_file) {
-        // choose the smallest config_m with config_m >=m
-        if(tuned_config.total_rows <= total_rows && tuned_config.n==gemm_n && tuned_config.k==gemm_k && tuned_config.num_experts==num_experts) {
-          best_total_rows=tuned_config.total_rows;
-          best_n=tuned_config.n;
-          best_k=tuned_config.k;
-          best_num_experts=tuned_config.num_experts;
-          gemm_config_from_file_and_param.tile_config = tuned_config.tile_config;
-          gemm_config_from_file_and_param.split_k_style = tuned_config.split_k_style;
-          gemm_config_from_file_and_param.split_k_factor = tuned_config.split_k_factor;
-          gemm_config_from_file_and_param.stages = tuned_config.stages;
-          match=true;
-        }
-        if(tuned_config.total_rows > max_config_total_rows_in_file && tuned_config.n==gemm_n && tuned_config.k==gemm_k && tuned_config.num_experts==num_experts){
-            max_config_total_rows_in_file = tuned_config.total_rows;
-            max_total_rows_config = tuned_config;
-        }
+    for (const auto& tuned_config : tuned_configs_from_file) {
+      // choose the smallest config_m with config_m >=m
+      if (tuned_config.total_rows <= total_rows && tuned_config.n == gemm_n &&
+          tuned_config.k == gemm_k && tuned_config.num_experts == num_experts) {
+        best_total_rows = tuned_config.total_rows;
+        best_n = tuned_config.n;
+        best_k = tuned_config.k;
+        best_num_experts = tuned_config.num_experts;
+        gemm_config_from_file_and_param.tile_config = tuned_config.tile_config;
+        gemm_config_from_file_and_param.split_k_style =
+            tuned_config.split_k_style;
+        gemm_config_from_file_and_param.split_k_factor =
+            tuned_config.split_k_factor;
+        gemm_config_from_file_and_param.stages = tuned_config.stages;
+        match = true;
+      }
+      if (tuned_config.total_rows > max_config_total_rows_in_file &&
+          tuned_config.n == gemm_n && tuned_config.k == gemm_k &&
+          tuned_config.num_experts == num_experts) {
+        max_config_total_rows_in_file = tuned_config.total_rows;
+        max_total_rows_config = tuned_config;
+      }
     }
-    if(!match){
-      if (max_total_rows_config.n==gemm_n && max_total_rows_config.k==gemm_k && max_total_rows_config.num_experts==num_experts) {
+    if (!match) {
+      if (max_total_rows_config.n == gemm_n &&
+          max_total_rows_config.k == gemm_k &&
+          max_total_rows_config.num_experts == num_experts) {
         best_total_rows = max_config_total_rows_in_file;
-        gemm_config_from_file_and_param.tile_config = max_total_rows_config.tile_config;
-        gemm_config_from_file_and_param.split_k_style = max_total_rows_config.split_k_style;
-        gemm_config_from_file_and_param.split_k_factor = max_total_rows_config.split_k_factor;
+        gemm_config_from_file_and_param.tile_config =
+            max_total_rows_config.tile_config;
+        gemm_config_from_file_and_param.split_k_style =
+            max_total_rows_config.split_k_style;
+        gemm_config_from_file_and_param.split_k_factor =
+            max_total_rows_config.split_k_factor;
         gemm_config_from_file_and_param.stages = max_total_rows_config.stages;
       }
     }
-    VLOG(1) <<"W4A8 moe gemm "
-            <<"total_rows: "<<total_rows<<" n: "<<gemm_n<<" k: "<<gemm_k
-            <<"Using gemm config from config file: config_total_rows: "<< best_total_rows<<" config_n: "<< best_n << " config_k: "<< best_k
-            <<"tile_config: "<<static_cast<int>(gemm_config_from_file_and_param.tile_config)
-            <<"split_k_style: "<<static_cast<int>(gemm_config_from_file_and_param.split_k_style)
-            <<"split_k_factor: "<<static_cast<int>(gemm_config_from_file_and_param.split_k_factor)
-            <<"stages: "<<static_cast<int>(gemm_config_from_file_and_param.stages);
+    VLOG(1) << "W4A8 moe gemm "
+            << "total_rows: " << total_rows << " n: " << gemm_n
+            << " k: " << gemm_k
+            << "Using gemm config from config file: config_total_rows: "
+            << best_total_rows << " config_n: " << best_n
+            << " config_k: " << best_k << "tile_config: "
+            << static_cast<int>(gemm_config_from_file_and_param.tile_config)
+            << "split_k_style: "
+            << static_cast<int>(gemm_config_from_file_and_param.split_k_style)
+            << "split_k_factor: "
+            << static_cast<int>(gemm_config_from_file_and_param.split_k_factor)
+            << "stages: "
+            << static_cast<int>(gemm_config_from_file_and_param.stages);
   } else {
-    VLOG(1) << "tuned_configs_from_file is empty, use W4A8 gemm config in param";
+    VLOG(1)
+        << "tuned_configs_from_file is empty, use W4A8 gemm config in param";
   }
   run_gemm<EpilogueOpNoBias>(A,
                              B,
@@ -849,7 +913,10 @@ void W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::moe_gemm(
 }
 
 template <typename OutputType, typename IntAType, typename IntBType>
-std::vector<typename W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::W4A8MoeGEMMConfig> W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::tuned_configs_from_file = {};
+std::vector<typename W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::
+                W4A8MoeGEMMConfig>
+    W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::tuned_configs_from_file =
+        {};
 
 template <typename OutputType, typename IntAType, typename IntBType>
 int W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::getWorkspaceSize(
@@ -863,6 +930,5 @@ int W4A8MoeGemmRunner<OutputType, IntAType, IntBType>::getWorkspaceSize(
   return max_grid_m * max_grid_n * split_k_limit * 4;
 }
 
-
 template class W4A8MoeGemmRunner<half, int8_t, cutlass::uint4b_t>;
 template class W4A8MoeGemmRunner<__nv_bfloat16, int8_t, cutlass::uint4b_t>;
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_kernel.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_kernel.h
index 7d06b59d65..6ec9121634 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_kernel.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_kernel.h
@@ -30,45 +30,46 @@ class W4A8MoeGemmRunner {
   ~W4A8MoeGemmRunner();
 
   void moe_gemm_bias_act(const IntAType* A,
-                     const IntBType* B,
-                     cutlass::epilogue::QuantMode quant_mode,
-                     const OutputType* col_scale,
-                     const OutputType* row_scale,
-                     const OutputType* biases,
-                     const int32_t* nf4_look_up_table,
-                     OutputType* C,
-                     int64_t* total_rows_before_expert,
-                     int m,
-                     int n,
-                     int k,
-                     int num_experts,
-                     std::string activation_type,
-                     char* workspace_ptr,
-                     const size_t workspace_bytes,
-                     cudaStream_t stream);
+                         const IntBType* B,
+                         cutlass::epilogue::QuantMode quant_mode,
+                         const OutputType* col_scale,
+                         const OutputType* row_scale,
+                         const OutputType* biases,
+                         const int32_t* nf4_look_up_table,
+                         OutputType* C,
+                         int64_t* total_rows_before_expert,
+                         int m,
+                         int n,
+                         int k,
+                         int num_experts,
+                         std::string activation_type,
+                         char* workspace_ptr,
+                         const size_t workspace_bytes,
+                         cudaStream_t stream);
 
   void moe_gemm(const IntAType* A,
-            const IntBType* B,
-            cutlass::epilogue::QuantMode quant_mode,
-            const OutputType* col_scale,
-            const OutputType* row_scale,
-            const int32_t* nf4_look_up_table,
-            OutputType* C,
-            int64_t* total_rows_before_expert,
-            int64_t total_rows_in_ll_else_minus1,
-            int64_t total_rows,
-            int64_t gemm_n,
-            int64_t gemm_k,
-            char* workspace_ptr,
-            const size_t workspace_bytes,
-            int num_experts,
-            cudaStream_t stream,
-            CutlassGemmConfig gemm_config = CutlassGemmConfig{CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64,
-                                                              SplitKStyle::NO_SPLIT_K,
-                                                              1,
-                                                              5});
- private:
+                const IntBType* B,
+                cutlass::epilogue::QuantMode quant_mode,
+                const OutputType* col_scale,
+                const OutputType* row_scale,
+                const int32_t* nf4_look_up_table,
+                OutputType* C,
+                int64_t* total_rows_before_expert,
+                int64_t total_rows_in_ll_else_minus1,
+                int64_t total_rows,
+                int64_t gemm_n,
+                int64_t gemm_k,
+                char* workspace_ptr,
+                const size_t workspace_bytes,
+                int num_experts,
+                cudaStream_t stream,
+                CutlassGemmConfig gemm_config = CutlassGemmConfig{
+                    CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64,
+                    SplitKStyle::NO_SPLIT_K,
+                    1,
+                    5});
 
+ private:
   template <typename EpilogueTag>
   void dispatch_to_arch(const IntAType* A,
                         const IntBType* B,
@@ -108,8 +109,7 @@ class W4A8MoeGemmRunner {
                 cudaStream_t stream,
                 CutlassGemmConfig gemm_config);
 
-  int getWorkspaceSize(
-    const int m, const int n, const int k);
+  int getWorkspaceSize(const int m, const int n, const int k);
 
  private:
   static constexpr int split_k_limit = 4;
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_kernel_template.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_kernel_template.h
index 80ffb06de4..3ab51307d6 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_kernel_template.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_kernel_template.h
@@ -29,7 +29,8 @@ template <typename OutputType,
           typename ThreadblockShape,
           typename WarpShape,
           int Stages>
-void generic_w4a8_moe_gemm_kernelLauncher(const IntAType* A,
+void generic_w4a8_moe_gemm_kernelLauncher(
+    const IntAType* A,
     const IntBType* B,
     cutlass::epilogue::QuantMode quant_mode,
     const OutputType* col_scale,
@@ -48,7 +49,6 @@ void generic_w4a8_moe_gemm_kernelLauncher(const IntAType* A,
     cudaStream_t stream,
     int* occupancy);
 
-
 template <typename OutputType,
           typename IntAType,
           typename IntBType,
@@ -79,8 +79,7 @@ struct dispatch_stages {
     std::string err_msg = "Cutlass fpA_intB gemm. Not instantiates for arch " +
                           std::to_string(arch::kMinComputeCapability) +
                           " with stages set to " + std::to_string(Stages);
-    throw std::runtime_error("[dispatch_stages::dispatch] " +
-                             err_msg);
+    throw std::runtime_error("[dispatch_stages::dispatch] " + err_msg);
   }
 };
 
@@ -92,13 +91,13 @@ template <typename OutputType,
           typename ThreadblockShape,
           typename WarpShape>
 struct dispatch_stages<OutputType,
-                        IntAType,
-                        IntBType,
-                        arch,
-                        EpilogueTag,
-                        ThreadblockShape,
-                        WarpShape,
-                        2> {
+                       IntAType,
+                       IntBType,
+                       arch,
+                       EpilogueTag,
+                       ThreadblockShape,
+                       WarpShape,
+                       2> {
   static void dispatch(const IntAType* A,
                        const IntBType* B,
                        cutlass::epilogue::QuantMode quant_mode,
@@ -119,31 +118,31 @@ struct dispatch_stages<OutputType,
                        cudaStream_t stream,
                        int* occupancy = nullptr) {
     generic_w4a8_moe_gemm_kernelLauncher<OutputType,
-                                            IntAType,
-                                            IntBType,
-                                            arch,
-                                            EpilogueTag,
-                                            ThreadblockShape,
-                                            WarpShape,
-                                            2>(A,
-                                       B,
-                                       quant_mode,
-                                       col_scale,
-                                       row_scale,
-                                       nf4_look_up_table,
-                                       C,
-                                       total_rows_before_expert,
-                                       total_rows,
-                                       total_rows_in_ll_else_minus1,
-                                       n,
-                                       k,
-                                       num_experts,
-                                       gemm_config,
-                                       workspace,
-                                       workspace_bytes,
-                                       multi_processor_count,
-                                       stream,
-                                       occupancy);
+                                         IntAType,
+                                         IntBType,
+                                         arch,
+                                         EpilogueTag,
+                                         ThreadblockShape,
+                                         WarpShape,
+                                         2>(A,
+                                            B,
+                                            quant_mode,
+                                            col_scale,
+                                            row_scale,
+                                            nf4_look_up_table,
+                                            C,
+                                            total_rows_before_expert,
+                                            total_rows,
+                                            total_rows_in_ll_else_minus1,
+                                            n,
+                                            k,
+                                            num_experts,
+                                            gemm_config,
+                                            workspace,
+                                            workspace_bytes,
+                                            multi_processor_count,
+                                            stream,
+                                            occupancy);
   }
 };
 
@@ -155,12 +154,12 @@ template <typename OutputType,
           typename WarpShape,
           int Stages>
 struct dispatch_stages<OutputType,
-                        IntAType,
-                        IntBType,
-                        cutlass::arch::Sm80,
-                        EpilogueTag,
-                        ThreadblockShape,
-                        WarpShape,
+                       IntAType,
+                       IntBType,
+                       cutlass::arch::Sm80,
+                       EpilogueTag,
+                       ThreadblockShape,
+                       WarpShape,
                        Stages,
                        typename std::enable_if<(Stages > 2)>::type> {
   static void dispatch(const IntAType* A,
@@ -183,35 +182,34 @@ struct dispatch_stages<OutputType,
                        cudaStream_t stream,
                        int* occupancy = nullptr) {
     generic_w4a8_moe_gemm_kernelLauncher<OutputType,
-                                            IntAType,
-                                            IntBType,
-                                            cutlass::arch::Sm80,
-                                            EpilogueTag,
-                                            ThreadblockShape,
-                                            WarpShape,
-                                            Stages>(A,
-                                            B,
-                                            quant_mode,
-                                             col_scale,
-                                             row_scale,
-                                             nf4_look_up_table,
-                                            C,
-                                            total_rows_before_expert,
-                                            total_rows,
-                                            total_rows_in_ll_else_minus1,
-                                            n,
-                                            k,
-                                            num_experts,
-                                            gemm_config,
-                                             workspace,
-                                             workspace_bytes,
-                                            multi_processor_count,
-                                            stream,
-                                            occupancy);
+                                         IntAType,
+                                         IntBType,
+                                         cutlass::arch::Sm80,
+                                         EpilogueTag,
+                                         ThreadblockShape,
+                                         WarpShape,
+                                         Stages>(A,
+                                                 B,
+                                                 quant_mode,
+                                                 col_scale,
+                                                 row_scale,
+                                                 nf4_look_up_table,
+                                                 C,
+                                                 total_rows_before_expert,
+                                                 total_rows,
+                                                 total_rows_in_ll_else_minus1,
+                                                 n,
+                                                 k,
+                                                 num_experts,
+                                                 gemm_config,
+                                                 workspace,
+                                                 workspace_bytes,
+                                                 multi_processor_count,
+                                                 stream,
+                                                 occupancy);
   }
 };
 
-
 template <typename OutputType,
           typename IntAType,
           typename IntBType,
@@ -231,8 +229,8 @@ void dispatch_gemm_config(const IntAType* A,
                           int64_t gemm_k,
                           int num_experts,
                           CutlassGemmConfig gemm_config,
-                           char* workspace,
-                           size_t workspace_bytes,
+                          char* workspace,
+                          size_t workspace_bytes,
                           int multi_processor_count,
                           cudaStream_t stream,
                           int* occupancy = nullptr);
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_test.cu b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_test.cu
index 76e0195af3..15893137af 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_test.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_test.cu
@@ -43,26 +43,31 @@ limitations under the License. */
 #include "nvToolsExt.h"
 #endif
 
-const uint32_t colors[] = {0xff00ff00, 0xff0000ff, 0xffffff00, 0xffff00ff,
-                           0xff00ffff, 0xffff0000, 0xffffffff};
+const uint32_t colors[] = {0xff00ff00,
+                           0xff0000ff,
+                           0xffffff00,
+                           0xffff00ff,
+                           0xff00ffff,
+                           0xffff0000,
+                           0xffffffff};
 const int num_colors = sizeof(colors) / sizeof(uint32_t);
 
 using CutlassTileConfig = CutlassTileConfig;
 using SplitKStyle = SplitKStyle;
 using CutlassGemmConfig = CutlassGemmConfig;
 
-#define PUSH_RANGE(name, cid)                                                  \
-  {                                                                            \
-    int color_id = cid;                                                        \
-    color_id = color_id % num_colors;                                          \
-    nvtxEventAttributes_t eventAttrib = {0};                                   \
-    eventAttrib.version = NVTX_VERSION;                                        \
-    eventAttrib.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;                          \
-    eventAttrib.colorType = NVTX_COLOR_ARGB;                                   \
-    eventAttrib.color = colors[color_id];                                      \
-    eventAttrib.messageType = NVTX_MESSAGE_TYPE_ASCII;                         \
-    eventAttrib.message.ascii = name;                                          \
-    nvtxRangePushEx(&eventAttrib);                                             \
+#define PUSH_RANGE(name, cid)                          \
+  {                                                    \
+    int color_id = cid;                                \
+    color_id = color_id % num_colors;                  \
+    nvtxEventAttributes_t eventAttrib = {0};           \
+    eventAttrib.version = NVTX_VERSION;                \
+    eventAttrib.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;  \
+    eventAttrib.colorType = NVTX_COLOR_ARGB;           \
+    eventAttrib.color = colors[color_id];              \
+    eventAttrib.messageType = NVTX_MESSAGE_TYPE_ASCII; \
+    eventAttrib.message.ascii = name;                  \
+    nvtxRangePushEx(&eventAttrib);                     \
   }
 #define POP_RANGE nvtxRangePop();
 #else
@@ -79,7 +84,9 @@ using CutlassGemmConfig = CutlassGemmConfig;
 // }
 
 template <typename T>
-static void PrintMatrix(const T *mat_d, int num, std::string name,
+static void PrintMatrix(const T *mat_d,
+                        int num,
+                        std::string name,
                         int numOfCols) {
   std::vector<T> tmp(num);
   cudaMemcpy(tmp.data(), mat_d, sizeof(T) * num, cudaMemcpyDeviceToHost);
@@ -104,17 +111,17 @@ static void PrintMatrix(const T *mat_d, int num, std::string name,
 
 uint as_uint(const float x) { return *(uint *)&x; }
 uint16_t ConvertFloat2Half(const float x) {
-  const uint b = as_uint(x) + 0x00001000; // round-to-nearest-even: add last
-                                          // bit after truncated mantissa
-  const uint e = (b & 0x7F800000) >> 23;  // exponent
-  const uint m = b & 0x007FFFFF; // mantissa; in line below: 0x007FF000 =
-                                 // 0x00800000-0x00001000 = decimal indicator
-                                 // flag - initial rounding
+  const uint b = as_uint(x) + 0x00001000;  // round-to-nearest-even: add last
+                                           // bit after truncated mantissa
+  const uint e = (b & 0x7F800000) >> 23;   // exponent
+  const uint m = b & 0x007FFFFF;  // mantissa; in line below: 0x007FF000 =
+                                  // 0x00800000-0x00001000 = decimal indicator
+                                  // flag - initial rounding
   return (b & 0x80000000) >> 16 |
          (e > 112) * ((((e - 112) << 10) & 0x7C00) | m >> 13) |
          ((e < 113) & (e > 101)) *
              ((((0x007FF000 + m) >> (125 - e)) + 1) >> 1) |
-         (e > 143) * 0x7FFF; // sign : normalized : denormalized : saturate
+         (e > 143) * 0x7FFF;  // sign : normalized : denormalized : saturate
 }
 
 inline float fp32_from_bits(uint32_t w) {
@@ -274,7 +281,9 @@ float CPUHalfConvert2Float(const uint16_t h) {
   return fp32_from_bits(result);
 }
 
-static void PrintHalfMatrix(const int16_t *mat_d, int num, std::string name,
+static void PrintHalfMatrix(const int16_t *mat_d,
+                            int num,
+                            std::string name,
                             int numOfCols) {
   std::vector<int16_t> tmp(num);
   cudaMemcpy(tmp.data(), mat_d, sizeof(int16_t) * num, cudaMemcpyDeviceToHost);
@@ -296,7 +305,9 @@ static void PrintHalfMatrix(const int16_t *mat_d, int num, std::string name,
 }
 
 template <typename T>
-static void PrintMatrixCPU(const T *mat, int num, std::string name,
+static void PrintMatrixCPU(const T *mat,
+                           int num,
+                           std::string name,
                            int numOfCols) {
   std::ofstream outfile;
   outfile.open(name + ".txt", std::ios::out);
@@ -315,7 +326,9 @@ static void PrintMatrixCPU(const T *mat, int num, std::string name,
   outfile.close();
 }
 
-static void PrintMatrixCPU_int4(const int8_t *mat, int num, std::string name,
+static void PrintMatrixCPU_int4(const int8_t *mat,
+                                int num,
+                                std::string name,
                                 int numOfCols) {
   std::ofstream outfile;
   outfile.open(name + ".txt", std::ios::out);
@@ -333,7 +346,9 @@ static void PrintMatrixCPU_int4(const int8_t *mat, int num, std::string name,
   outfile.close();
 }
 template <typename T>
-static void PrintHalfMatrixCPU(const T *mat, int num, std::string name,
+static void PrintHalfMatrixCPU(const T *mat,
+                               int num,
+                               std::string name,
                                int numOfCols) {
   std::ofstream outfile;
   outfile.open(name + ".txt", std::ios::out);
@@ -349,8 +364,8 @@ static void PrintHalfMatrixCPU(const T *mat, int num, std::string name,
 }
 
 template <typename T, typename outputT>
-void naive_matmul(const T *a, const T *b, outputT *c, size_t m, size_t n,
-                  size_t k) {
+void naive_matmul(
+    const T *a, const T *b, outputT *c, size_t m, size_t n, size_t k) {
   for (int ik = 0; ik < k; ik++) {
     for (int im = 0; im < m; im++) {
       for (int in = 0; in < n; in++) {
@@ -361,13 +376,17 @@ void naive_matmul(const T *a, const T *b, outputT *c, size_t m, size_t n,
 }
 
 template <typename T, typename outputT, typename ScaleType = uint16_t>
-void naive_matmul_fused_dequantize_nf4(const T *a, const T *b,
+void naive_matmul_fused_dequantize_nf4(const T *a,
+                                       const T *b,
                                        const ScaleType *col_scale,
                                        const ScaleType *row_scale,
                                        const int32_t *nf4_look_up_table,
-                                       outputT *c, size_t num_experts,
+                                       outputT *c,
+                                       size_t num_experts,
                                        int64_t *total_rows_before_experts,
-                                       size_t total_rows, size_t n, size_t k) {
+                                       size_t total_rows,
+                                       size_t n,
+                                       size_t k) {
   // PrintMatrixCPU<T>(
   //     a, total_rows * k, "naive_matmul_a", k);
   // PrintMatrixCPU<T>(
@@ -442,10 +461,15 @@ void naive_matmul_fused_dequantize_nf4(const T *a, const T *b,
 }
 
 // Author (zhengzekang): we use float to monitor half matmul in CPU.
-void CheckHalfDiff(int16_t *device_res, float *host_result, size_t elem_cnt,
-                   float atol, float rtol) {
+void CheckHalfDiff(int16_t *device_res,
+                   float *host_result,
+                   size_t elem_cnt,
+                   float atol,
+                   float rtol) {
   std::vector<int16_t> device_data(elem_cnt);
-  cudaMemcpy(device_data.data(), device_res, sizeof(int16_t) * elem_cnt,
+  cudaMemcpy(device_data.data(),
+             device_res,
+             sizeof(int16_t) * elem_cnt,
              cudaMemcpyDeviceToHost);
 
   for (size_t i = 0; i < elem_cnt; i++) {
@@ -459,7 +483,10 @@ void CheckHalfDiff(int16_t *device_res, float *host_result, size_t elem_cnt,
       printf(
           "Here in Idx: %d, CUDA result is: %f, Host result is: %f, absolute "
           "diff val is: %f \n",
-          i, device_res_val, host_res_val, absolute_diff);
+          i,
+          device_res_val,
+          host_res_val,
+          absolute_diff);
       return;
     }
   }
@@ -508,11 +535,11 @@ CutlassGemmConfig GetGemmConfig(int token_nums,
   // gemm_config_tuple:[m,n,k,tile_config,split_k_style,split_k_factor,stages]
   for (int i = 0; i < len_of_gemm_config_tuple; i += 7) {
     gemm_config.tile_config =
-        CutlassTileConfig(gemm_config_tuple[i + 3]); // tile_config
+        CutlassTileConfig(gemm_config_tuple[i + 3]);  // tile_config
     gemm_config.split_k_style =
-        SplitKStyle(gemm_config_tuple[i + 4]);             // split_k_style
-    gemm_config.split_k_factor = gemm_config_tuple[i + 5]; // split_k_factor
-    gemm_config.stages = gemm_config_tuple[i + 6];         // stages
+        SplitKStyle(gemm_config_tuple[i + 4]);              // split_k_style
+    gemm_config.split_k_factor = gemm_config_tuple[i + 5];  // split_k_factor
+    gemm_config.stages = gemm_config_tuple[i + 6];          // stages
     // make sure we have at least one tuned config
     if (token_nums <= gemm_config_tuple[i + 0]) {
       break;
@@ -522,7 +549,8 @@ CutlassGemmConfig GetGemmConfig(int token_nums,
 }
 
 template <typename T, typename U = T>
-void get_tensor_from_file(const std::string file_path, int64_t numel,
+void get_tensor_from_file(const std::string file_path,
+                          int64_t numel,
                           T *tensor_ptr) {
   std::fstream datafile;
   datafile.open(file_path, std::ios_base::in | std::ios_base::out);
@@ -641,7 +669,7 @@ int main(int argc, char *argv[]) {
   auto mixed_gemm_runner = W4A8MoeGemmRunner<half, int8_t, cutlass::uint4b_t>();
 
   // int mixgemm_max_size = std::max(m, k);
-  int mixgemm_workspace_size_bytes = 1 * 1024 * 1024 * 1024; // 1G workspace
+  int mixgemm_workspace_size_bytes = 1 * 1024 * 1024 * 1024;  // 1G workspace
   std::cout << "mixgemm_workspace_size_bytes: " << mixgemm_workspace_size_bytes
             << std::endl;
   char *mixgemm_workspace_data;
@@ -663,8 +691,8 @@ int main(int argc, char *argv[]) {
       }
     } else {
       std::cout << "get a data from: " << a_data_file << std::endl;
-      get_tensor_from_file<int8_t, int32_t>(a_data_file, total_rows * k,
-                                            a_int.data());
+      get_tensor_from_file<int8_t, int32_t>(
+          a_data_file, total_rows * k, a_int.data());
     }
     // PrintMatrixCPU<int8_t>(a_int.data(),total_rows*k,"a_int8_cpu",n);
   }
@@ -752,7 +780,8 @@ int main(int argc, char *argv[]) {
         // PrintMatrixCPU_int4(packed_b_int.data(),num_experts*k*n,"w4a8_packed_b_int4",n);
         permute_B_rows_for_mixed_gemm_int4<4>(
             b_int_processed.data() + ie * k * n / 2,
-            packed_b_int.data() + ie * k * n / 2, std::vector<size_t>{k, n},
+            packed_b_int.data() + ie * k * n / 2,
+            std::vector<size_t>{k, n},
             (int64_t)80);
 
         // PrintMatrixCPU_int4(b_int_processed.data(),num_experts*k*n,"w4a8_permuted_int4",n);
@@ -820,8 +849,8 @@ int main(int argc, char *argv[]) {
         // }
       }
     } else {
-      get_tensor_from_file<float>(row_scale_data_file, total_rows,
-                                  row_scale_float.data());
+      get_tensor_from_file<float>(
+          row_scale_data_file, total_rows, row_scale_float.data());
     }
     // PrintMatrixCPU<float>(row_scale_float.data(),total_rows,"row_scale_float_cpu",total_rows);
   }
@@ -839,8 +868,8 @@ int main(int argc, char *argv[]) {
         // }
       }
     } else {
-      get_tensor_from_file<float>(col_scale_data_file, num_experts * n,
-                                  col_scale_float.data());
+      get_tensor_from_file<float>(
+          col_scale_data_file, num_experts * n, col_scale_float.data());
     }
     // PrintMatrixCPU<float>(col_scale_float.data(),num_experts*n,"col_scale_float_cpu",n);
   }
@@ -878,21 +907,35 @@ int main(int argc, char *argv[]) {
   cudaMalloc(&d_col_scale_half, num_experts * n * sizeof(uint16_t));
   cudaMalloc(&d_nf4_look_up_table, 4 * sizeof(uint32_t));
   cudaMalloc(&d_total_rows_before_experts, num_experts * sizeof(int64_t));
-  cudaMemcpy(d_a_int, a_int.data(), total_rows * k * sizeof(int8_t),
+  cudaMemcpy(d_a_int,
+             a_int.data(),
+             total_rows * k * sizeof(int8_t),
+             cudaMemcpyHostToDevice);
+  cudaMemcpy(d_b_int,
+             b_int_processed_3.data(),
+             num_experts * k * n / 2 * sizeof(int8_t),
              cudaMemcpyHostToDevice);
-  cudaMemcpy(d_b_int, b_int_processed_3.data(),
-             num_experts * k * n / 2 * sizeof(int8_t), cudaMemcpyHostToDevice);
 
-  cudaMemcpy(d_row_scale_half, row_scale_half.data(),
-             total_rows * sizeof(uint16_t), cudaMemcpyHostToDevice);
-  cudaMemcpy(d_col_scale_half, col_scale_half.data(),
-             num_experts * n * sizeof(uint16_t), cudaMemcpyHostToDevice);
-  cudaMemcpy(d_nf4_look_up_table, nf4_look_up_table_compress.data(),
-             4 * sizeof(uint32_t), cudaMemcpyHostToDevice);
-  cudaMemcpy(d_c_int, c_half.data(), total_rows * n * sizeof(uint16_t),
+  cudaMemcpy(d_row_scale_half,
+             row_scale_half.data(),
+             total_rows * sizeof(uint16_t),
+             cudaMemcpyHostToDevice);
+  cudaMemcpy(d_col_scale_half,
+             col_scale_half.data(),
+             num_experts * n * sizeof(uint16_t),
+             cudaMemcpyHostToDevice);
+  cudaMemcpy(d_nf4_look_up_table,
+             nf4_look_up_table_compress.data(),
+             4 * sizeof(uint32_t),
+             cudaMemcpyHostToDevice);
+  cudaMemcpy(d_c_int,
+             c_half.data(),
+             total_rows * n * sizeof(uint16_t),
+             cudaMemcpyHostToDevice);
+  cudaMemcpy(d_total_rows_before_experts,
+             total_rows_before_experts.data(),
+             num_experts * sizeof(int64_t),
              cudaMemcpyHostToDevice);
-  cudaMemcpy(d_total_rows_before_experts, total_rows_before_experts.data(),
-             num_experts * sizeof(int64_t), cudaMemcpyHostToDevice);
 
   cudaDeviceSynchronize();
   cudaError_t result = cudaGetLastError();
@@ -908,7 +951,9 @@ int main(int argc, char *argv[]) {
   std::cout << "=== do warm up for " << kWarmTime << " times" << std::endl;
   auto test_config =
       CutlassGemmConfig{CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64,
-                        SplitKStyle::NO_SPLIT_K, 1, 5};
+                        SplitKStyle::NO_SPLIT_K,
+                        1,
+                        5};
   std::cout << "=== do warm up end" << std::endl;
   for (int i = 0; i < kWarmTime; i++) {
     printf("warm up %d\n", i);
@@ -920,9 +965,16 @@ int main(int argc, char *argv[]) {
         reinterpret_cast<const half *>(d_row_scale_half),
         reinterpret_cast<const int32_t *>(d_nf4_look_up_table),
         reinterpret_cast<half *>(d_c_int),
-        reinterpret_cast<int64_t *>(d_total_rows_before_experts), -1,
-        total_rows, n, k, mixgemm_workspace_data, mixgemm_workspace_size_bytes,
-        num_experts, 0, test_config);
+        reinterpret_cast<int64_t *>(d_total_rows_before_experts),
+        -1,
+        total_rows,
+        n,
+        k,
+        mixgemm_workspace_data,
+        mixgemm_workspace_size_bytes,
+        num_experts,
+        0,
+        test_config);
     cudaError_t err = cudaGetLastError();
     if (err != cudaSuccess) {
       std::cout << "error: " << cudaGetErrorString(err) << std::endl;
@@ -1012,7 +1064,6 @@ int main(int argc, char *argv[]) {
             cudaEventCreate(&end);
             cudaEventRecord(begin, 0);
             for (int i = 0; i < kTestTime; ++i) {
-
               mixed_gemm_runner.moe_gemm(
                   reinterpret_cast<const int8_t *>(d_a_int),
                   reinterpret_cast<const cutlass::uint4b_t *>((void *)d_b_int),
@@ -1021,9 +1072,15 @@ int main(int argc, char *argv[]) {
                   reinterpret_cast<const half *>(d_row_scale_half),
                   reinterpret_cast<const int32_t *>(d_nf4_look_up_table),
                   reinterpret_cast<half *>(d_c_int),
-                  reinterpret_cast<int64_t *>(d_total_rows_before_experts), -1,
-                  total_rows, n, k, mixgemm_workspace_data,
-                  mixgemm_workspace_size_bytes, num_experts, 0,
+                  reinterpret_cast<int64_t *>(d_total_rows_before_experts),
+                  -1,
+                  total_rows,
+                  n,
+                  k,
+                  mixgemm_workspace_data,
+                  mixgemm_workspace_size_bytes,
+                  num_experts,
+                  0,
                   test_gemm_config);
             }
             cudaEventRecord(end, 0);
@@ -1149,37 +1206,58 @@ int main(int argc, char *argv[]) {
   if (do_check) {
     std::cout << "=== do accuracy check " << std::endl;
     cudaMemset(d_c_int, 0, total_rows * n * sizeof(uint16_t));
-    PrintHalfMatrix(static_cast<int16_t *>(d_c_int), total_rows * n,
-                    "CUDA_c_dequantize_fp16_output_before_gemm", n);
+    PrintHalfMatrix(static_cast<int16_t *>(d_c_int),
+                    total_rows * n,
+                    "CUDA_c_dequantize_fp16_output_before_gemm",
+                    n);
 
     mixed_gemm_runner.moe_gemm(
         reinterpret_cast<const int8_t *>(d_a_int),
         reinterpret_cast<const cutlass::uint4b_t *>((void *)d_b_int),
         cutlass::epilogue::QuantMode::PerChannelQuant,
         reinterpret_cast<const half *>(d_col_scale_half),
-        nullptr, // reinterpret_cast<const half*>(d_row_scale_half),
-        nullptr, // reinterpret_cast<const int32_t*>(d_nf4_look_up_table),
+        nullptr,  // reinterpret_cast<const half*>(d_row_scale_half),
+        nullptr,  // reinterpret_cast<const int32_t*>(d_nf4_look_up_table),
         reinterpret_cast<half *>(d_c_int),
-        reinterpret_cast<int64_t *>(d_total_rows_before_experts), -1,
-        total_rows, n, k, mixgemm_workspace_data, mixgemm_workspace_size_bytes,
-        num_experts, 0);
+        reinterpret_cast<int64_t *>(d_total_rows_before_experts),
+        -1,
+        total_rows,
+        n,
+        k,
+        mixgemm_workspace_data,
+        mixgemm_workspace_size_bytes,
+        num_experts,
+        0);
     cudaDeviceSynchronize();
     // PrintMatrix<int32_t>(reinterpret_cast<const
     // int32_t*>(d_nf4_look_up_table),4,"d_nf4_look_up_table",1);
     printf("##### d_nf4_look_up_table address: %p \n", d_nf4_look_up_table);
 
     naive_matmul_fused_dequantize_nf4<int8_t, float, float>(
-        a_int.data(), b_int.data(), col_scale_float.data(),
-        nullptr, // row_scale_float.data(),
-        nullptr, // nf4_look_up_table.data(),
-        c_float.data(), num_experts, total_rows_before_experts.data(),
-        total_rows, n, k);
-    PrintMatrixCPU<float>(c_float.data(), total_rows * n,
-                          "CPU_c_fake_fp16_dequantize_output_base", n);
-    PrintHalfMatrix(static_cast<int16_t *>(d_c_int), total_rows * n,
-                    "CUDA_c_dequantize_fp16_output", n);
-    CheckHalfDiff(static_cast<int16_t *>(d_c_int), c_float.data(),
-                  total_rows * n, 1e-4, 1e-2);
+        a_int.data(),
+        b_int.data(),
+        col_scale_float.data(),
+        nullptr,  // row_scale_float.data(),
+        nullptr,  // nf4_look_up_table.data(),
+        c_float.data(),
+        num_experts,
+        total_rows_before_experts.data(),
+        total_rows,
+        n,
+        k);
+    PrintMatrixCPU<float>(c_float.data(),
+                          total_rows * n,
+                          "CPU_c_fake_fp16_dequantize_output_base",
+                          n);
+    PrintHalfMatrix(static_cast<int16_t *>(d_c_int),
+                    total_rows * n,
+                    "CUDA_c_dequantize_fp16_output",
+                    n);
+    CheckHalfDiff(static_cast<int16_t *>(d_c_int),
+                  c_float.data(),
+                  total_rows * n,
+                  1e-4,
+                  1e-2);
   }
 
   // if(kTestTime > 0){
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_with_epilogue_visitor.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_with_epilogue_visitor.h
index 648a21a353..804dcfdf91 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_with_epilogue_visitor.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/w4a8_moe_gemm_with_epilogue_visitor.h
@@ -33,10 +33,9 @@ namespace gemm {
 namespace kernel {
 
 template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
-          typename Epilogue_,           ///! Epilogue
+          typename Epilogue_,            ///! Epilogue
           typename ThreadblockSwizzle_,  ///! Threadblock swizzling function
-          GroupScheduleMode GroupScheduleMode_
-          >
+          GroupScheduleMode GroupScheduleMode_>
 struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
  public:
   using Mma = Mma_;
@@ -98,10 +97,10 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
   static bool const kTransposed = false;
 
   using ProblemVisitor = GemmMoeProblemVisitor<ThreadblockShape,
-                                              kGroupScheduleMode,
-                                              kThreadCount,
-                                              kThreadCount,
-                                              kTransposed>;
+                                               kGroupScheduleMode,
+                                               kThreadCount,
+                                               kThreadCount,
+                                               kTransposed>;
 
   //
   // Structures
@@ -206,7 +205,6 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
                                       ElementC,
                                       LayoutA,
                                       LayoutB> {
-
     using ParamsBase = UniversalParamsBase<ThreadblockSwizzle,
                                            ThreadblockShape,
                                            ElementA,
@@ -226,7 +224,6 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
     typename EpilogueVisitor::OutputTileIterator::Params params_C;
     typename EpilogueVisitor::OutputTileIterator::Params params_D;
 
-
     void* ptr_A;
     void* ptr_B;
     epilogue::QuantMode quant_mode;
@@ -248,7 +245,11 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
     CUTLASS_HOST_DEVICE
     Params() = default;
 
-    Params(Arguments const& args, int device_sms, int sm_occupancy, void* workspace = nullptr, int tile_count=0)
+    Params(Arguments const& args,
+           int device_sms,
+           int sm_occupancy,
+           void* workspace = nullptr,
+           int tile_count = 0)
         : ParamsBase(args, device_sms, sm_occupancy),
           problem_visitor(args.total_rows_before_expert,
                           args.total_rows,
@@ -274,7 +275,7 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
           ptr_D(args.ref_D.data()),
           batch_stride_A(args.batch_stride_A),
           batch_stride_B(args.batch_stride_B),
-          epilogue_visitor(args.epilogue_visitor) { }
+          epilogue_visitor(args.epilogue_visitor) {}
   };
 
   /// Shared memory storage structure
@@ -298,7 +299,8 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
 
   /// Determines whether kernel satisfies alignment
   static Status can_implement(cutlass::gemm::GemmCoord const& problem_size) {
-    CUTLASS_TRACE_HOST("MoeW4A8GemmWithEpilogueVisitorInterleavedNf4::can_implement()");
+    CUTLASS_TRACE_HOST(
+        "MoeW4A8GemmWithEpilogueVisitorInterleavedNf4::can_implement()");
 
     static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
     static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
@@ -332,7 +334,7 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
     }
 
     if (platform::is_same<LayoutC, layout::RowMajor>::value) {
-         isCMisaligned  = problem_size.n() % kAlignmentC;
+      isCMisaligned = problem_size.n() % kAlignmentC;
     } else if (platform::is_same<LayoutC, layout::ColumnMajor>::value) {
       isCMisaligned = problem_size.m() % kAlignmentC;
     } else if (platform::is_same<LayoutC,
@@ -381,89 +383,92 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
   /// Executes one GEMM
   CUTLASS_DEVICE
   void operator()(Params const& params, SharedStorage& shared_storage) {
+    using ElementA = typename Mma::IteratorA::Element;
+    using LayoutA = typename Mma::IteratorA::Layout;
+    using ElementB = typename Mma::IteratorB::Element;
+    using LayoutB = typename Mma::IteratorB::Layout;
 
-      using ElementA = typename Mma::IteratorA::Element;
-      using LayoutA = typename Mma::IteratorA::Layout;
-      using ElementB = typename Mma::IteratorB::Element;
-      using LayoutB = typename Mma::IteratorB::Layout;
+    static constexpr int kInterleave =
+        Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
+    static_assert(platform::is_same<LayoutB, layout::RowMajor>::value &&
+                          kInterleave == 1 ||
+                      platform::is_same<LayoutB, layout::ColumnMajor>::value &&
+                          kInterleave >= 1,
+                  "B must be row major/col major OR col major interleaved.");
 
-      static constexpr int kInterleave =
-          Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
-      static_assert(
-          platform::is_same<LayoutB, layout::RowMajor>::value &&
-                  kInterleave == 1 ||
-              platform::is_same<LayoutB, layout::ColumnMajor>::value &&
-                  kInterleave >= 1,
-          "B must be row major/col major OR col major interleaved.");
+    //
+    // Problem visitor.
+    //
+    ProblemVisitor problem_visitor(
+        params.problem_visitor, shared_storage.problem_visitor, blockIdx.x);
+    const int64_t gemm_k = params.problem_visitor.gemm_k;
+    const int64_t gemm_n = params.problem_visitor.gemm_n;
+    int64_t bytes_per_expert_matrix =
+        (gemm_k * gemm_n / 8) * cutlass::sizeof_bits<ElementB>::value;
 
-      //
-      // Problem visitor.
-      //
-      ProblemVisitor problem_visitor(
-          params.problem_visitor, shared_storage.problem_visitor, blockIdx.x);
-      const int64_t gemm_k = params.problem_visitor.gemm_k;
-      const int64_t gemm_n = params.problem_visitor.gemm_n;
-      int64_t bytes_per_expert_matrix =
-          (gemm_k * gemm_n / 8) * cutlass::sizeof_bits<ElementB>::value;
-
-      // Outer 'persistent' loop to iterate over tiles
-      while (problem_visitor.next_tile()) {
+    // Outer 'persistent' loop to iterate over tiles
+    while (problem_visitor.next_tile()) {
       // // Compute threadblock location
       ThreadblockSwizzle threadblock_swizzle;
 
       cutlass::gemm::GemmCoord threadblock_tile_offset =
           threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-        GemmCoord problem_size = problem_visitor.problem_size();
-        int32_t problem_idx = problem_visitor.problem_index();
-        int32_t cta_idx = int32_t(problem_visitor.threadblock_idx());
+      GemmCoord problem_size = problem_visitor.problem_size();
+      int32_t problem_idx = problem_visitor.problem_index();
+      int32_t cta_idx = int32_t(problem_visitor.threadblock_idx());
 
-        GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
+      GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
 
-        cutlass::MatrixCoord threadblock_offset(
-            int(cta_idx / grid_shape.n()) * Mma::Shape::kM,  // NOLINT
-            int(cta_idx % grid_shape.n()) * Mma::Shape::kN  // NOLINT
-            );
+      cutlass::MatrixCoord threadblock_offset(
+          int(cta_idx / grid_shape.n()) * Mma::Shape::kM,  // NOLINT
+          int(cta_idx % grid_shape.n()) * Mma::Shape::kN   // NOLINT
+      );
 
-        // if (threadIdx.x == 0) {
-        //   printf("%d-%d-%d problem_size: %d, %d problem_idx: %d, cta_idx: %d\n", blockIdx.x,blockIdx.y,blockIdx.z, problem_size.m(), problem_size.n(), problem_idx, cta_idx);
-        // }
+      // if (threadIdx.x == 0) {
+      //   printf("%d-%d-%d problem_size: %d, %d problem_idx: %d, cta_idx:
+      //   %d\n", blockIdx.x,blockIdx.y,blockIdx.z, problem_size.m(),
+      //   problem_size.n(), problem_idx, cta_idx);
+      // }
 
-        // Load element pointers. Exchange pointers and strides if working on
-        // the transpose
-        int64_t rows_to_jump = 0;
-        if (params.problem_visitor.total_rows < 0) {
-          rows_to_jump = problem_idx == 0 ? 0 : params.problem_visitor.last_row_for_problem[problem_idx - 1];
-        } else {
-          rows_to_jump = problem_idx * (params.problem_visitor.total_rows / params.problem_visitor.problem_count);
-        }
+      // Load element pointers. Exchange pointers and strides if working on
+      // the transpose
+      int64_t rows_to_jump = 0;
+      if (params.problem_visitor.total_rows < 0) {
+        rows_to_jump =
+            problem_idx == 0
+                ? 0
+                : params.problem_visitor.last_row_for_problem[problem_idx - 1];
+      } else {
+        rows_to_jump = problem_idx * (params.problem_visitor.total_rows /
+                                      params.problem_visitor.problem_count);
+      }
 
-        ElementA* ptr_A =
-            reinterpret_cast<ElementA*>(params.ptr_A) + rows_to_jump * gemm_k;
-        typename LayoutA::LongIndex ldm_A = gemm_k;
-
-        char* byte_ptr_B = ((char*)params.ptr_B) +                 // NOLINT
-                           problem_idx * bytes_per_expert_matrix;  // NOLINT
-        ElementB* ptr_B = reinterpret_cast<ElementB*>(byte_ptr_B);
-        typename LayoutB::LongIndex ldm_B =
-            platform::is_same<layout::RowMajor, LayoutB>::value
-                ? gemm_n
-                : gemm_k * kInterleave;
+      ElementA* ptr_A =
+          reinterpret_cast<ElementA*>(params.ptr_A) + rows_to_jump * gemm_k;
+      typename LayoutA::LongIndex ldm_A = gemm_k;
 
+      char* byte_ptr_B = ((char*)params.ptr_B) +                 // NOLINT
+                         problem_idx * bytes_per_expert_matrix;  // NOLINT
+      ElementB* ptr_B = reinterpret_cast<ElementB*>(byte_ptr_B);
+      typename LayoutB::LongIndex ldm_B =
+          platform::is_same<layout::RowMajor, LayoutB>::value
+              ? gemm_n
+              : gemm_k * kInterleave;
 
       int offset_k = 0;
       int problem_size_k = params.problem_size.k();
 
-
-  // Maybe need to modify? Author zhengzekang.
-  #if SPLIT_K_ENABLED
+// Maybe need to modify? Author zhengzekang.
+#if SPLIT_K_ENABLED
       //
       // Fetch pointers based on mode.
       //
       if (params.mode == GemmUniversalMode::kGemm ||
           params.mode == GemmUniversalMode::kGemmSplitKParallel) {
         if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
-          problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
+          problem_size_k =
+              (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
         }
         offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
       } else if (params.mode == GemmUniversalMode::kBatched) {
@@ -475,17 +480,19 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
         ptr_B = static_cast<ElementB* const*>(
             params.ptr_B)[threadblock_tile_offset.k()];
       }
-  #endif
-    // if(threadIdx.x==0){
-    //     printf("##### block: %d-%d-%d, offset_k:%d, threadblock_tile_offset.m-n-k():%d-%d-%d, params.gemm_k_size:%d \n",
-    //             blockIdx.x, blockIdx.y, blockIdx.z,
-    //             offset_k,
-    //             threadblock_tile_offset.m(),
-    //             threadblock_tile_offset.n(),
-    //             threadblock_tile_offset.k(),
-    //             params.gemm_k_size
-    //             );
-    //   }
+#endif
+      // if(threadIdx.x==0){
+      //     printf("##### block: %d-%d-%d, offset_k:%d,
+      //     threadblock_tile_offset.m-n-k():%d-%d-%d, params.gemm_k_size:%d
+      //     \n",
+      //             blockIdx.x, blockIdx.y, blockIdx.z,
+      //             offset_k,
+      //             threadblock_tile_offset.m(),
+      //             threadblock_tile_offset.n(),
+      //             threadblock_tile_offset.k(),
+      //             params.gemm_k_size
+      //             );
+      //   }
 
       // Compute initial location in logical coordinates
       cutlass::MatrixCoord tb_offset_A{
@@ -493,21 +500,18 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
           0,
       };
 
-
       cutlass::MatrixCoord tb_offset_B{
-          0,
-          threadblock_offset.column() / kInterleave};
+          0, threadblock_offset.column() / kInterleave};
 
       // Compute position within threadblock
       int thread_idx = threadIdx.x;
 
       // Construct iterators to A and B operands
-      typename Mma::IteratorA iterator_A(
-          params.params_A,
-          ptr_A,
-          {problem_size.m(), problem_size_k},
-          thread_idx,
-          tb_offset_A);
+      typename Mma::IteratorA iterator_A(params.params_A,
+                                         ptr_A,
+                                         {problem_size.m(), problem_size_k},
+                                         thread_idx,
+                                         tb_offset_A);
 
       typename Mma::IteratorB iterator_B(
           params.params_B,
@@ -516,13 +520,11 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
           thread_idx,
           tb_offset_B);
       typename Mma::IteratorNF4LookUpTable iterator_nf4_look_up_table =
-        Mma::IteratorNF4LookUpTable(
-          params.params_nf4_look_up_table,
-        params.ref_nf4_look_up_table.data(),
-          {0,16},
-          threadIdx.x,
-          {0,0}
-        );
+          Mma::IteratorNF4LookUpTable(params.params_nf4_look_up_table,
+                                      params.ref_nf4_look_up_table.data(),
+                                      {0, 16},
+                                      threadIdx.x,
+                                      {0, 0});
 
       // Broadcast the warp_id computed by lane 0 to ensure dependent code
       // is compiled as warp-uniform.
@@ -542,9 +544,14 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
 
       // Compute threadblock-scoped matrix multiply-add
       int gemm_k_iterations =
-          (problem_size_k  + Mma::Shape::kK - 1) / Mma::Shape::kK;
+          (problem_size_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
       // Compute threadblock-scoped matrix multiply-add
-      mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, iterator_nf4_look_up_table, accumulators);
+      mma(gemm_k_iterations,
+          accumulators,
+          iterator_A,
+          iterator_B,
+          iterator_nf4_look_up_table,
+          accumulators);
       // if(threadIdx.x==0){
       //   printf("##### block: %d-%d-%d, offset-m-n-k:%d-%d-%d \n",
       //           blockIdx.x, blockIdx.y, blockIdx.z,
@@ -560,44 +567,54 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
       threadblock_tile_offset =
           threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-        ElementC* ptr_C =
-            reinterpret_cast<ElementC*>(params.ptr_C) + rows_to_jump * gemm_n;
-        ElementC* ptr_D =
-            reinterpret_cast<ElementC*>(params.ptr_D) + rows_to_jump * gemm_n;
+      ElementC* ptr_C =
+          reinterpret_cast<ElementC*>(params.ptr_C) + rows_to_jump * gemm_n;
+      ElementC* ptr_D =
+          reinterpret_cast<ElementC*>(params.ptr_D) + rows_to_jump * gemm_n;
 
-      using Element_scale = typename EpilogueVisitor::ScaleTileIterator::Element;
-      Element_scale* ptr_alpha_row = params.ptr_alpha_row == nullptr ? params.ptr_alpha_row : reinterpret_cast<Element_scale*>(params.ptr_alpha_row) + rows_to_jump;
-      Element_scale* ptr_alpha_col = reinterpret_cast<Element_scale*>(params.ptr_alpha_col) + problem_idx * params.problem_size.n();
+      using Element_scale =
+          typename EpilogueVisitor::ScaleTileIterator::Element;
+      Element_scale* ptr_alpha_row =
+          params.ptr_alpha_row == nullptr
+              ? params.ptr_alpha_row
+              : reinterpret_cast<Element_scale*>(params.ptr_alpha_row) +
+                    rows_to_jump;
+      Element_scale* ptr_alpha_col =
+          reinterpret_cast<Element_scale*>(params.ptr_alpha_col) +
+          problem_idx * params.problem_size.n();
       // if (threadIdx.x == 0)
-      // printf("##### block: %d-%d-%d, ptr_alpha_row:%p,(%f) ptr_alpha_col:%p,(%f)\n", blockIdx.x, blockIdx.y, blockIdx.z, ptr_alpha_row, static_cast<float>(*ptr_alpha_row), ptr_alpha_col, static_cast<float>(*ptr_alpha_col));
+      // printf("##### block: %d-%d-%d, ptr_alpha_row:%p,(%f)
+      // ptr_alpha_col:%p,(%f)\n", blockIdx.x, blockIdx.y, blockIdx.z,
+      // ptr_alpha_row, static_cast<float>(*ptr_alpha_row), ptr_alpha_col,
+      // static_cast<float>(*ptr_alpha_col));
       //
       // Construct the epilogue visitor
       //
 
       EpilogueVisitor epilogue_visitor(params.epilogue_visitor,
-                                     shared_storage.epilogue.visitor,
-                                     problem_size.mn(),
-                                     thread_idx,
-                                     warp_idx,
-                                     lane_idx,
-                                     params.params_alpha_col,
-                                     params.params_C,
-                                     params.params_D,
-                                     params.quant_mode,
-                                     ptr_alpha_row,
-                                     ptr_alpha_col,
-                                     ptr_C,
-                                     ptr_D,
-                                     threadblock_offset,
-                                     blockIdx.y * params.problem_size.m());
+                                       shared_storage.epilogue.visitor,
+                                       problem_size.mn(),
+                                       thread_idx,
+                                       warp_idx,
+                                       lane_idx,
+                                       params.params_alpha_col,
+                                       params.params_C,
+                                       params.params_D,
+                                       params.quant_mode,
+                                       ptr_alpha_row,
+                                       ptr_alpha_col,
+                                       ptr_C,
+                                       ptr_D,
+                                       threadblock_offset,
+                                       blockIdx.y * params.problem_size.m());
 
       if (params.mode == GemmUniversalMode::kGemm) {
         // Indicate which position in a serial reduction the output operator is
         // currently updating
         epilogue_visitor.set_k_partition(threadblock_tile_offset.k(),
-                                        params.grid_tiled_shape.k());
+                                         params.grid_tiled_shape.k());
       } else if (params.mode == GemmUniversalMode::kBatched ||
-                params.mode == GemmUniversalMode::kArray) {
+                 params.mode == GemmUniversalMode::kArray) {
         epilogue_visitor.set_batch_index(threadblock_tile_offset.k());
       }
 
@@ -608,9 +625,9 @@ struct MoeW4A8GemmWithEpilogueVisitorInterleavedNf4 {
       // Execute the epilogue operator to update the destination tensor.
       epilogue(epilogue_visitor, accumulators);
 
-        // Next tile
-        problem_visitor.advance(gridDim.x);
-      }
+      // Next tile
+      problem_visitor.advance(gridDim.x);
+    }
   }
 };
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/weight_process_utils.h b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/weight_process_utils.h
index 68779ba28c..dd6536c762 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/weight_process_utils.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/w4a8_moe/weight_process_utils.h
@@ -32,482 +32,496 @@ limitations under the License. */
 
 void row_major_to_column_major(int8_t* col_major_tensor,
                                const int8_t* row_major_tensor,
-                               const std::vector<size_t>& shape){
-    size_t m = shape[0];
-    size_t n = shape[1];
-    for(auto i=0;i<m*n;i++){
-        size_t im = i / n;
-        size_t in = i % n;
-        col_major_tensor[in*m+im] = row_major_tensor[im*n+in];
-    }
+                               const std::vector<size_t>& shape) {
+  size_t m = shape[0];
+  size_t n = shape[1];
+  for (auto i = 0; i < m * n; i++) {
+    size_t im = i / n;
+    size_t in = i % n;
+    col_major_tensor[in * m + im] = row_major_tensor[im * n + in];
+  }
 }
 
 void add_bias_and_interleave_int8s_inplace(int8_t* int8_tensor_ptr,
-                                           int64_t num_elts)
-{
-    int8_t* int8_tensor = reinterpret_cast<int8_t *>(int8_tensor_ptr);
-    for (int ii = 0; ii < num_elts; ++ii) {
-        int8_tensor[ii] = int8_t(int(int8_tensor[ii]) + 128);
-        // int8_tensor[ii] = int8_t(int(int8_tensor[ii]));
-    }
+                                           int64_t num_elts) {
+  int8_t* int8_tensor = reinterpret_cast<int8_t*>(int8_tensor_ptr);
+  for (int ii = 0; ii < num_elts; ++ii) {
+    int8_tensor[ii] = int8_t(int(int8_tensor[ii]) + 128);
+    // int8_tensor[ii] = int8_t(int(int8_tensor[ii]));
+  }
 
-    // Step 2 will transform the layout of a 32-bit register in CUDA in order to match the int4 layout. This has no
-    // performance benefit and is purely so that int4 and int8 have the same layout.
-    // Pictorially, this does the following:
-    // bit 32                                                      0
-    //      [elt_3  elt_2  elt_1  elt_0] (each elt occupies 8 bits)
-    //
-    // And it will rearrange the output 32 bit register to be the following:
-    // bit 32                                                      0
-    //      [elt_3  elt_1  elt_2  elt_0] (each elt occupies 8 bits)
+  // Step 2 will transform the layout of a 32-bit register in CUDA in order to
+  // match the int4 layout. This has no performance benefit and is purely so
+  // that int4 and int8 have the same layout. Pictorially, this does the
+  // following: bit 32                                                      0
+  //      [elt_3  elt_2  elt_1  elt_0] (each elt occupies 8 bits)
+  //
+  // And it will rearrange the output 32 bit register to be the following:
+  // bit 32                                                      0
+  //      [elt_3  elt_1  elt_2  elt_0] (each elt occupies 8 bits)
 
-    for (int64_t base = 0; base < num_elts; base += 4) {
-        std::swap(int8_tensor[base + 1], int8_tensor[base + 2]);
-    }
+  for (int64_t base = 0; base < num_elts; base += 4) {
+    std::swap(int8_tensor[base + 1], int8_tensor[base + 2]);
+  }
 }
 
+void subbyte_transpose_impl_int4(int8_t* transposed_quantized_tensor,
+                                 const int8_t* quantized_tensor,
+                                 const std::vector<size_t>& shape) {
+  const int bits_per_elt = 4;
 
-void subbyte_transpose_impl_int4(int8_t*                    transposed_quantized_tensor,
-                            const int8_t*              quantized_tensor,
-                            const std::vector<size_t>& shape)
-{
-    const int bits_per_elt = 4;
+  const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const size_t col_bytes = num_cols * bits_per_elt / 8;
+  const size_t col_bytes_trans = num_rows * bits_per_elt / 8;
+  const size_t num_bytes = size_t(num_experts) * num_rows * col_bytes;
 
-    const size_t col_bytes       = num_cols * bits_per_elt / 8;
-    const size_t col_bytes_trans = num_rows * bits_per_elt / 8;
-    const size_t num_bytes       = size_t(num_experts) * num_rows * col_bytes;
+  const uint8_t* input_byte_ptr =
+      reinterpret_cast<const uint8_t*>(quantized_tensor);
+  uint8_t* output_byte_ptr =
+      reinterpret_cast<uint8_t*>(transposed_quantized_tensor);
 
-    const uint8_t* input_byte_ptr  = reinterpret_cast<const uint8_t*>(quantized_tensor);
-    uint8_t*       output_byte_ptr = reinterpret_cast<uint8_t*>(transposed_quantized_tensor);
+  // static_assert(quant_type == QuantType::INT8_WEIGHT_ONLY || quant_type ==
+  // QuantType::PACKED_INT4_WEIGHT_ONLY, "");
+  static constexpr int ELTS_PER_BYTE = 2;
 
-    // static_assert(quant_type == QuantType::INT8_WEIGHT_ONLY || quant_type == QuantType::PACKED_INT4_WEIGHT_ONLY, "");
-    static constexpr int ELTS_PER_BYTE = 2;
+  static constexpr int M_TILE_L1 = 64;
+  static constexpr int N_TILE_L1 = M_TILE_L1 / ELTS_PER_BYTE;
+  uint8_t cache_buf[M_TILE_L1][N_TILE_L1];
 
-    static constexpr int M_TILE_L1 = 64;
-    static constexpr int N_TILE_L1 = M_TILE_L1 / ELTS_PER_BYTE;
-    uint8_t              cache_buf[M_TILE_L1][N_TILE_L1];
+  static constexpr int VECTOR_WIDTH = std::min(32, N_TILE_L1);
 
-    static constexpr int VECTOR_WIDTH = std::min(32, N_TILE_L1);
+  // We assume the dims are a multiple of vector width. Our kernels only handle
+  // dims which are multiples of 64 for weight-only quantization. As a result,
+  // this seemed like a reasonable tradeoff because it allows GCC to emit vector
+  // instructions.
 
-    // We assume the dims are a multiple of vector width. Our kernels only handle dims which are multiples
-    // of 64 for weight-only quantization. As a result, this seemed like a reasonable tradeoff because it
-    // allows GCC to emit vector instructions.
+  const int num_m_tiles = (num_rows + M_TILE_L1 - 1) / M_TILE_L1;
+  const int num_n_tiles = (col_bytes + N_TILE_L1 - 1) / N_TILE_L1;
 
-    const int num_m_tiles = (num_rows + M_TILE_L1 - 1) / M_TILE_L1;
-    const int num_n_tiles = (col_bytes + N_TILE_L1 - 1) / N_TILE_L1;
+  for (size_t expert = 0; expert < num_experts; ++expert) {
+    const size_t matrix_offset = expert * num_rows * col_bytes;
+    for (size_t row_tile_start = 0; row_tile_start < num_rows;
+         row_tile_start += M_TILE_L1) {
+      for (size_t col_tile_start_byte = 0; col_tile_start_byte < col_bytes;
+           col_tile_start_byte += N_TILE_L1) {
+        const int row_limit = std::min(row_tile_start + M_TILE_L1, num_rows);
+        const int col_limit =
+            std::min(col_tile_start_byte + N_TILE_L1, col_bytes);
 
-    for (size_t expert = 0; expert < num_experts; ++expert) {
-        const size_t matrix_offset = expert * num_rows * col_bytes;
-        for (size_t row_tile_start = 0; row_tile_start < num_rows; row_tile_start += M_TILE_L1) {
-            for (size_t col_tile_start_byte = 0; col_tile_start_byte < col_bytes; col_tile_start_byte += N_TILE_L1) {
+        for (int ii = 0; ii < M_TILE_L1; ++ii) {
+          const int row = row_tile_start + ii;
 
-                const int row_limit = std::min(row_tile_start + M_TILE_L1, num_rows);
-                const int col_limit = std::min(col_tile_start_byte + N_TILE_L1, col_bytes);
+          for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
+            const int col = col_tile_start_byte + jj;
 
-                for (int ii = 0; ii < M_TILE_L1; ++ii) {
-                    const int row = row_tile_start + ii;
+            const size_t logical_src_offset =
+                matrix_offset + row * col_bytes + col;
 
-                    for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
-                        const int col = col_tile_start_byte + jj;
-
-                        const size_t logical_src_offset = matrix_offset + row * col_bytes + col;
-
-                        if (row < row_limit && col < col_limit) {
-                            for (int v = 0; v < VECTOR_WIDTH; ++v) {
-                                cache_buf[ii][jj + v] = input_byte_ptr[logical_src_offset + v];
-                            }
-                        }
-                    }
-                }
-
-
-                for (int ii = 0; ii < M_TILE_L1; ++ii) {
-                    // Using M_TILE_L1 here is deliberate since we assume that the cache tile
-                    // is square in the number of elements (not necessarily the number of bytes).
-                    for (int jj = ii + 1; jj < M_TILE_L1; ++jj) {
-                        const int ii_byte       = ii / ELTS_PER_BYTE;
-                        const int ii_bit_offset = ii % ELTS_PER_BYTE;
-
-                        const int jj_byte       = jj / ELTS_PER_BYTE;
-                        const int jj_bit_offset = jj % ELTS_PER_BYTE;
-
-                        uint8_t src_elt = 0xF & (cache_buf[ii][jj_byte] >> (4 * jj_bit_offset));
-                        uint8_t tgt_elt = 0xF & (cache_buf[jj][ii_byte] >> (4 * ii_bit_offset));
-
-                        cache_buf[ii][jj_byte] &= (0xF0 >> (4 * jj_bit_offset));
-                        cache_buf[jj][ii_byte] &= (0xF0 >> (4 * ii_bit_offset));
-
-                        cache_buf[ii][jj_byte] |= (tgt_elt << (4 * jj_bit_offset));
-                        cache_buf[jj][ii_byte] |= (src_elt << (4 * ii_bit_offset));
-                    }
-                }
-
-
-                const size_t row_tile_start_trans      = col_tile_start_byte * ELTS_PER_BYTE;
-                const size_t col_tile_start_byte_trans = row_tile_start / ELTS_PER_BYTE;
-
-                const int row_limit_trans = std::min(row_tile_start_trans + M_TILE_L1, num_cols);
-                const int col_limit_trans = std::min(col_tile_start_byte_trans + N_TILE_L1, col_bytes_trans);
-
-                for (int ii = 0; ii < M_TILE_L1; ++ii) {
-                    const int row = row_tile_start_trans + ii;
-                    for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
-                        const int col = col_tile_start_byte_trans + jj;
-
-                        const size_t logical_tgt_offset = matrix_offset + row * col_bytes_trans + col;
-
-                        if (row < row_limit_trans && col < col_limit_trans) {
-                            for (int v = 0; v < VECTOR_WIDTH; ++v) {
-                                output_byte_ptr[logical_tgt_offset + v] = cache_buf[ii][jj + v];
-                            }
-                        }
-                    }
-                }
+            if (row < row_limit && col < col_limit) {
+              for (int v = 0; v < VECTOR_WIDTH; ++v) {
+                cache_buf[ii][jj + v] = input_byte_ptr[logical_src_offset + v];
+              }
             }
+          }
         }
+
+        for (int ii = 0; ii < M_TILE_L1; ++ii) {
+          // Using M_TILE_L1 here is deliberate since we assume that the cache
+          // tile is square in the number of elements (not necessarily the
+          // number of bytes).
+          for (int jj = ii + 1; jj < M_TILE_L1; ++jj) {
+            const int ii_byte = ii / ELTS_PER_BYTE;
+            const int ii_bit_offset = ii % ELTS_PER_BYTE;
+
+            const int jj_byte = jj / ELTS_PER_BYTE;
+            const int jj_bit_offset = jj % ELTS_PER_BYTE;
+
+            uint8_t src_elt =
+                0xF & (cache_buf[ii][jj_byte] >> (4 * jj_bit_offset));
+            uint8_t tgt_elt =
+                0xF & (cache_buf[jj][ii_byte] >> (4 * ii_bit_offset));
+
+            cache_buf[ii][jj_byte] &= (0xF0 >> (4 * jj_bit_offset));
+            cache_buf[jj][ii_byte] &= (0xF0 >> (4 * ii_bit_offset));
+
+            cache_buf[ii][jj_byte] |= (tgt_elt << (4 * jj_bit_offset));
+            cache_buf[jj][ii_byte] |= (src_elt << (4 * ii_bit_offset));
+          }
+        }
+
+        const size_t row_tile_start_trans = col_tile_start_byte * ELTS_PER_BYTE;
+        const size_t col_tile_start_byte_trans = row_tile_start / ELTS_PER_BYTE;
+
+        const int row_limit_trans =
+            std::min(row_tile_start_trans + M_TILE_L1, num_cols);
+        const int col_limit_trans =
+            std::min(col_tile_start_byte_trans + N_TILE_L1, col_bytes_trans);
+
+        for (int ii = 0; ii < M_TILE_L1; ++ii) {
+          const int row = row_tile_start_trans + ii;
+          for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
+            const int col = col_tile_start_byte_trans + jj;
+
+            const size_t logical_tgt_offset =
+                matrix_offset + row * col_bytes_trans + col;
+
+            if (row < row_limit_trans && col < col_limit_trans) {
+              for (int v = 0; v < VECTOR_WIDTH; ++v) {
+                output_byte_ptr[logical_tgt_offset + v] = cache_buf[ii][jj + v];
+              }
+            }
+          }
+        }
+      }
     }
+  }
 }
 
+void add_bias_and_interleave_int4s_inplace(int8_t* packed_int4_tensor,
+                                           const size_t num_elts) {
+  const int num_bytes = num_elts / 2;
 
-void add_bias_and_interleave_int4s_inplace(int8_t* packed_int4_tensor, const size_t num_elts)
-{
-    const int num_bytes = num_elts / 2;
+  // Step 1 will be to transform all the int4s to unsigned in order to make the
+  // dequantize take as little instructions as possible in the CUDA code.
+  for (size_t ii = 0; ii < num_bytes; ++ii) {
+    int8_t transformed_packed_int4s = 0;
+    // We don't need to mask in these ops since everything should be in the
+    // range 0-15
+    int8_t transformed_first_elt = (packed_int4_tensor[ii] & 0x0F);
+    int8_t transformed_second_elt = (packed_int4_tensor[ii] >> 4);
 
-    // Step 1 will be to transform all the int4s to unsigned in order to make the dequantize take as little
-    // instructions as possible in the CUDA code.
-    for (size_t ii = 0; ii < num_bytes; ++ii) {
-        int8_t transformed_packed_int4s = 0;
-        // We don't need to mask in these ops since everything should be in the range 0-15
-        int8_t transformed_first_elt = (packed_int4_tensor[ii] & 0x0F);
-        int8_t transformed_second_elt = (packed_int4_tensor[ii] >> 4);
+    transformed_packed_int4s |= transformed_first_elt;
+    transformed_packed_int4s |= (transformed_second_elt << 4);
+    packed_int4_tensor[ii] = transformed_packed_int4s;
+  }
 
-        transformed_packed_int4s |= transformed_first_elt;
-        transformed_packed_int4s |= (transformed_second_elt << 4);
-        packed_int4_tensor[ii] = transformed_packed_int4s;
-    }
-
-    // Step 2 will transform the layout of a 32-bit register in CUDA in order to minimize the number of shift & logical
-    // instructions That are needed to extract the int4s in the GEMM main loop. Pictorially, the loop below will do the
-    // following: Take as input a 32 bit register with layout: bit 32 0
-    //      [elt_7  elt_6  elt_5  elt_4  elt_3  elt_2  elt_1  elt_0] (each elt occupies 4 bits)
-    //
-    // And it will rearrange the output 32 bit register to be the following:
-    // bit 32                                                      0
-    //      [elt_7  elt_5  elt_3  elt_1  elt_6  elt_4  elt_2  elt_0] (each elt occupies 4 bits)
-
-    // FT_CHECK_WITH_INFO(num_bytes % 4 == 0, "Dimensions of int4 tensor must be a multiple of 8 for register relayout");
-    const size_t num_registers = num_bytes / 4;
-
-    uint32_t* register_ptr = reinterpret_cast<uint32_t*>(packed_int4_tensor);
-    for (size_t ii = 0; ii < num_registers; ++ii) {
-        const uint32_t current_register     = register_ptr[ii];
-        uint32_t       transformed_register = 0;
-
-        for (int dest_idx = 0; dest_idx < 8; ++dest_idx) {
-            const int src_idx    = dest_idx < 4 ? 2 * dest_idx : 2 * (dest_idx - 4) + 1;
-            const int src_shift  = 4 * src_idx;
-            const int dest_shift = 4 * dest_idx;
-
-            const uint32_t src_bits = (current_register >> src_shift) & 0xF;
-            transformed_register |= (src_bits << dest_shift);
-
-        }
-        register_ptr[ii] = transformed_register;
+  // Step 2 will transform the layout of a 32-bit register in CUDA in order to
+  // minimize the number of shift & logical instructions That are needed to
+  // extract the int4s in the GEMM main loop. Pictorially, the loop below will
+  // do the following: Take as input a 32 bit register with layout: bit 32 0
+  //      [elt_7  elt_6  elt_5  elt_4  elt_3  elt_2  elt_1  elt_0] (each elt
+  //      occupies 4 bits)
+  //
+  // And it will rearrange the output 32 bit register to be the following:
+  // bit 32                                                      0
+  //      [elt_7  elt_5  elt_3  elt_1  elt_6  elt_4  elt_2  elt_0] (each elt
+  //      occupies 4 bits)
+
+  // FT_CHECK_WITH_INFO(num_bytes % 4 == 0, "Dimensions of int4 tensor must be a
+  // multiple of 8 for register relayout");
+  const size_t num_registers = num_bytes / 4;
+
+  uint32_t* register_ptr = reinterpret_cast<uint32_t*>(packed_int4_tensor);
+  for (size_t ii = 0; ii < num_registers; ++ii) {
+    const uint32_t current_register = register_ptr[ii];
+    uint32_t transformed_register = 0;
+
+    for (int dest_idx = 0; dest_idx < 8; ++dest_idx) {
+      const int src_idx = dest_idx < 4 ? 2 * dest_idx : 2 * (dest_idx - 4) + 1;
+      const int src_shift = 4 * src_idx;
+      const int dest_shift = 4 * dest_idx;
+
+      const uint32_t src_bits = (current_register >> src_shift) & 0xF;
+      transformed_register |= (src_bits << dest_shift);
     }
+    register_ptr[ii] = transformed_register;
+  }
 }
 
-void permute_B_rows_for_mixed_and_int8_gemm(int8_t*                    permuted_quantized_tensor,
-                                            const int8_t*              quantized_tensor,
+void permute_B_rows_for_mixed_and_int8_gemm(int8_t* permuted_quantized_tensor,
+                                            const int8_t* quantized_tensor,
                                             const std::vector<size_t>& shape,
-                                            const int64_t              arch_version)
-{
+                                            const int64_t arch_version) {
+  // We only want to run this step for weight only quant.
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    // We only want to run this step for weight only quant.
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const int BITS_PER_ELT = 8;
+  const int K = 16 / BITS_PER_ELT;
+  const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;
+  const int ELTS_PER_REG = 32 / BITS_PER_ELT;
 
-    const int BITS_PER_ELT  = 8;
-    const int K             = 16 / BITS_PER_ELT;
-    const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;
-    const int ELTS_PER_REG  = 32 / BITS_PER_ELT;
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
+  int MMA_SHAPE_N = 8;
+  int B_ROWS_PER_MMA = 8 * K;
+  const int elts_in_int32 = 32 / BITS_PER_ELT;
 
-    int       MMA_SHAPE_N    = 8;
-    int       B_ROWS_PER_MMA = 8 * K;
-    const int elts_in_int32  = 32 / BITS_PER_ELT;
+  const int num_vec_cols = num_cols / elts_in_int32;
 
-    const int num_vec_cols = num_cols / elts_in_int32;
+  // The code is written as below so it works for both int8 and packed int4.
+  for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
+    for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+      for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
+        const int write_row = base_row + tile_row;
+        const int tile_read_row = 4 * (((tile_row % ELTS_PER_REG) / 2)) +
+                                  tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
 
-    // The code is written as below so it works for both int8 and packed int4.
-    for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
-        for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+        const int read_row = base_row + tile_read_row;
+        const int read_col = write_col;
 
-            for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
-                const int write_row = base_row + tile_row;
-                const int tile_read_row =
-                    4 * (((tile_row % ELTS_PER_REG) / 2)) + tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
+        const int64_t read_offset = int64_t(read_row) * num_vec_cols + read_col;
+        const int64_t write_offset =
+            int64_t(write_row) * num_vec_cols + write_col;
 
-                const int read_row = base_row + tile_read_row;
-                const int read_col = write_col;
-
-                const int64_t read_offset  = int64_t(read_row) * num_vec_cols + read_col;
-                const int64_t write_offset = int64_t(write_row) * num_vec_cols + write_col;
-
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
-        }
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
 
-// Permutes the rows of B for Turing and Ampere. Throws an error for other architectures.
-// The data is permuted such that:
-// For int8, each group of 16 rows is permuted using the map below:
+// Permutes the rows of B for Turing and Ampere. Throws an error for other
+// architectures. The data is permuted such that: For int8, each group of 16
+// rows is permuted using the map below:
 //  0 1 8 9 2 3 10 11 4 5 12 13 6 7 14 15
 //  0 1 2 3 4 5 6 7
-template<int bits=8>
-void permute_B_rows_for_mixed_gemm(int8_t*                    permuted_quantized_tensor,
-                                   const int8_t*              quantized_tensor,
+template <int bits = 8>
+void permute_B_rows_for_mixed_gemm(int8_t* permuted_quantized_tensor,
+                                   const int8_t* quantized_tensor,
                                    const std::vector<size_t>& shape,
-                                   const int64_t              arch_version)
-{
+                                   const int64_t arch_version) {
+  // We only want to run this step for weight only quant.
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    // We only want to run this step for weight only quant.
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const int BITS_PER_ELT = bits;
+  const int K = 16 / BITS_PER_ELT;
+  const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;
+  const int ELTS_PER_REG = 32 / BITS_PER_ELT;
 
-    const int BITS_PER_ELT  = bits;
-    const int K             = 16 / BITS_PER_ELT;
-    const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;
-    const int ELTS_PER_REG  = 32 / BITS_PER_ELT;
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
+  int MMA_SHAPE_N = 8;
+  int B_ROWS_PER_MMA = 8 * K;
+  const int elts_in_int32 = 32 / BITS_PER_ELT;
 
-    int       MMA_SHAPE_N    = 8;
-    int       B_ROWS_PER_MMA = 8 * K;
-    const int elts_in_int32  = 32 / BITS_PER_ELT;
+  const int num_vec_cols = num_cols / elts_in_int32;
 
-    const int num_vec_cols = num_cols / elts_in_int32;
-
-    // The code is written as below so it works for both int8 and packed int4.
-    for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
-        for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
-
-            for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
-                const int write_row = base_row + tile_row;
-                const int tile_read_row =
-                    8 * (((tile_row % ELTS_PER_REG) / 2)) + tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
-                if(base_row == 0 && write_col == 0){
-                    std::cout<<"tile_read_row:"<<tile_read_row<<std::endl;
-                }
-                const int read_row = base_row + tile_read_row;
-                const int read_col = write_col;
-
-                const int64_t read_offset  = int64_t(read_row) * num_vec_cols + read_col;
-                const int64_t write_offset = int64_t(write_row) * num_vec_cols + write_col;
-
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
+  // The code is written as below so it works for both int8 and packed int4.
+  for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
+    for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+      for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
+        const int write_row = base_row + tile_row;
+        const int tile_read_row = 8 * (((tile_row % ELTS_PER_REG) / 2)) +
+                                  tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
+        if (base_row == 0 && write_col == 0) {
+          std::cout << "tile_read_row:" << tile_read_row << std::endl;
         }
+        const int read_row = base_row + tile_read_row;
+        const int read_col = write_col;
+
+        const int64_t read_offset = int64_t(read_row) * num_vec_cols + read_col;
+        const int64_t write_offset =
+            int64_t(write_row) * num_vec_cols + write_col;
+
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
 
-template<int bits=4>
-void permute_B_rows_for_mixed_gemm_int4(int8_t*                    permuted_quantized_tensor,
-                                   const int8_t*              quantized_tensor,
-                                   const std::vector<size_t>& shape,
-                                   const int64_t              arch_version)
-{
+template <int bits = 4>
+void permute_B_rows_for_mixed_gemm_int4(int8_t* permuted_quantized_tensor,
+                                        const int8_t* quantized_tensor,
+                                        const std::vector<size_t>& shape,
+                                        const int64_t arch_version) {
+  // We only want to run this step for weight only quant.
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    // We only want to run this step for weight only quant.
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const int BITS_PER_ELT = bits;               // 4
+  const int K = 16 / BITS_PER_ELT;             // 4
+  const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;  // 2
+  const int ELTS_PER_REG = 32 / BITS_PER_ELT;  // 8
 
-    const int BITS_PER_ELT  = bits; //4
-    const int K             = 16 / BITS_PER_ELT; // 4
-    const int ELTS_PER_BYTE = 8 / BITS_PER_ELT; // 2
-    const int ELTS_PER_REG  = 32 / BITS_PER_ELT; // 8
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
+  int MMA_SHAPE_N = 8;
+  int B_ROWS_PER_MMA = 8 * K;  // 32
+  const int elts_in_int32 = 32 / BITS_PER_ELT;
 
-    int       MMA_SHAPE_N    = 8;
-    int       B_ROWS_PER_MMA = 8 * K; // 32
-    const int elts_in_int32  = 32 / BITS_PER_ELT;
+  const int num_vec_cols = num_cols / elts_in_int32;
+  const std::vector<int> tile_col_map{0,  2,  16, 18, 1,  3,  17, 19, 4,  6, 20,
+                                      22, 5,  7,  21, 23, 8,  10, 24, 26, 9, 11,
+                                      25, 27, 12, 14, 28, 30, 13, 15, 29, 31};
 
-    const int num_vec_cols = num_cols / elts_in_int32;
-    const std::vector<int> tile_col_map{
-                                       0,2,16,18,
-                                       1,3,17,19,
-                                       4,6,20,22,
-                                       5,7,21,23,
-                                       8,10,24,26,
-                                       9,11,25,27,
-                                       12,14,28,30,
-                                       13,15,29,31};
+  // const std::vector<int> tile_col_map{
+  //                   0                   0,2,16,18,
+  //                   4                   1,3,17,19,
+  //                   8                   4,6,20,22,
+  //                   12                  5,7,21,23,
+  //                   16                  8,10,24,26,
+  //                   20                  9,11,25,27,
+  //                   24                  12,14,28,30,
+  //                   28                  13,15,29,31};
+  // std::vector<int> tile_col_map(32);
+  // for(int i=0;i<32;i++){
+  //     tile_col_map[i]=i;
+  // }
+  // // tile_col_map[1]=4;
+  // tile_col_map[0]=0;
+  // tile_col_map[4]=1;
+  // tile_col_map[1]=2;
+  // tile_col_map[5]=3;
+  // tile_col_map[8]=4;
+  // tile_col_map[12]=5;
+  // tile_col_map[9]=6;
+  // tile_col_map[13]=7;
+  // tile_col_map[16]=8;
+  // tile_col_map[20]=9;
+  // tile_col_map[17]=10;
+  // tile_col_map[21]=11;
+  // tile_col_map[24]=12;
+  // tile_col_map[28]=13;
+  // tile_col_map[25]=14;
+  // tile_col_map[29]=15;
 
-    // const std::vector<int> tile_col_map{
-    //                   0                   0,2,16,18,
-    //                   4                   1,3,17,19,
-    //                   8                   4,6,20,22,
-    //                   12                  5,7,21,23,
-    //                   16                  8,10,24,26,
-    //                   20                  9,11,25,27,
-    //                   24                  12,14,28,30,
-    //                   28                  13,15,29,31};
-    // std::vector<int> tile_col_map(32);
-    // for(int i=0;i<32;i++){
-    //     tile_col_map[i]=i;
-    // }
-    // // tile_col_map[1]=4;
-    // tile_col_map[0]=0;
-    // tile_col_map[4]=1;
-    // tile_col_map[1]=2;
-    // tile_col_map[5]=3;
-    // tile_col_map[8]=4;
-    // tile_col_map[12]=5;
-    // tile_col_map[9]=6;
-    // tile_col_map[13]=7;
-    // tile_col_map[16]=8;
-    // tile_col_map[20]=9;
-    // tile_col_map[17]=10;
-    // tile_col_map[21]=11;
-    // tile_col_map[24]=12;
-    // tile_col_map[28]=13;
-    // tile_col_map[25]=14;
-    // tile_col_map[29]=15;
+  // tile_col_map[4]=1;
+  // tile_col_map[4]=1;
+  // tile_col_map[4]=2;
 
-    // tile_col_map[4]=1;
-    // tile_col_map[4]=1;
-    // tile_col_map[4]=2;
-
-    // The code is written as below so it works for both int8 and packed int4.
-    for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
-        for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
-
-            for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
-                const int write_row = base_row + tile_row;
-                // const int tile_read_row =
-                //     8 * (((tile_row % ELTS_PER_REG) / 2)) + tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
-                // const int tile_read_row = std::distance(tile_col_map.begin(), std::find(tile_col_map.begin(),tile_col_map.end(), tile_row));
-                const int tile_read_row = tile_col_map[tile_row];
-                if(base_row == 0 && write_col == 0){
-                    std::cout<<" write_row:"<<tile_row<<" tile_read_row:"<<tile_read_row<<std::endl;
-                }
-                const int read_row = base_row + tile_read_row;
-                const int read_col = write_col;
-
-                const int64_t read_offset  = int64_t(read_row) * num_vec_cols + read_col;
-                const int64_t write_offset = int64_t(write_row) * num_vec_cols + write_col;
-
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
+  // The code is written as below so it works for both int8 and packed int4.
+  for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
+    for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+      for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
+        const int write_row = base_row + tile_row;
+        // const int tile_read_row =
+        //     8 * (((tile_row % ELTS_PER_REG) / 2)) + tile_row % 2 + 2 *
+        //     (tile_row / ELTS_PER_REG);
+        // const int tile_read_row = std::distance(tile_col_map.begin(),
+        // std::find(tile_col_map.begin(),tile_col_map.end(), tile_row));
+        const int tile_read_row = tile_col_map[tile_row];
+        if (base_row == 0 && write_col == 0) {
+          std::cout << " write_row:" << tile_row
+                    << " tile_read_row:" << tile_read_row << std::endl;
         }
+        const int read_row = base_row + tile_read_row;
+        const int read_col = write_col;
+
+        const int64_t read_offset = int64_t(read_row) * num_vec_cols + read_col;
+        const int64_t write_offset =
+            int64_t(write_row) * num_vec_cols + write_col;
+
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
 
+void interleave_column_major_tensor(int8_t* interleaved_quantized_tensor,
+                                    const int8_t* quantized_tensor,
+                                    const std::vector<size_t>& shape) {
+  // We only want to run this step for weight only quant.
+  std::cout << "### in interleave_column_major_tensor" << std::endl;
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-void interleave_column_major_tensor(int8_t*                    interleaved_quantized_tensor,
-                                    const int8_t*              quantized_tensor,
-                                    const std::vector<size_t>& shape)
-{
+  const size_t BITS_PER_ELT = 8;
+  const size_t elts_in_int32 = 32 / BITS_PER_ELT;
 
-    // We only want to run this step for weight only quant.
-    std::cout<<"### in interleave_column_major_tensor"<<std::endl;
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const size_t rows_per_tile = 64;
+  std::cout << "running interleave_column_major_tensor" << std::endl;
+  std::cout << "num_rows:" << num_rows << ","
+            << "num_cols:" << num_cols << ","
+            << "BITS_PER_ELT:" << BITS_PER_ELT << ","
+            << "elts_in_int32:" << elts_in_int32 << ","
+            << "rows_per_tile:" << rows_per_tile << std::endl;
 
-    const size_t BITS_PER_ELT  = 8;
-    const size_t elts_in_int32 = 32 / BITS_PER_ELT;
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
 
-    const size_t rows_per_tile = 64;
-    std::cout<<"running interleave_column_major_tensor"<<std::endl;
-    std::cout<<"num_rows:"<<num_rows<<","
-             <<"num_cols:"<<num_cols<<","
-             <<"BITS_PER_ELT:"<<BITS_PER_ELT<<","
-             <<"elts_in_int32:"<<elts_in_int32<<","
-             <<"rows_per_tile:"<<rows_per_tile<<std::endl;
+  const size_t num_vec_rows = num_rows / elts_in_int32;
+  const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
+  const size_t interleave = 2;
+  std::cout << "num_vec_rows:" << num_vec_rows << ","
+            << "vec_rows_per_tile:" << vec_rows_per_tile << ","
+            << "interleave:" << interleave << std::endl;
+  for (int read_col = 0; read_col < num_cols; ++read_col) {
+    const size_t write_col = read_col / interleave;
+    for (int base_vec_row = 0; base_vec_row < num_vec_rows;
+         base_vec_row += vec_rows_per_tile) {
+      for (int vec_read_row = base_vec_row;
+           vec_read_row <
+           std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
+           ++vec_read_row) {
+        const size_t vec_write_row =
+            interleave * base_vec_row +
+            vec_rows_per_tile * (read_col % interleave) +
+            vec_read_row % vec_rows_per_tile;
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
-
-
-    const size_t num_vec_rows      = num_rows / elts_in_int32;
-    const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
-    const size_t interleave        = 2;
-    std::cout<<"num_vec_rows:"<<num_vec_rows<<","
-             <<"vec_rows_per_tile:"<<vec_rows_per_tile<<","
-             <<"interleave:"<<interleave<<std::endl;
-    for (int read_col = 0; read_col < num_cols; ++read_col) {
-        const size_t write_col = read_col / interleave;
-        for (int base_vec_row = 0; base_vec_row < num_vec_rows; base_vec_row += vec_rows_per_tile) {
-            for (int vec_read_row = base_vec_row;
-                    vec_read_row < std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
-                    ++vec_read_row) {
-                const size_t vec_write_row = interleave * base_vec_row
-                                                + vec_rows_per_tile * (read_col % interleave)
-                                                + vec_read_row % vec_rows_per_tile;
-
-                const size_t read_offset = size_t(read_col) * num_vec_rows + vec_read_row;
-                const size_t write_offset = size_t(write_col) * num_vec_rows * interleave + vec_write_row;
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
-        }
+        const size_t read_offset =
+            size_t(read_col) * num_vec_rows + vec_read_row;
+        const size_t write_offset =
+            size_t(write_col) * num_vec_rows * interleave + vec_write_row;
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
 
+void interleave_column_major_tensor_int4(int8_t* interleaved_quantized_tensor,
+                                         const int8_t* quantized_tensor,
+                                         const std::vector<size_t>& shape) {
+  // We only want to run this step for weight only quant.
+  std::cout << "### in interleave_column_major_tensor" << std::endl;
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-void interleave_column_major_tensor_int4(int8_t*                    interleaved_quantized_tensor,
-                                    const int8_t*              quantized_tensor,
-                                    const std::vector<size_t>& shape)
-{
+  const size_t BITS_PER_ELT = 4;
+  const size_t elts_in_int32 = 32 / BITS_PER_ELT;
 
-    // We only want to run this step for weight only quant.
-    std::cout<<"### in interleave_column_major_tensor"<<std::endl;
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const size_t rows_per_tile = 64;
+  std::cout << "running interleave_column_major_tensor" << std::endl;
+  std::cout << "num_rows:" << num_rows << ","
+            << "num_cols:" << num_cols << ","
+            << "BITS_PER_ELT:" << BITS_PER_ELT << ","
+            << "elts_in_int32:" << elts_in_int32 << ","
+            << "rows_per_tile:" << rows_per_tile << std::endl;
 
-    const size_t BITS_PER_ELT  = 4;
-    const size_t elts_in_int32 = 32 / BITS_PER_ELT;
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
 
-    const size_t rows_per_tile = 64;
-    std::cout<<"running interleave_column_major_tensor"<<std::endl;
-    std::cout<<"num_rows:"<<num_rows<<","
-             <<"num_cols:"<<num_cols<<","
-             <<"BITS_PER_ELT:"<<BITS_PER_ELT<<","
-             <<"elts_in_int32:"<<elts_in_int32<<","
-             <<"rows_per_tile:"<<rows_per_tile<<std::endl;
+  const size_t num_vec_rows = num_rows / elts_in_int32;
+  const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
+  const size_t interleave = 4;
+  std::cout << "num_vec_rows:" << num_vec_rows << ","
+            << "vec_rows_per_tile:" << vec_rows_per_tile << ","
+            << "interleave:" << interleave << std::endl;
+  for (int read_col = 0; read_col < num_cols; ++read_col) {
+    const size_t write_col = read_col / interleave;
+    for (int base_vec_row = 0; base_vec_row < num_vec_rows;
+         base_vec_row += vec_rows_per_tile) {
+      for (int vec_read_row = base_vec_row;
+           vec_read_row <
+           std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
+           ++vec_read_row) {
+        const size_t vec_write_row =
+            interleave * base_vec_row +
+            vec_rows_per_tile * (read_col % interleave) +
+            vec_read_row % vec_rows_per_tile;
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
-
-
-    const size_t num_vec_rows      = num_rows / elts_in_int32;
-    const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
-    const size_t interleave        = 4;
-    std::cout<<"num_vec_rows:"<<num_vec_rows<<","
-             <<"vec_rows_per_tile:"<<vec_rows_per_tile<<","
-             <<"interleave:"<<interleave<<std::endl;
-    for (int read_col = 0; read_col < num_cols; ++read_col) {
-        const size_t write_col = read_col / interleave;
-        for (int base_vec_row = 0; base_vec_row < num_vec_rows; base_vec_row += vec_rows_per_tile) {
-            for (int vec_read_row = base_vec_row;
-                    vec_read_row < std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
-                    ++vec_read_row) {
-                const size_t vec_write_row = interleave * base_vec_row
-                                                + vec_rows_per_tile * (read_col % interleave)
-                                                + vec_read_row % vec_rows_per_tile;
-
-                const size_t read_offset = size_t(read_col) * num_vec_rows + vec_read_row;
-                const size_t write_offset = size_t(write_col) * num_vec_rows * interleave + vec_write_row;
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
-        }
+        const size_t read_offset =
+            size_t(read_col) * num_vec_rows + vec_read_row;
+        const size_t write_offset =
+            size_t(write_col) * num_vec_rows * interleave + vec_write_row;
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/cutlass_gemm_caller.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/cutlass_gemm_caller.cuh
index 1944716c3e..e35018a26f 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/cutlass_gemm_caller.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/cutlass_gemm_caller.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
 #pragma once
 
 // clang-format will break include orders
@@ -21,15 +22,16 @@
 
 namespace fastdeploy::c3x {
 
-static inline cute::Shape<int, int, int, int>
-get_problem_shape(paddle::Tensor const &a, paddle::Tensor const &b) {
+static inline cute::Shape<int, int, int, int> get_problem_shape(
+    paddle::Tensor const &a, paddle::Tensor const &b) {
   int32_t m = a.dims()[0], n = b.dims()[0], k = a.dims()[1];
   return {m, n, k, 1};
 }
 
 template <typename GemmKernel>
 void cutlass_gemm_caller(
-    phi::Place device, cute::Shape<int, int, int, int> prob_shape,
+    phi::Place device,
+    cute::Shape<int, int, int, int> prob_shape,
     typename GemmKernel::MainloopArguments mainloop_args,
     typename GemmKernel::EpilogueArguments epilogue_args,
     typename GemmKernel::TileSchedulerArguments scheduler = {}) {
@@ -57,7 +59,8 @@ void cutlass_gemm_caller(
 }
 
 template <typename Gemm, typename... EpilogueArgs>
-void cutlass_gemm_caller(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_gemm_caller(paddle::Tensor &out,
+                         paddle::Tensor const &a,
                          paddle::Tensor const &b,
                          EpilogueArgs &&...epilogue_params) {
   using ElementAB = typename Gemm::ElementAB;
@@ -86,17 +89,20 @@ void cutlass_gemm_caller(paddle::Tensor &out, paddle::Tensor const &a,
 
   auto a_ptr = static_cast<ElementAB *>(const_cast<void *>(a.data()));
   auto b_ptr = static_cast<ElementAB *>(const_cast<void *>(b.data()));
-  typename GemmKernel::MainloopArguments mainloop_args{a_ptr, a_stride, b_ptr,
-                                                       b_stride};
+  typename GemmKernel::MainloopArguments mainloop_args{
+      a_ptr, a_stride, b_ptr, b_stride};
 
   auto c_ptr = static_cast<ElementD *>(const_cast<void *>(out.data()));
   typename GemmKernel::EpilogueArguments epilogue_args{
       Gemm::Epilogue::prepare_args(
           std::forward<EpilogueArgs>(epilogue_params)...),
-      c_ptr, c_stride, c_ptr, d_stride};
+      c_ptr,
+      c_stride,
+      c_ptr,
+      d_stride};
 
-  cutlass_gemm_caller<GemmKernel>(a.place(), prob_shape, mainloop_args,
-                                  epilogue_args);
+  cutlass_gemm_caller<GemmKernel>(
+      a.place(), prob_shape, mainloop_args, epilogue_args);
 }
 
-} // namespace fastdeploy::c3x
+}  // namespace fastdeploy::c3x
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm.cuh
index 26278a79fd..4d8edbd621 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm.cuh
 
 #pragma once
 
@@ -31,16 +32,19 @@ using namespace cute;
 
 namespace fastdeploy {
 
-template <typename ElementAB_, typename ElementD_,
-          template <typename, typename, typename> typename Epilogue_,
-          typename TileShape, typename ClusterShape, typename KernelSchedule,
+template <typename ElementAB_,
+          typename ElementD_,
+          template <typename, typename, typename>
+          typename Epilogue_,
+          typename TileShape,
+          typename ClusterShape,
+          typename KernelSchedule,
           typename EpilogueSchedule>
 struct cutlass_3x_gemm {
   using ElementAB = ElementAB_;
   using ElementD = ElementD_;
-  using ElementAcc =
-      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
-                                float>::type;
+  using ElementAcc = typename std::
+      conditional<std::is_same_v<ElementAB, int8_t>, int32_t, float>::type;
 
   using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
 
@@ -57,10 +61,21 @@ struct cutlass_3x_gemm {
 
   using CollectiveEpilogue =
       typename cutlass::epilogue::collective::CollectiveBuilder<
-          cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape,
-          ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
-          ElementAcc, float, ElementC, StrideC, AlignmentCD, ElementD, StrideD,
-          AlignmentCD, EpilogueSchedule, EVTCompute>::CollectiveOp;
+          cutlass::arch::Sm90,
+          cutlass::arch::OpClassTensorOp,
+          TileShape,
+          ClusterShape,
+          cutlass::epilogue::collective::EpilogueTileAuto,
+          ElementAcc,
+          float,
+          ElementC,
+          StrideC,
+          AlignmentCD,
+          ElementD,
+          StrideD,
+          AlignmentCD,
+          EpilogueSchedule,
+          EVTCompute>::CollectiveOp;
 
   static constexpr size_t CEStorageSize =
       sizeof(typename CollectiveEpilogue::SharedStorage);
@@ -78,16 +93,22 @@ struct cutlass_3x_gemm {
           KernelSchedule>::CollectiveOp;
   // clang-format on
 
-  using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
-      cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
-      cutlass::gemm::PersistentScheduler>>;
+  using KernelType = enable_sm90_or_later<
+      cutlass::gemm::kernel::GemmUniversal<cute::Shape<int, int, int, int>,
+                                           CollectiveMainloop,
+                                           CollectiveEpilogue,
+                                           cutlass::gemm::PersistentScheduler>>;
 
   struct GemmKernel : public KernelType {};
 };
 
-template <typename ElementAB_, typename ElementD_,
-          template <typename, typename, typename> typename Epilogue_,
-          typename TileShape, typename ClusterShape, typename KernelSchedule,
+template <typename ElementAB_,
+          typename ElementD_,
+          template <typename, typename, typename>
+          typename Epilogue_,
+          typename TileShape,
+          typename ClusterShape,
+          typename KernelSchedule,
           typename EpilogueSchedule>
 struct cutlass_3x_gemm_sm100 {
   using ElementAB = ElementAB_;
@@ -108,9 +129,8 @@ struct cutlass_3x_gemm_sm100 {
   using LayoutD = cutlass::layout::RowMajor;
   static constexpr int AlignmentD = AlignmentC;
 
-  using ElementAcc =
-      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
-                                float>::type;
+  using ElementAcc = typename std::
+      conditional<std::is_same_v<ElementAB, int8_t>, int32_t, float>::type;
   using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
 
   // MMA type
@@ -127,23 +147,44 @@ struct cutlass_3x_gemm_sm100 {
 
   using CollectiveEpilogue =
       typename cutlass::epilogue::collective::CollectiveBuilder<
-          cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, TileShape,
-          ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
-          ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC,
-          ElementD, LayoutD, AlignmentD, EpilogueSchedule,
+          cutlass::arch::Sm100,
+          cutlass::arch::OpClassTensorOp,
+          TileShape,
+          ClusterShape,
+          cutlass::epilogue::collective::EpilogueTileAuto,
+          ElementAccumulator,
+          ElementCompute,
+          ElementC,
+          LayoutC,
+          AlignmentC,
+          ElementD,
+          LayoutD,
+          AlignmentD,
+          EpilogueSchedule,
           EVTCompute>::CollectiveOp;
 
   using CollectiveMainloop =
       typename cutlass::gemm::collective::CollectiveBuilder<
-          cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, ElementAB,
-          LayoutA, AlignmentA, ElementAB, LayoutB, AlignmentB,
-          ElementAccumulator, TileShape, ClusterShape,
+          cutlass::arch::Sm100,
+          cutlass::arch::OpClassTensorOp,
+          ElementAB,
+          LayoutA,
+          AlignmentA,
+          ElementAB,
+          LayoutB,
+          AlignmentB,
+          ElementAccumulator,
+          TileShape,
+          ClusterShape,
           cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
               sizeof(typename CollectiveEpilogue::SharedStorage))>,
           KernelSchedule>::CollectiveOp;
 
-  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
-      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>,
+                                           CollectiveMainloop,
+                                           CollectiveEpilogue,
+                                           void>;
 };
 
-} // namespace fastdeploy
+}  // namespace fastdeploy
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_azp_sm90_int8.cu b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_azp_sm90_int8.cu
index f5d4d6aa28..704f1021f8 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_azp_sm90_int8.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_azp_sm90_int8.cu
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_azp_sm90_int8.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_azp_sm90_int8.cu
 
 // clang-format will break include orders
 // clang-format off
@@ -10,18 +11,22 @@
 namespace fastdeploy {
 
 void cutlass_scaled_mm_azp_sm90_int8(
-    paddle::Tensor &out, paddle::Tensor const &a, paddle::Tensor const &b,
-    paddle::Tensor const &a_scales, paddle::Tensor const &b_scales,
-    paddle::Tensor const &azp_adj, paddle::optional<paddle::Tensor> const &azp,
+    paddle::Tensor &out,
+    paddle::Tensor const &a,
+    paddle::Tensor const &b,
+    paddle::Tensor const &a_scales,
+    paddle::Tensor const &b_scales,
+    paddle::Tensor const &azp_adj,
+    paddle::optional<paddle::Tensor> const &azp,
     paddle::optional<paddle::Tensor> const &bias) {
   if (azp) {
     return cutlass_scaled_mm_sm90_int8_epilogue<
-        c3x::ScaledEpilogueBiasAzpToken>(out, a, b, a_scales, b_scales, azp_adj,
-                                         *azp, bias);
+        c3x::ScaledEpilogueBiasAzpToken>(
+        out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
   } else {
     return cutlass_scaled_mm_sm90_int8_epilogue<c3x::ScaledEpilogueBiasAzp>(
         out, a, b, a_scales, b_scales, azp_adj, bias);
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_helper.hpp b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_helper.hpp
index 9a601f75ad..2bfa58231f 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_helper.hpp
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_helper.hpp
@@ -1,34 +1,38 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_helper.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_helper.hpp
 
 #include "helper.h"
 
 template <typename Fp8Func, typename Int8Func>
-void dispatch_scaled_mm(paddle::Tensor &c, paddle::Tensor const &a,
-                        paddle::Tensor const &b, paddle::Tensor const &a_scales,
+void dispatch_scaled_mm(paddle::Tensor &c,
+                        paddle::Tensor const &a,
+                        paddle::Tensor const &b,
+                        paddle::Tensor const &a_scales,
                         paddle::Tensor const &b_scales,
                         paddle::optional<paddle::Tensor> const &bias,
-                        Fp8Func fp8_func, Int8Func int8_func) {
-    PD_CHECK(a_scales.dtype() == paddle::DataType::FLOAT32);
-    PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
+                        Fp8Func fp8_func,
+                        Int8Func int8_func) {
+  PD_CHECK(a_scales.dtype() == paddle::DataType::FLOAT32);
+  PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
 
-    int M = a.dims()[0], N = b.dims()[0], K = a.dims()[1];
+  int M = a.dims()[0], N = b.dims()[0], K = a.dims()[1];
 
-    if ((a_scales.numel() == 1 || a_scales.numel() == a.dims()[0]) &&
-        (b_scales.numel() == 1 || b_scales.numel() == b.dims()[0])) {
-      // Standard per-tensor/per-token/per-channel scaling
-      PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
-      if (a.dtype() == phi::DataType::FLOAT8_E4M3FN) {
-        fp8_func(c, a, b, a_scales, b_scales, bias);
-      } else {
-        PD_CHECK(a.dtype() == paddle::DataType::INT8);
-        if constexpr (!std::is_same_v<Int8Func, std::nullptr_t>) {
-          int8_func(c, a, b, a_scales, b_scales, bias);
-        } else {
-          PD_CHECK(false, "Int8 not supported for this architecture");
-        }
-      }
+  if ((a_scales.numel() == 1 || a_scales.numel() == a.dims()[0]) &&
+      (b_scales.numel() == 1 || b_scales.numel() == b.dims()[0])) {
+    // Standard per-tensor/per-token/per-channel scaling
+    PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
+    if (a.dtype() == phi::DataType::FLOAT8_E4M3FN) {
+      fp8_func(c, a, b, a_scales, b_scales, bias);
     } else {
-      PADDLE_THROW(phi::errors::Unimplemented(
-          "No kernel for this combination of input dtypes is implemented."));
+      PD_CHECK(a.dtype() == paddle::DataType::INT8);
+      if constexpr (!std::is_same_v<Int8Func, std::nullptr_t>) {
+        int8_func(c, a, b, a_scales, b_scales, bias);
+      } else {
+        PD_CHECK(false, "Int8 not supported for this architecture");
+      }
     }
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "No kernel for this combination of input dtypes is implemented."));
+  }
 }
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_kernels.hpp b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_kernels.hpp
index 75472ea805..4e3e67ac5a 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_kernels.hpp
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_kernels.hpp
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_kernels.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_kernels.hpp
 
 #pragma once
 
@@ -6,30 +7,35 @@
 
 namespace fastdeploy {
 
-void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out,
+                                paddle::Tensor const &a,
                                 paddle::Tensor const &b,
                                 paddle::Tensor const &a_scales,
                                 paddle::Tensor const &b_scales,
                                 paddle::optional<paddle::Tensor> const &bias);
 
-void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out,
+                                 paddle::Tensor const &a,
                                  paddle::Tensor const &b,
                                  paddle::Tensor const &a_scales,
                                  paddle::Tensor const &b_scales,
                                  paddle::optional<paddle::Tensor> const &bias);
 
-void cutlass_scaled_mm_azp_sm90_int8(paddle::Tensor& out, paddle::Tensor const& a,
-                                     paddle::Tensor const& b,
-                                     paddle::Tensor const& a_scales,
-                                     paddle::Tensor const& b_scales,
-                                     paddle::Tensor const& azp_adj,
-                                     paddle::optional<paddle::Tensor> const& azp,
-                                     paddle::optional<paddle::Tensor> const& bias);
+void cutlass_scaled_mm_azp_sm90_int8(
+    paddle::Tensor &out,
+    paddle::Tensor const &a,
+    paddle::Tensor const &b,
+    paddle::Tensor const &a_scales,
+    paddle::Tensor const &b_scales,
+    paddle::Tensor const &azp_adj,
+    paddle::optional<paddle::Tensor> const &azp,
+    paddle::optional<paddle::Tensor> const &bias);
 
-void cutlass_scaled_mm_sm100_fp8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm100_fp8(paddle::Tensor &out,
+                                 paddle::Tensor const &a,
                                  paddle::Tensor const &b,
                                  paddle::Tensor const &a_scales,
                                  paddle::Tensor const &b_scales,
                                  paddle::optional<paddle::Tensor> const &bias);
 
-} // namespace fastdeploy
+}  // namespace fastdeploy
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8.cu b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8.cu
index 801e90fd73..1b197b8028 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8.cu
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8.cu
 
 // clang-format will break include orders
 // clang-format off
@@ -9,7 +10,8 @@
 
 namespace fastdeploy {
 
-void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out,
+                                paddle::Tensor const &a,
                                 paddle::Tensor const &b,
                                 paddle::Tensor const &a_scales,
                                 paddle::Tensor const &b_scales,
@@ -17,7 +19,8 @@ void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out, paddle::Tensor const &a,
   PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-             "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm90_fp8_epilogue<c3x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -25,4 +28,4 @@ void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out, paddle::Tensor const &a,
         out, a, b, a_scales, b_scales);
   }
 }
-} // namespace fastdeploy
+}  // namespace fastdeploy
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh
index ac86aeba85..cd0eda3ad9 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh
 
 #pragma once
 
@@ -17,8 +18,10 @@ namespace fastdeploy {
 
 using c3x::cutlass_gemm_caller;
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_fp8_config_default {
   // M in (128, inf)
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
@@ -27,13 +30,19 @@ struct sm90_fp8_config_default {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_128, _128, _128>;
   using ClusterShape = Shape<_2, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_fp8_config_M128 {
   // M in (64, 128]
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
@@ -42,13 +51,19 @@ struct sm90_fp8_config_M128 {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _128, _128>;
   using ClusterShape = Shape<_2, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_fp8_config_M64 {
   // M in [1, 64]
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
@@ -58,13 +73,19 @@ struct sm90_fp8_config_M64 {
   using TileShape = Shape<_64, _64, _128>;
   using ClusterShape = Shape<_1, _8, _1>;
 
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm90_fp8_dispatch(paddle::Tensor &out,
                                            paddle::Tensor const &a,
@@ -75,8 +96,8 @@ inline void cutlass_gemm_sm90_fp8_dispatch(paddle::Tensor &out,
   PD_CHECK(b.dtype() == phi::DataType::FLOAT8_E4M3FN);
 
   using Cutlass3xGemmDefault =
-      typename sm90_fp8_config_default<InType, OutType,
-                                       Epilogue>::Cutlass3xGemm;
+      typename sm90_fp8_config_default<InType, OutType, Epilogue>::
+          Cutlass3xGemm;
   using Cutlass3xGemmM64 =
       typename sm90_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
   using Cutlass3xGemmM128 =
@@ -84,7 +105,7 @@ inline void cutlass_gemm_sm90_fp8_dispatch(paddle::Tensor &out,
 
   uint32_t const m = a.dims()[0];
   uint32_t const mp2 =
-      std::max(static_cast<uint32_t>(64), next_pow_2(m)); // next power of 2
+      std::max(static_cast<uint32_t>(64), next_pow_2(m));  // next power of 2
 
   if (mp2 <= 64) {
     // m in [1, 64]
@@ -112,14 +133,16 @@ void cutlass_scaled_mm_sm90_fp8_epilogue(paddle::Tensor &out,
 
   if (out.dtype() == paddle::DataType::BFLOAT16) {
     return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
-                                          cutlass::bfloat16_t, Epilogue>(
+                                          cutlass::bfloat16_t,
+                                          Epilogue>(
         out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
   } else {
     PD_CHECK(out.dtype() == paddle::DataType::FLOAT16);
     return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
-                                          cutlass::half_t, Epilogue>(
+                                          cutlass::half_t,
+                                          Epilogue>(
         out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8.cu b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8.cu
index 633f76fd88..5256b27c14 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8.cu
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8.cu
 
 // clang-format will break include orders
 // clang-format off
@@ -9,7 +10,8 @@
 
 namespace fastdeploy {
 
-void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out,
+                                 paddle::Tensor const &a,
                                  paddle::Tensor const &b,
                                  paddle::Tensor const &a_scales,
                                  paddle::Tensor const &b_scales,
@@ -17,7 +19,8 @@ void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out, paddle::Tensor const &a,
   PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-             "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm90_int8_epilogue<c3x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -26,4 +29,4 @@ void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out, paddle::Tensor const &a,
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh
index df63de0fa6..1b14e1b749 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh
 
 #pragma once
 
@@ -17,8 +18,10 @@ namespace fastdeploy {
 
 using c3x::cutlass_gemm_caller;
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_default {
   // For M > 128 and any N
   static_assert(std::is_same<InType, int8_t>());
@@ -27,13 +30,19 @@ struct sm90_int8_config_default {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_128, _128, _128>;
   using ClusterShape = Shape<_2, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_M128 {
   // For M in (64, 128] and any N
   static_assert(std::is_same<InType, int8_t>());
@@ -42,13 +51,19 @@ struct sm90_int8_config_M128 {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _128, _128>;
   using ClusterShape = Shape<_2, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_M64 {
   // For M in (32, 64] and any N
   static_assert(std::is_same<InType, int8_t>());
@@ -56,13 +71,19 @@ struct sm90_int8_config_M64 {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _64, _256>;
   using ClusterShape = Shape<_1, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_M32_NBig {
   // For M in [1, 32] and N >= 8192
   static_assert(std::is_same<InType, int8_t>());
@@ -70,13 +91,19 @@ struct sm90_int8_config_M32_NBig {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _128, _256>;
   using ClusterShape = Shape<_1, _4, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_M32_NSmall {
   // For M in [1, 32] and N < 8192
   static_assert(std::is_same<InType, int8_t>());
@@ -84,13 +111,19 @@ struct sm90_int8_config_M32_NSmall {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _64, _256>;
   using ClusterShape = Shape<_1, _8, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm90_int8_dispatch(paddle::Tensor &out,
                                             paddle::Tensor const &a,
@@ -101,25 +134,25 @@ inline void cutlass_gemm_sm90_int8_dispatch(paddle::Tensor &out,
   PD_CHECK(b.dtype() == paddle::DataType::INT8);
 
   using Cutlass3xGemmDefault =
-      typename sm90_int8_config_default<InType, OutType,
-                                        Epilogue>::Cutlass3xGemm;
+      typename sm90_int8_config_default<InType, OutType, Epilogue>::
+          Cutlass3xGemm;
   using Cutlass3xGemmM128 =
       typename sm90_int8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
   using Cutlass3xGemmM64 =
       typename sm90_int8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
   using Cutlass3xGemmM32NBig =
-      typename sm90_int8_config_M32_NBig<InType, OutType,
-                                         Epilogue>::Cutlass3xGemm;
+      typename sm90_int8_config_M32_NBig<InType, OutType, Epilogue>::
+          Cutlass3xGemm;
   using Cutlass3xGemmM32NSmall =
-      typename sm90_int8_config_M32_NSmall<InType, OutType,
-                                           Epilogue>::Cutlass3xGemm;
+      typename sm90_int8_config_M32_NSmall<InType, OutType, Epilogue>::
+          Cutlass3xGemm;
 
   uint32_t const n = out.dims()[1];
   bool const is_small_n = n < 8192;
 
   uint32_t const m = a.dims()[0];
   uint32_t const mp2 =
-      std::max(static_cast<uint32_t>(32), next_pow_2(m)); // next power of 2
+      std::max(static_cast<uint32_t>(32), next_pow_2(m));  // next power of 2
 
   if (mp2 <= 32) {
     // m in [1, 32]
@@ -155,7 +188,8 @@ void cutlass_scaled_mm_sm90_int8_epilogue(paddle::Tensor &out,
   PD_CHECK(b.dtype() == paddle::DataType::INT8);
 
   if (out.dtype() == paddle::DataType::BFLOAT16) {
-    return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
+    return cutlass_gemm_sm90_int8_dispatch<int8_t,
+                                           cutlass::bfloat16_t,
                                            Epilogue>(
         out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
   } else {
@@ -165,4 +199,4 @@ void cutlass_scaled_mm_sm90_int8_epilogue(paddle::Tensor &out,
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cu b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cu
index 55015ea3e9..e3ce7e1fcb 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cu
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
 
 #include "helper.h"
 #include <stddef.h>
@@ -20,7 +21,8 @@ using namespace fastdeploy;
 
 template <template <typename, typename> typename Epilogue,
           typename... EpilogueArgs>
-void cutlass_scaled_mm_sm75_epilogue(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm75_epilogue(paddle::Tensor& out,
+                                     paddle::Tensor const& a,
                                      paddle::Tensor const& b,
                                      EpilogueArgs&&... epilogue_args) {
   PD_CHECK(a.dtype() == paddle::DataType::INT8);
@@ -36,7 +38,8 @@ void cutlass_scaled_mm_sm75_epilogue(paddle::Tensor& out, paddle::Tensor const&
   }
 }
 
-void cutlass_scaled_mm_sm75(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm75(paddle::Tensor& out,
+                            paddle::Tensor const& a,
                             paddle::Tensor const& b,
                             paddle::Tensor const& a_scales,
                             paddle::Tensor const& b_scales,
@@ -45,7 +48,8 @@ void cutlass_scaled_mm_sm75(paddle::Tensor& out, paddle::Tensor const& a,
   PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-                "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -54,7 +58,8 @@ void cutlass_scaled_mm_sm75(paddle::Tensor& out, paddle::Tensor const& a,
   }
 }
 
-void cutlass_scaled_mm_azp_sm75(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm75(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -75,7 +80,8 @@ void cutlass_scaled_mm_azp_sm75(paddle::Tensor& out, paddle::Tensor const& a,
 
 template <template <typename, typename> typename Epilogue,
           typename... EpilogueArgs>
-void cutlass_scaled_mm_sm80_epilogue(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm80_epilogue(paddle::Tensor& out,
+                                     paddle::Tensor const& a,
                                      paddle::Tensor const& b,
                                      EpilogueArgs&&... epilogue_args) {
   PD_CHECK(a.dtype() == paddle::DataType::INT8);
@@ -91,7 +97,8 @@ void cutlass_scaled_mm_sm80_epilogue(paddle::Tensor& out, paddle::Tensor const&
   }
 }
 
-void cutlass_scaled_mm_sm80(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm80(paddle::Tensor& out,
+                            paddle::Tensor const& a,
                             paddle::Tensor const& b,
                             paddle::Tensor const& a_scales,
                             paddle::Tensor const& b_scales,
@@ -100,7 +107,8 @@ void cutlass_scaled_mm_sm80(paddle::Tensor& out, paddle::Tensor const& a,
   PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-                "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -109,7 +117,8 @@ void cutlass_scaled_mm_sm80(paddle::Tensor& out, paddle::Tensor const& a,
   }
 }
 
-void cutlass_scaled_mm_azp_sm80(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm80(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -130,14 +139,16 @@ void cutlass_scaled_mm_azp_sm80(paddle::Tensor& out, paddle::Tensor const& a,
 
 template <template <typename, typename> typename Epilogue,
           typename... EpilogueArgs>
-void cutlass_scaled_mm_sm89_epilogue(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm89_epilogue(paddle::Tensor& out,
+                                     paddle::Tensor const& a,
                                      paddle::Tensor const& b,
                                      EpilogueArgs&&... epilogue_args) {
   if (a.dtype() == paddle::DataType::INT8) {
     PD_CHECK(b.dtype() == paddle::DataType::INT8);
 
     if (out.dtype() == paddle::DataType::BFLOAT16) {
-      return cutlass_gemm_sm89_int8_dispatch<int8_t, cutlass::bfloat16_t,
+      return cutlass_gemm_sm89_int8_dispatch<int8_t,
+                                             cutlass::bfloat16_t,
                                              Epilogue>(
           out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
     } else {
@@ -151,18 +162,21 @@ void cutlass_scaled_mm_sm89_epilogue(paddle::Tensor& out, paddle::Tensor const&
 
     if (out.dtype() == paddle::DataType::BFLOAT16) {
       return cutlass_gemm_sm89_fp8_dispatch<cutlass::float_e4m3_t,
-                                            cutlass::bfloat16_t, Epilogue>(
+                                            cutlass::bfloat16_t,
+                                            Epilogue>(
           out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
     } else {
       PD_CHECK(out.dtype() == paddle::DataType::FLOAT16);
       return cutlass_gemm_sm89_fp8_dispatch<cutlass::float_e4m3_t,
-                                            cutlass::half_t, Epilogue>(
+                                            cutlass::half_t,
+                                            Epilogue>(
           out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
     }
   }
 }
 
-void cutlass_scaled_mm_sm89(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm89(paddle::Tensor& out,
+                            paddle::Tensor const& a,
                             paddle::Tensor const& b,
                             paddle::Tensor const& a_scales,
                             paddle::Tensor const& b_scales,
@@ -171,7 +185,8 @@ void cutlass_scaled_mm_sm89(paddle::Tensor& out, paddle::Tensor const& a,
   PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-                "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -180,7 +195,8 @@ void cutlass_scaled_mm_sm89(paddle::Tensor& out, paddle::Tensor const& a,
   }
 }
 
-void cutlass_scaled_mm_azp_sm89(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm89(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cuh
index b66041cc59..2527e037b0 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
 
 #pragma once
 #include <stddef.h>
@@ -70,34 +71,42 @@ struct enable_sm89_to_sm90 : Kernel {
 #endif
   }
 };
-template <typename Arch, template <typename> typename ArchGuard,
-          typename ElementAB_, typename ElementD_,
-          template <typename, typename> typename Epilogue_, typename TileShape,
-          typename WarpShape, typename InstructionShape, int32_t MainLoopStages,
+template <typename Arch,
+          template <typename>
+          typename ArchGuard,
+          typename ElementAB_,
+          typename ElementD_,
+          template <typename, typename>
+          typename Epilogue_,
+          typename TileShape,
+          typename WarpShape,
+          typename InstructionShape,
+          int32_t MainLoopStages,
           typename FP8MathOperator = cutlass::arch::OpMultiplyAdd>
 struct cutlass_2x_gemm {
   using ElementAB = ElementAB_;
   using ElementD = ElementD_;
 
-  using ElementAcc =
-      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
-                                float>::type;
+  using ElementAcc = typename std::
+      conditional<std::is_same_v<ElementAB, int8_t>, int32_t, float>::type;
 
   using Operator =
       typename std::conditional<std::is_same_v<ElementAB, int8_t>,
                                 cutlass::arch::OpMultiplyAddSaturate,
                                 FP8MathOperator>::type;
 
-  using OutputTileThreadMap =
-      cutlass::epilogue::threadblock::OutputTileThreadLayout<
-          TileShape, WarpShape, float, 4, 1 /* epilogue stages */
-          >;
+  using OutputTileThreadMap = cutlass::epilogue::threadblock::
+      OutputTileThreadLayout<TileShape, WarpShape, float, 4, 1 /* epilogue
+                                                                  stages */
+                             >;
 
   using Epilogue = Epilogue_<ElementD, OutputTileThreadMap>;
   using EVTCompute = typename Epilogue::EVTCompute;
 
   using D = cutlass::epilogue::threadblock::VisitorAuxStore<
-      OutputTileThreadMap, ElementD, cutlass::FloatRoundStyle::round_to_nearest,
+      OutputTileThreadMap,
+      ElementD,
+      cutlass::FloatRoundStyle::round_to_nearest,
       Stride<int64_t, Int<1>, Int<0>>>;
 
   using EVTD = cutlass::epilogue::threadblock::Sm80EVT<D, EVTCompute>;
@@ -129,7 +138,8 @@ struct cutlass_2x_gemm {
 };
 
 template <typename Gemm, typename... EpilogueArgs>
-inline void cutlass_gemm_caller(paddle::Tensor& out, paddle::Tensor const& a,
+inline void cutlass_gemm_caller(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 EpilogueArgs&&... epilogue_params) {
   using ElementAB = typename Gemm::ElementAB;
@@ -183,7 +193,7 @@ inline void cutlass_gemm_caller(paddle::Tensor& out, paddle::Tensor const& a,
   // Launch the CUTLASS GEMM kernel.
   typename Gemm::Op gemm_op;
   size_t workspace_size = gemm_op.get_workspace_size(args);
-  phi::Allocator *allocator = paddle::GetAllocator(a.place());
+  phi::Allocator* allocator = paddle::GetAllocator(a.place());
   auto workspace = allocator->Allocate(workspace_size);
 
   auto stream = a.stream();
@@ -210,11 +220,10 @@ inline void fallback_cutlass_gemm_caller(paddle::Tensor& out,
       sizeof(typename FallbackGemm::KernelType::SharedStorage);
 
   if (gemm_shared_mem_size <= max_shared_mem_per_block_opt_in) {
-    return cutlass_gemm_caller<Gemm>(out, a, b,
-                                     std::forward<EpilogueArgs>(args)...);
+    return cutlass_gemm_caller<Gemm>(
+        out, a, b, std::forward<EpilogueArgs>(args)...);
   } else {
-    PD_CHECK(fallback_gemm_shared_mem_size <=
-                max_shared_mem_per_block_opt_in);
+    PD_CHECK(fallback_gemm_shared_mem_size <= max_shared_mem_per_block_opt_in);
     return cutlass_gemm_caller<FallbackGemm>(
         out, a, b, std::forward<EpilogueArgs>(args)...);
   }
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm75_dispatch.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm75_dispatch.cuh
index a84205f146..38994a91bd 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm75_dispatch.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm75_dispatch.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm75_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm75_dispatch.cuh
 
 #pragma once
 
@@ -11,8 +12,10 @@
 
 namespace fastdeploy {
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm75_config_default {
   // This config is used in 2 cases,
   // - M in (256, inf]
@@ -22,13 +25,21 @@ struct sm75_config_default {
   using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<8, 8, 16>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 2>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm75,
+                                        enable_sm75_to_sm80,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        2>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm75_config_M256 {
   // M in (128, 256]
   // Shared memory required by this Gemm 65536
@@ -36,13 +47,21 @@ struct sm75_config_M256 {
   using TileShape = typename cutlass::gemm::GemmShape<128, 128, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<8, 8, 16>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 2>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm75,
+                                        enable_sm75_to_sm80,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        2>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm75_config_M64 {
   // M in (32, 64]
   // Shared memory required by this Gemm 49152
@@ -50,13 +69,21 @@ struct sm75_config_M64 {
   using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<8, 8, 16>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 2>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm75,
+                                        enable_sm75_to_sm80,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        2>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm75_config_M32 {
   // M in [1, 32]
   // Shared memory required by this Gemm 49152
@@ -64,13 +91,21 @@ struct sm75_config_M32 {
   using TileShape = typename cutlass::gemm::GemmShape<32, 128, 64>;
   using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<8, 8, 16>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 2>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm75,
+                                        enable_sm75_to_sm80,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        2>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm75_dispatch(paddle::Tensor& out,
                                        paddle::Tensor const& a,
@@ -96,7 +131,8 @@ inline void cutlass_gemm_sm75_dispatch(paddle::Tensor& out,
   // sm75_config_default has the least shared-memory requirements.
   using FallbackGemm = Cutlass2xGemmDefault;
 
-  uint32_t const m = a.dims()[0];;
+  uint32_t const m = a.dims()[0];
+  ;
   uint32_t const mp2 =
       std::max(static_cast<uint32_t>(32), next_pow_2(m));  // next power of 2
   if (mp2 <= 32) {
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm80_dispatch.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm80_dispatch.cuh
index 1e3056d9b4..1c5e6eb26d 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm80_dispatch.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm80_dispatch.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm80_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm80_dispatch.cuh
 
 #pragma once
 
@@ -11,8 +12,10 @@
 
 namespace fastdeploy {
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm80_config_default {
   // This config is used in 2 cases,
   //  - M in (128, inf)
@@ -22,13 +25,21 @@ struct sm80_config_default {
   using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm80, enable_sm80_to_sm89, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm80,
+                                        enable_sm80_to_sm89,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm80_config_M64 {
   // This config is used in 2 cases,
   // - M in (32, 64]
@@ -38,13 +49,21 @@ struct sm80_config_M64 {
   using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm80, enable_sm80_to_sm89, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm80,
+                                        enable_sm80_to_sm89,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm80_config_M32 {
   // M in (16, 32]
   // Shared Memory required by this Gemm - 61440 bytes
@@ -52,13 +71,21 @@ struct sm80_config_M32 {
   using TileShape = typename cutlass::gemm::GemmShape<32, 64, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm80, enable_sm80_to_sm89, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm80,
+                                        enable_sm80_to_sm89,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm80_config_M16 {
   // M in [1, 16]
   // Shared Memory required by this Gemm - 51200 bytes
@@ -66,13 +93,21 @@ struct sm80_config_M16 {
   using TileShape = typename cutlass::gemm::GemmShape<16, 64, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<16, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm80, enable_sm80_to_sm89, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm80,
+                                        enable_sm80_to_sm89,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm80_dispatch(paddle::Tensor& out,
                                        paddle::Tensor const& a,
@@ -104,7 +139,8 @@ inline void cutlass_gemm_sm80_dispatch(paddle::Tensor& out,
   using FallbackGemm =
       typename sm80_config_M32<InType, OutType, Epilogue>::Cutlass2xGemm;
 
-  uint32_t const m = a.dims()[0];;
+  uint32_t const m = a.dims()[0];
+  ;
   uint32_t const mp2 =
       std::max(static_cast<uint32_t>(16), next_pow_2(m));  // next power of 2
   if (mp2 <= 16) {
@@ -121,7 +157,8 @@ inline void cutlass_gemm_sm80_dispatch(paddle::Tensor& out,
         out, a, b, std::forward<EpilogueArgs>(args)...);
   } else if (mp2 <= 128) {
     // M in (64, 128]
-    uint32_t const n = out.dims()[1];;
+    uint32_t const n = out.dims()[1];
+    ;
     bool const small_n = n < 8192;
     if (small_n) {
       return fallback_cutlass_gemm_caller<Cutlass2xGemmM128SmallN,
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh
index d67ebc5058..54ed72d23a 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh
 
 #pragma once
 
@@ -12,8 +13,10 @@
 
 namespace fastdeploy {
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm89_fp8_fallback_gemm {
   // Shared Memory required by this Gemm - 61440 bytes
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
@@ -21,10 +24,16 @@ struct sm89_fp8_fallback_gemm {
   using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAdd;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm89, enable_sm89_to_sm90, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5,
-                      FP8MathOperator>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm89,
+                                        enable_sm89_to_sm90,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5,
+                                        FP8MathOperator>;
 };
 
 struct sm89_fp8_config_default {
@@ -33,17 +42,21 @@ struct sm89_fp8_config_default {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -52,26 +65,47 @@ struct sm89_fp8_config_default {
       using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 8192) {
       using TileShape = typename cutlass::gemm::GemmShape<256, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
 
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -83,17 +117,21 @@ struct sm89_fp8_config_M256 {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -102,17 +140,31 @@ struct sm89_fp8_config_M256 {
       using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -124,17 +176,21 @@ struct sm89_fp8_config_M128 {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -143,26 +199,47 @@ struct sm89_fp8_config_M128 {
       using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
 
     } else if (np2 <= 16384) {
       using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<128, 64, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -172,17 +249,21 @@ struct sm89_fp8_config_M64 {
   // M in (32, 64]
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -193,9 +274,16 @@ struct sm89_fp8_config_M64 {
       using FP8MathOperator = typename cutlass::arch::OpMultiplyAdd;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
@@ -203,9 +291,16 @@ struct sm89_fp8_config_M64 {
       using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<64, 64, 128>;
@@ -213,9 +308,16 @@ struct sm89_fp8_config_M64 {
       using FP8MathOperator = typename cutlass::arch::OpMultiplyAdd;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -226,17 +328,21 @@ struct sm89_fp8_config_M32 {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -246,27 +352,48 @@ struct sm89_fp8_config_M32 {
       using WarpShape = typename cutlass::gemm::GemmShape<16, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = typename cutlass::gemm::GemmShape<32, 128, 128>;
       using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 4, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      4,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<32, 64, 128>;
       using WarpShape = typename cutlass::gemm::GemmShape<16, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -279,17 +406,21 @@ struct sm89_fp8_config_M16 {
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
   static const int32_t MainLoopStages = 5;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -298,35 +429,55 @@ struct sm89_fp8_config_M16 {
       using TileShape = typename cutlass::gemm::GemmShape<16, 64, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, MainLoopStages,
-                                FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      MainLoopStages,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 24576) {
       using TileShape = typename cutlass::gemm::GemmShape<16, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, MainLoopStages,
-                                FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      MainLoopStages,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<32, 64, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, MainLoopStages,
-                                FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      MainLoopStages,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm89_fp8_dispatch(paddle::Tensor& out,
                                            paddle::Tensor const& a,
@@ -336,7 +487,8 @@ inline void cutlass_gemm_sm89_fp8_dispatch(paddle::Tensor& out,
   PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
   PD_CHECK(b.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
-  uint32_t const m = a.dims()[0];;
+  uint32_t const m = a.dims()[0];
+  ;
   uint32_t const mp2 =
       std::max(static_cast<uint32_t>(16), next_pow_2(m));  // next power of 2
 
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh
index cd4de126b0..d8d4f9a573 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh
 
 #pragma once
 
@@ -11,8 +12,10 @@
 
 namespace fastdeploy {
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm89_int8_fallback_gemm {
   // Shared mem requirement : 61440
   static_assert(std::is_same<InType, int8_t>());
@@ -21,9 +24,15 @@ struct sm89_int8_fallback_gemm {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   static int32_t const MainLoopStages = 5;
 
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm89, enable_sm89_to_sm90, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm89,
+                                        enable_sm89_to_sm90,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
 struct sm89_int8_config_default {
@@ -31,17 +40,21 @@ struct sm89_int8_config_default {
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -50,33 +63,57 @@ struct sm89_int8_config_default {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 8192) {
       using TileShape = cutlass::gemm::GemmShape<256, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<256, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -87,17 +124,21 @@ struct sm89_int8_config_M256 {
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -106,33 +147,57 @@ struct sm89_int8_config_M256 {
       using TileShape = cutlass::gemm::GemmShape<64, 128, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 8192) {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = cutlass::gemm::GemmShape<256, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -142,17 +207,21 @@ struct sm89_int8_config_M128 {
   // M in (64, 128]
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -162,27 +231,45 @@ struct sm89_int8_config_M128 {
       using WarpShape = cutlass::gemm::GemmShape<64, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
       using WarpShape = cutlass::gemm::GemmShape<64, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<64, 64, 128>;
       using WarpShape = cutlass::gemm::GemmShape<32, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -192,17 +279,21 @@ struct sm89_int8_config_M64 {
   // M in (32, 64]
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -212,18 +303,30 @@ struct sm89_int8_config_M64 {
       using WarpShape = cutlass::gemm::GemmShape<32, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<64, 128, 128>;
       using WarpShape = cutlass::gemm::GemmShape<64, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -233,17 +336,21 @@ struct sm89_int8_config_M32 {
   // M in (16, 32]
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -253,18 +360,30 @@ struct sm89_int8_config_M32 {
       using WarpShape = cutlass::gemm::GemmShape<16, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<32, 128, 128>;
       using WarpShape = cutlass::gemm::GemmShape<32, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 4>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      4>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -275,17 +394,21 @@ struct sm89_int8_config_M16 {
   using WarpShape = typename cutlass::gemm::GemmShape<16, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[0];
     uint32_t const np2 = next_pow_2(n);
@@ -294,24 +417,38 @@ struct sm89_int8_config_M16 {
       using TileShape = cutlass::gemm::GemmShape<16, 64, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<16, 128, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 4>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      4>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm89_int8_dispatch(paddle::Tensor& out,
                                             paddle::Tensor const& a,
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c3x_sm90.cu b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c3x_sm90.cu
index bb8576abc5..9dfe781a31 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c3x_sm90.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c3x_sm90.cu
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm90.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm90.cu
 
 #include "c3x/scaled_mm_helper.hpp"
 #include "c3x/scaled_mm_kernels.hpp"
@@ -10,17 +11,24 @@
 
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
 
-void cutlass_scaled_mm_sm90(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias) {
-  dispatch_scaled_mm(c, a, b, a_scales, b_scales, bias,
+void cutlass_scaled_mm_sm90(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias) {
+  dispatch_scaled_mm(c,
+                     a,
+                     b,
+                     a_scales,
+                     b_scales,
+                     bias,
                      fastdeploy::cutlass_scaled_mm_sm90_fp8,
                      fastdeploy::cutlass_scaled_mm_sm90_int8);
 }
 
-void cutlass_scaled_mm_azp_sm90(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm90(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -30,8 +38,8 @@ void cutlass_scaled_mm_azp_sm90(paddle::Tensor& out, paddle::Tensor const& a,
   PD_CHECK(a_scales.dtype() == paddle::DataType::FLOAT32);
   PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
 
-  fastdeploy::cutlass_scaled_mm_azp_sm90_int8(out, a, b, a_scales, b_scales, azp_adj,
-                                        azp, bias);
+  fastdeploy::cutlass_scaled_mm_azp_sm90_int8(
+      out, a, b, a_scales, b_scales, azp_adj, azp, bias);
 }
 
 #endif
diff --git a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_entry.cu b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_entry.cu
index 23d07ed4d6..2565b9e189 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_entry.cu
+++ b/custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_entry.cu
@@ -1,36 +1,42 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
 
 #pragma once
 #include "helper.h"
 #include <iostream>
 
-void cutlass_scaled_mm_sm75(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias);
+void cutlass_scaled_mm_sm75(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias);
 
-void cutlass_scaled_mm_sm80(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias);
+void cutlass_scaled_mm_sm80(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias);
 
-void cutlass_scaled_mm_sm89(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias);
+void cutlass_scaled_mm_sm89(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias);
 
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
-void cutlass_scaled_mm_sm90(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias);
+void cutlass_scaled_mm_sm90(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias);
 #endif
 
-void cutlass_scaled_mm_azp_sm75(paddle::Tensor& c, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm75(paddle::Tensor& c,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -38,7 +44,8 @@ void cutlass_scaled_mm_azp_sm75(paddle::Tensor& c, paddle::Tensor const& a,
                                 paddle::optional<paddle::Tensor> const& azp,
                                 paddle::optional<paddle::Tensor> const& bias);
 
-void cutlass_scaled_mm_azp_sm80(paddle::Tensor& c, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm80(paddle::Tensor& c,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -46,7 +53,8 @@ void cutlass_scaled_mm_azp_sm80(paddle::Tensor& c, paddle::Tensor const& a,
                                 paddle::optional<paddle::Tensor> const& azp,
                                 paddle::optional<paddle::Tensor> const& bias);
 
-void cutlass_scaled_mm_azp_sm89(paddle::Tensor& c, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm89(paddle::Tensor& c,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -55,7 +63,8 @@ void cutlass_scaled_mm_azp_sm89(paddle::Tensor& c, paddle::Tensor const& a,
                                 paddle::optional<paddle::Tensor> const& bias);
 
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
-void cutlass_scaled_mm_azp_sm90(paddle::Tensor& c, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm90(paddle::Tensor& c,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -80,10 +89,12 @@ bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability) {
   return false;
 }
 
-void CutlassScaledMm(paddle::Tensor &c, paddle::Tensor const &a,
-                     paddle::Tensor const &b, paddle::Tensor const &a_scales,
-                     paddle::Tensor const &b_scales,
-                     paddle::optional<paddle::Tensor> const &bias) {
+void CutlassScaledMm(paddle::Tensor& c,
+                     paddle::Tensor const& a,
+                     paddle::Tensor const& b,
+                     paddle::Tensor const& a_scales,
+                     paddle::Tensor const& b_scales,
+                     paddle::optional<paddle::Tensor> const& bias) {
   // Checks for conformality
   PD_CHECK(a.dims().size() == 2 && b.dims().size() == 2 &&
            c.dims().size() == 2);
@@ -91,10 +102,10 @@ void CutlassScaledMm(paddle::Tensor &c, paddle::Tensor const &a,
            b.dims()[0] == c.dims()[1]);
 
   // Check for strides and alignment
-  PD_CHECK(a.strides()[1] == 1 && c.strides()[1] == 1); // Row-major
-  PD_CHECK(b.strides()[1] == 1);                        // Column-major
+  PD_CHECK(a.strides()[1] == 1 && c.strides()[1] == 1);  // Row-major
+  PD_CHECK(b.strides()[1] == 1);                         // Column-major
   PD_CHECK(c.strides()[0] % 16 == 0 &&
-           b.strides()[0] % 16 == 0); // 16 Byte Alignment
+           b.strides()[0] % 16 == 0);  // 16 Byte Alignment
 
   if (bias) {
     PD_CHECK(bias->numel() == b.dims()[0] && bias->is_contiguous() &&
@@ -138,13 +149,14 @@ void CutlassScaledMm(paddle::Tensor &c, paddle::Tensor const &a,
       version_num));
 }
 
-void CutlassScaledMmAzp(paddle::Tensor& c, paddle::Tensor const& a,
-                           paddle::Tensor const& b,
-                           paddle::Tensor const& a_scales,
-                           paddle::Tensor const& b_scales,
-                           paddle::Tensor const& azp_adj,
-                           paddle::optional<paddle::Tensor> const& azp,
-                           paddle::optional<paddle::Tensor> const& bias) {
+void CutlassScaledMmAzp(paddle::Tensor& c,
+                        paddle::Tensor const& a,
+                        paddle::Tensor const& b,
+                        paddle::Tensor const& a_scales,
+                        paddle::Tensor const& b_scales,
+                        paddle::Tensor const& azp_adj,
+                        paddle::optional<paddle::Tensor> const& azp,
+                        paddle::optional<paddle::Tensor> const& bias) {
   // Checks for conformality
   PD_CHECK(a.dims().size() == 2 && b.dims().size() == 2 &&
            c.dims().size() == 2);
@@ -154,10 +166,10 @@ void CutlassScaledMmAzp(paddle::Tensor& c, paddle::Tensor const& a,
   PD_CHECK(b_scales.numel() == 1 || b_scales.numel() == b.dims()[0]);
 
   // Check for strides and alignment
-  PD_CHECK(a.strides()[1] == 1 && c.strides()[1] == 1); // Row-major
-  PD_CHECK(b.strides()[1] == 1);                        // Column-major
+  PD_CHECK(a.strides()[1] == 1 && c.strides()[1] == 1);  // Row-major
+  PD_CHECK(b.strides()[1] == 1);                         // Column-major
   PD_CHECK(c.strides()[0] % 16 == 0 &&
-           b.strides()[0] % 16 == 0); // 16 Byte Alignment
+           b.strides()[0] % 16 == 0);  // 16 Byte Alignment
   PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
 
   // bias, azp, azp_adj are all 1d
@@ -174,7 +186,8 @@ void CutlassScaledMmAzp(paddle::Tensor& c, paddle::Tensor const& a,
   PD_CHECK(azp_adj.dtype() == paddle::DataType::INT32);
   PD_CHECK(!azp || azp->dtype() == paddle::DataType::INT32);
   PD_CHECK(!bias || bias->dtype() == c.dtype(),
-              "currently bias dtype must match output dtype ", c.dtype());
+           "currently bias dtype must match output dtype ",
+           c.dtype());
 
   int32_t version_num = GetGPUComputeCapability(a.place().GetDeviceId());
 
@@ -210,7 +223,6 @@ void CutlassScaledMmAzp(paddle::Tensor& c, paddle::Tensor const& a,
       version_num));
 }
 
-
 PD_BUILD_STATIC_OP(cutlass_scaled_mm)
     .Inputs({"c", "a", "b", "a_scales", "b_scales", paddle::Optional("bias")})
     .Outputs({"c_out"})
@@ -218,7 +230,14 @@ PD_BUILD_STATIC_OP(cutlass_scaled_mm)
     .SetKernelFn(PD_KERNEL(CutlassScaledMm));
 
 PD_BUILD_STATIC_OP(cutlass_scaled_mm_azp)
-    .Inputs({"c", "a", "b", "a_scales", "b_scales", "azp_adj", paddle::Optional("azp"), paddle::Optional("bias")})
+    .Inputs({"c",
+             "a",
+             "b",
+             "a_scales",
+             "b_scales",
+             "azp_adj",
+             paddle::Optional("azp"),
+             paddle::Optional("bias")})
     .Outputs({"c_out"})
     .SetInplaceMap({{"c", "c_out"}})
     .SetKernelFn(PD_KERNEL(CutlassScaledMmAzp));
diff --git a/custom_ops/gpu_ops/cutlass_kernels/weight_process_utils.h b/custom_ops/gpu_ops/cutlass_kernels/weight_process_utils.h
index 0a26cb6c92..01a9201a0e 100644
--- a/custom_ops/gpu_ops/cutlass_kernels/weight_process_utils.h
+++ b/custom_ops/gpu_ops/cutlass_kernels/weight_process_utils.h
@@ -33,484 +33,498 @@ limitations under the License. */
 
 void row_major_to_column_major(int8_t* col_major_tensor,
                                const int8_t* row_major_tensor,
-                               const std::vector<size_t>& shape){
-    size_t m = shape[0];
-    size_t n = shape[1];
-    for(auto i=0;i<m*n;i++){
-        size_t im = i / n;
-        size_t in = i % n;
-        col_major_tensor[in*m+im] = row_major_tensor[im*n+in];
-    }
+                               const std::vector<size_t>& shape) {
+  size_t m = shape[0];
+  size_t n = shape[1];
+  for (auto i = 0; i < m * n; i++) {
+    size_t im = i / n;
+    size_t in = i % n;
+    col_major_tensor[in * m + im] = row_major_tensor[im * n + in];
+  }
 }
 
 void add_bias_and_interleave_int8s_inplace(int8_t* int8_tensor_ptr,
-                                           int64_t num_elts)
-{
-    int8_t* int8_tensor = reinterpret_cast<int8_t *>(int8_tensor_ptr);
-    for (int ii = 0; ii < num_elts; ++ii) {
-        int8_tensor[ii] = int8_t(int(int8_tensor[ii]) + 128);
-        // int8_tensor[ii] = int8_t(int(int8_tensor[ii]));
-    }
+                                           int64_t num_elts) {
+  int8_t* int8_tensor = reinterpret_cast<int8_t*>(int8_tensor_ptr);
+  for (int ii = 0; ii < num_elts; ++ii) {
+    int8_tensor[ii] = int8_t(int(int8_tensor[ii]) + 128);
+    // int8_tensor[ii] = int8_t(int(int8_tensor[ii]));
+  }
 
-    // Step 2 will transform the layout of a 32-bit register in CUDA in order to match the int4 layout. This has no
-    // performance benefit and is purely so that int4 and int8 have the same layout.
-    // Pictorially, this does the following:
-    // bit 32                                                      0
-    //      [elt_3  elt_2  elt_1  elt_0] (each elt occupies 8 bits)
-    //
-    // And it will rearrange the output 32 bit register to be the following:
-    // bit 32                                                      0
-    //      [elt_3  elt_1  elt_2  elt_0] (each elt occupies 8 bits)
+  // Step 2 will transform the layout of a 32-bit register in CUDA in order to
+  // match the int4 layout. This has no performance benefit and is purely so
+  // that int4 and int8 have the same layout. Pictorially, this does the
+  // following: bit 32                                                      0
+  //      [elt_3  elt_2  elt_1  elt_0] (each elt occupies 8 bits)
+  //
+  // And it will rearrange the output 32 bit register to be the following:
+  // bit 32                                                      0
+  //      [elt_3  elt_1  elt_2  elt_0] (each elt occupies 8 bits)
 
-    for (int64_t base = 0; base < num_elts; base += 4) {
-        std::swap(int8_tensor[base + 1], int8_tensor[base + 2]);
-    }
+  for (int64_t base = 0; base < num_elts; base += 4) {
+    std::swap(int8_tensor[base + 1], int8_tensor[base + 2]);
+  }
 }
 
+void subbyte_transpose_impl_int4(int8_t* transposed_quantized_tensor,
+                                 const int8_t* quantized_tensor,
+                                 const std::vector<size_t>& shape) {
+  const int bits_per_elt = 4;
 
-void subbyte_transpose_impl_int4(int8_t*                    transposed_quantized_tensor,
-                            const int8_t*              quantized_tensor,
-                            const std::vector<size_t>& shape)
-{
-    const int bits_per_elt = 4;
+  const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    const size_t num_experts = shape.size() == 2 ? 1 : shape[0];
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const size_t col_bytes = num_cols * bits_per_elt / 8;
+  const size_t col_bytes_trans = num_rows * bits_per_elt / 8;
+  const size_t num_bytes = size_t(num_experts) * num_rows * col_bytes;
 
-    const size_t col_bytes       = num_cols * bits_per_elt / 8;
-    const size_t col_bytes_trans = num_rows * bits_per_elt / 8;
-    const size_t num_bytes       = size_t(num_experts) * num_rows * col_bytes;
+  const uint8_t* input_byte_ptr =
+      reinterpret_cast<const uint8_t*>(quantized_tensor);
+  uint8_t* output_byte_ptr =
+      reinterpret_cast<uint8_t*>(transposed_quantized_tensor);
 
-    const uint8_t* input_byte_ptr  = reinterpret_cast<const uint8_t*>(quantized_tensor);
-    uint8_t*       output_byte_ptr = reinterpret_cast<uint8_t*>(transposed_quantized_tensor);
+  // static_assert(quant_type == QuantType::INT8_WEIGHT_ONLY || quant_type ==
+  // QuantType::PACKED_INT4_WEIGHT_ONLY, "");
+  static constexpr int ELTS_PER_BYTE = 2;
 
-    // static_assert(quant_type == QuantType::INT8_WEIGHT_ONLY || quant_type == QuantType::PACKED_INT4_WEIGHT_ONLY, "");
-    static constexpr int ELTS_PER_BYTE = 2;
+  static constexpr int M_TILE_L1 = 64;
+  static constexpr int N_TILE_L1 = M_TILE_L1 / ELTS_PER_BYTE;
+  uint8_t cache_buf[M_TILE_L1][N_TILE_L1];
 
-    static constexpr int M_TILE_L1 = 64;
-    static constexpr int N_TILE_L1 = M_TILE_L1 / ELTS_PER_BYTE;
-    uint8_t              cache_buf[M_TILE_L1][N_TILE_L1];
+  static constexpr int VECTOR_WIDTH = std::min(32, N_TILE_L1);
 
-    static constexpr int VECTOR_WIDTH = std::min(32, N_TILE_L1);
+  // We assume the dims are a multiple of vector width. Our kernels only handle
+  // dims which are multiples of 64 for weight-only quantization. As a result,
+  // this seemed like a reasonable tradeoff because it allows GCC to emit vector
+  // instructions.
 
-    // We assume the dims are a multiple of vector width. Our kernels only handle dims which are multiples
-    // of 64 for weight-only quantization. As a result, this seemed like a reasonable tradeoff because it
-    // allows GCC to emit vector instructions.
+  const int num_m_tiles = (num_rows + M_TILE_L1 - 1) / M_TILE_L1;
+  const int num_n_tiles = (col_bytes + N_TILE_L1 - 1) / N_TILE_L1;
 
-    const int num_m_tiles = (num_rows + M_TILE_L1 - 1) / M_TILE_L1;
-    const int num_n_tiles = (col_bytes + N_TILE_L1 - 1) / N_TILE_L1;
+  for (size_t expert = 0; expert < num_experts; ++expert) {
+    const size_t matrix_offset = expert * num_rows * col_bytes;
+    for (size_t row_tile_start = 0; row_tile_start < num_rows;
+         row_tile_start += M_TILE_L1) {
+      for (size_t col_tile_start_byte = 0; col_tile_start_byte < col_bytes;
+           col_tile_start_byte += N_TILE_L1) {
+        const int row_limit = std::min(row_tile_start + M_TILE_L1, num_rows);
+        const int col_limit =
+            std::min(col_tile_start_byte + N_TILE_L1, col_bytes);
 
-    for (size_t expert = 0; expert < num_experts; ++expert) {
-        const size_t matrix_offset = expert * num_rows * col_bytes;
-        for (size_t row_tile_start = 0; row_tile_start < num_rows; row_tile_start += M_TILE_L1) {
-            for (size_t col_tile_start_byte = 0; col_tile_start_byte < col_bytes; col_tile_start_byte += N_TILE_L1) {
+        for (int ii = 0; ii < M_TILE_L1; ++ii) {
+          const int row = row_tile_start + ii;
 
-                const int row_limit = std::min(row_tile_start + M_TILE_L1, num_rows);
-                const int col_limit = std::min(col_tile_start_byte + N_TILE_L1, col_bytes);
+          for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
+            const int col = col_tile_start_byte + jj;
 
-                for (int ii = 0; ii < M_TILE_L1; ++ii) {
-                    const int row = row_tile_start + ii;
+            const size_t logical_src_offset =
+                matrix_offset + row * col_bytes + col;
 
-                    for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
-                        const int col = col_tile_start_byte + jj;
-
-                        const size_t logical_src_offset = matrix_offset + row * col_bytes + col;
-
-                        if (row < row_limit && col < col_limit) {
-                            for (int v = 0; v < VECTOR_WIDTH; ++v) {
-                                cache_buf[ii][jj + v] = input_byte_ptr[logical_src_offset + v];
-                            }
-                        }
-                    }
-                }
-
-
-                for (int ii = 0; ii < M_TILE_L1; ++ii) {
-                    // Using M_TILE_L1 here is deliberate since we assume that the cache tile
-                    // is square in the number of elements (not necessarily the number of bytes).
-                    for (int jj = ii + 1; jj < M_TILE_L1; ++jj) {
-                        const int ii_byte       = ii / ELTS_PER_BYTE;
-                        const int ii_bit_offset = ii % ELTS_PER_BYTE;
-
-                        const int jj_byte       = jj / ELTS_PER_BYTE;
-                        const int jj_bit_offset = jj % ELTS_PER_BYTE;
-
-                        uint8_t src_elt = 0xF & (cache_buf[ii][jj_byte] >> (4 * jj_bit_offset));
-                        uint8_t tgt_elt = 0xF & (cache_buf[jj][ii_byte] >> (4 * ii_bit_offset));
-
-                        cache_buf[ii][jj_byte] &= (0xF0 >> (4 * jj_bit_offset));
-                        cache_buf[jj][ii_byte] &= (0xF0 >> (4 * ii_bit_offset));
-
-                        cache_buf[ii][jj_byte] |= (tgt_elt << (4 * jj_bit_offset));
-                        cache_buf[jj][ii_byte] |= (src_elt << (4 * ii_bit_offset));
-                    }
-                }
-
-
-                const size_t row_tile_start_trans      = col_tile_start_byte * ELTS_PER_BYTE;
-                const size_t col_tile_start_byte_trans = row_tile_start / ELTS_PER_BYTE;
-
-                const int row_limit_trans = std::min(row_tile_start_trans + M_TILE_L1, num_cols);
-                const int col_limit_trans = std::min(col_tile_start_byte_trans + N_TILE_L1, col_bytes_trans);
-
-                for (int ii = 0; ii < M_TILE_L1; ++ii) {
-                    const int row = row_tile_start_trans + ii;
-                    for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
-                        const int col = col_tile_start_byte_trans + jj;
-
-                        const size_t logical_tgt_offset = matrix_offset + row * col_bytes_trans + col;
-
-                        if (row < row_limit_trans && col < col_limit_trans) {
-                            for (int v = 0; v < VECTOR_WIDTH; ++v) {
-                                output_byte_ptr[logical_tgt_offset + v] = cache_buf[ii][jj + v];
-                            }
-                        }
-                    }
-                }
+            if (row < row_limit && col < col_limit) {
+              for (int v = 0; v < VECTOR_WIDTH; ++v) {
+                cache_buf[ii][jj + v] = input_byte_ptr[logical_src_offset + v];
+              }
             }
+          }
         }
+
+        for (int ii = 0; ii < M_TILE_L1; ++ii) {
+          // Using M_TILE_L1 here is deliberate since we assume that the cache
+          // tile is square in the number of elements (not necessarily the
+          // number of bytes).
+          for (int jj = ii + 1; jj < M_TILE_L1; ++jj) {
+            const int ii_byte = ii / ELTS_PER_BYTE;
+            const int ii_bit_offset = ii % ELTS_PER_BYTE;
+
+            const int jj_byte = jj / ELTS_PER_BYTE;
+            const int jj_bit_offset = jj % ELTS_PER_BYTE;
+
+            uint8_t src_elt =
+                0xF & (cache_buf[ii][jj_byte] >> (4 * jj_bit_offset));
+            uint8_t tgt_elt =
+                0xF & (cache_buf[jj][ii_byte] >> (4 * ii_bit_offset));
+
+            cache_buf[ii][jj_byte] &= (0xF0 >> (4 * jj_bit_offset));
+            cache_buf[jj][ii_byte] &= (0xF0 >> (4 * ii_bit_offset));
+
+            cache_buf[ii][jj_byte] |= (tgt_elt << (4 * jj_bit_offset));
+            cache_buf[jj][ii_byte] |= (src_elt << (4 * ii_bit_offset));
+          }
+        }
+
+        const size_t row_tile_start_trans = col_tile_start_byte * ELTS_PER_BYTE;
+        const size_t col_tile_start_byte_trans = row_tile_start / ELTS_PER_BYTE;
+
+        const int row_limit_trans =
+            std::min(row_tile_start_trans + M_TILE_L1, num_cols);
+        const int col_limit_trans =
+            std::min(col_tile_start_byte_trans + N_TILE_L1, col_bytes_trans);
+
+        for (int ii = 0; ii < M_TILE_L1; ++ii) {
+          const int row = row_tile_start_trans + ii;
+          for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
+            const int col = col_tile_start_byte_trans + jj;
+
+            const size_t logical_tgt_offset =
+                matrix_offset + row * col_bytes_trans + col;
+
+            if (row < row_limit_trans && col < col_limit_trans) {
+              for (int v = 0; v < VECTOR_WIDTH; ++v) {
+                output_byte_ptr[logical_tgt_offset + v] = cache_buf[ii][jj + v];
+              }
+            }
+          }
+        }
+      }
     }
+  }
 }
 
+void add_bias_and_interleave_int4s_inplace(int8_t* packed_int4_tensor,
+                                           const size_t num_elts) {
+  const int num_bytes = num_elts / 2;
 
-void add_bias_and_interleave_int4s_inplace(int8_t* packed_int4_tensor, const size_t num_elts)
-{
-    const int num_bytes = num_elts / 2;
+  // Step 1 will be to transform all the int4s to unsigned in order to make the
+  // dequantize take as little instructions as possible in the CUDA code.
+  for (size_t ii = 0; ii < num_bytes; ++ii) {
+    int8_t transformed_packed_int4s = 0;
+    // We don't need to mask in these ops since everything should be in the
+    // range 0-15
+    int8_t transformed_first_elt = (packed_int4_tensor[ii] & 0x0F);
+    int8_t transformed_second_elt = (packed_int4_tensor[ii] >> 4);
 
-    // Step 1 will be to transform all the int4s to unsigned in order to make the dequantize take as little
-    // instructions as possible in the CUDA code.
-    for (size_t ii = 0; ii < num_bytes; ++ii) {
-        int8_t transformed_packed_int4s = 0;
-        // We don't need to mask in these ops since everything should be in the range 0-15
-        int8_t transformed_first_elt = (packed_int4_tensor[ii] & 0x0F);
-        int8_t transformed_second_elt = (packed_int4_tensor[ii] >> 4);
+    transformed_packed_int4s |= transformed_first_elt;
+    transformed_packed_int4s |= (transformed_second_elt << 4);
+    packed_int4_tensor[ii] = transformed_packed_int4s;
+  }
 
-        transformed_packed_int4s |= transformed_first_elt;
-        transformed_packed_int4s |= (transformed_second_elt << 4);
-        packed_int4_tensor[ii] = transformed_packed_int4s;
-    }
-
-    // Step 2 will transform the layout of a 32-bit register in CUDA in order to minimize the number of shift & logical
-    // instructions That are needed to extract the int4s in the GEMM main loop. Pictorially, the loop below will do the
-    // following: Take as input a 32 bit register with layout: bit 32 0
-    //      [elt_7  elt_6  elt_5  elt_4  elt_3  elt_2  elt_1  elt_0] (each elt occupies 4 bits)
-    //
-    // And it will rearrange the output 32 bit register to be the following:
-    // bit 32                                                      0
-    //      [elt_7  elt_5  elt_3  elt_1  elt_6  elt_4  elt_2  elt_0] (each elt occupies 4 bits)
-
-    // FT_CHECK_WITH_INFO(num_bytes % 4 == 0, "Dimensions of int4 tensor must be a multiple of 8 for register relayout");
-    const size_t num_registers = num_bytes / 4;
-
-    uint32_t* register_ptr = reinterpret_cast<uint32_t*>(packed_int4_tensor);
-    for (size_t ii = 0; ii < num_registers; ++ii) {
-        const uint32_t current_register     = register_ptr[ii];
-        uint32_t       transformed_register = 0;
-
-        for (int dest_idx = 0; dest_idx < 8; ++dest_idx) {
-            const int src_idx    = dest_idx < 4 ? 2 * dest_idx : 2 * (dest_idx - 4) + 1;
-            const int src_shift  = 4 * src_idx;
-            const int dest_shift = 4 * dest_idx;
-
-            const uint32_t src_bits = (current_register >> src_shift) & 0xF;
-            transformed_register |= (src_bits << dest_shift);
-
-        }
-        register_ptr[ii] = transformed_register;
+  // Step 2 will transform the layout of a 32-bit register in CUDA in order to
+  // minimize the number of shift & logical instructions That are needed to
+  // extract the int4s in the GEMM main loop. Pictorially, the loop below will
+  // do the following: Take as input a 32 bit register with layout: bit 32 0
+  //      [elt_7  elt_6  elt_5  elt_4  elt_3  elt_2  elt_1  elt_0] (each elt
+  //      occupies 4 bits)
+  //
+  // And it will rearrange the output 32 bit register to be the following:
+  // bit 32                                                      0
+  //      [elt_7  elt_5  elt_3  elt_1  elt_6  elt_4  elt_2  elt_0] (each elt
+  //      occupies 4 bits)
+
+  // FT_CHECK_WITH_INFO(num_bytes % 4 == 0, "Dimensions of int4 tensor must be a
+  // multiple of 8 for register relayout");
+  const size_t num_registers = num_bytes / 4;
+
+  uint32_t* register_ptr = reinterpret_cast<uint32_t*>(packed_int4_tensor);
+  for (size_t ii = 0; ii < num_registers; ++ii) {
+    const uint32_t current_register = register_ptr[ii];
+    uint32_t transformed_register = 0;
+
+    for (int dest_idx = 0; dest_idx < 8; ++dest_idx) {
+      const int src_idx = dest_idx < 4 ? 2 * dest_idx : 2 * (dest_idx - 4) + 1;
+      const int src_shift = 4 * src_idx;
+      const int dest_shift = 4 * dest_idx;
+
+      const uint32_t src_bits = (current_register >> src_shift) & 0xF;
+      transformed_register |= (src_bits << dest_shift);
     }
+    register_ptr[ii] = transformed_register;
+  }
 }
 
-void permute_B_rows_for_mixed_and_int8_gemm(int8_t*                    permuted_quantized_tensor,
-                                            const int8_t*              quantized_tensor,
+void permute_B_rows_for_mixed_and_int8_gemm(int8_t* permuted_quantized_tensor,
+                                            const int8_t* quantized_tensor,
                                             const std::vector<size_t>& shape,
-                                            const int64_t              arch_version)
-{
+                                            const int64_t arch_version) {
+  // We only want to run this step for weight only quant.
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    // We only want to run this step for weight only quant.
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const int BITS_PER_ELT = 8;
+  const int K = 16 / BITS_PER_ELT;
+  const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;
+  const int ELTS_PER_REG = 32 / BITS_PER_ELT;
 
-    const int BITS_PER_ELT  = 8;
-    const int K             = 16 / BITS_PER_ELT;
-    const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;
-    const int ELTS_PER_REG  = 32 / BITS_PER_ELT;
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
+  int MMA_SHAPE_N = 8;
+  int B_ROWS_PER_MMA = 8 * K;
+  const int elts_in_int32 = 32 / BITS_PER_ELT;
 
-    int       MMA_SHAPE_N    = 8;
-    int       B_ROWS_PER_MMA = 8 * K;
-    const int elts_in_int32  = 32 / BITS_PER_ELT;
+  const int num_vec_cols = num_cols / elts_in_int32;
 
-    const int num_vec_cols = num_cols / elts_in_int32;
+  // The code is written as below so it works for both int8 and packed int4.
+  for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
+    for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+      for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
+        const int write_row = base_row + tile_row;
+        const int tile_read_row = 4 * (((tile_row % ELTS_PER_REG) / 2)) +
+                                  tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
 
-    // The code is written as below so it works for both int8 and packed int4.
-    for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
-        for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+        const int read_row = base_row + tile_read_row;
+        const int read_col = write_col;
 
-            for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
-                const int write_row = base_row + tile_row;
-                const int tile_read_row =
-                    4 * (((tile_row % ELTS_PER_REG) / 2)) + tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
+        const int64_t read_offset = int64_t(read_row) * num_vec_cols + read_col;
+        const int64_t write_offset =
+            int64_t(write_row) * num_vec_cols + write_col;
 
-                const int read_row = base_row + tile_read_row;
-                const int read_col = write_col;
-
-                const int64_t read_offset  = int64_t(read_row) * num_vec_cols + read_col;
-                const int64_t write_offset = int64_t(write_row) * num_vec_cols + write_col;
-
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
-        }
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
 
-// Permutes the rows of B for Turing and Ampere. Throws an error for other architectures.
-// The data is permuted such that:
-// For int8, each group of 16 rows is permuted using the map below:
+// Permutes the rows of B for Turing and Ampere. Throws an error for other
+// architectures. The data is permuted such that: For int8, each group of 16
+// rows is permuted using the map below:
 //  0 1 8 9 2 3 10 11 4 5 12 13 6 7 14 15
 //  0 1 2 3 4 5 6 7
-template<int bits=8>
-void permute_B_rows_for_mixed_gemm(int8_t*                    permuted_quantized_tensor,
-                                   const int8_t*              quantized_tensor,
+template <int bits = 8>
+void permute_B_rows_for_mixed_gemm(int8_t* permuted_quantized_tensor,
+                                   const int8_t* quantized_tensor,
                                    const std::vector<size_t>& shape,
-                                   const int64_t              arch_version)
-{
+                                   const int64_t arch_version) {
+  // We only want to run this step for weight only quant.
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    // We only want to run this step for weight only quant.
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const int BITS_PER_ELT = bits;
+  const int K = 16 / BITS_PER_ELT;
+  const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;
+  const int ELTS_PER_REG = 32 / BITS_PER_ELT;
 
-    const int BITS_PER_ELT  = bits;
-    const int K             = 16 / BITS_PER_ELT;
-    const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;
-    const int ELTS_PER_REG  = 32 / BITS_PER_ELT;
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
+  int MMA_SHAPE_N = 8;
+  int B_ROWS_PER_MMA = 8 * K;
+  const int elts_in_int32 = 32 / BITS_PER_ELT;
 
-    int       MMA_SHAPE_N    = 8;
-    int       B_ROWS_PER_MMA = 8 * K;
-    const int elts_in_int32  = 32 / BITS_PER_ELT;
+  const int num_vec_cols = num_cols / elts_in_int32;
 
-    const int num_vec_cols = num_cols / elts_in_int32;
-
-    // The code is written as below so it works for both int8 and packed int4.
-    for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
-        for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
-
-            for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
-                const int write_row = base_row + tile_row;
-                const int tile_read_row =
-                    8 * (((tile_row % ELTS_PER_REG) / 2)) + tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
-                if(base_row == 0 && write_col == 0){
-                    std::cout<<"tile_read_row:"<<tile_read_row<<std::endl;
-                }
-                const int read_row = base_row + tile_read_row;
-                const int read_col = write_col;
-
-                const int64_t read_offset  = int64_t(read_row) * num_vec_cols + read_col;
-                const int64_t write_offset = int64_t(write_row) * num_vec_cols + write_col;
-
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
+  // The code is written as below so it works for both int8 and packed int4.
+  for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
+    for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+      for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
+        const int write_row = base_row + tile_row;
+        const int tile_read_row = 8 * (((tile_row % ELTS_PER_REG) / 2)) +
+                                  tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
+        if (base_row == 0 && write_col == 0) {
+          std::cout << "tile_read_row:" << tile_read_row << std::endl;
         }
+        const int read_row = base_row + tile_read_row;
+        const int read_col = write_col;
+
+        const int64_t read_offset = int64_t(read_row) * num_vec_cols + read_col;
+        const int64_t write_offset =
+            int64_t(write_row) * num_vec_cols + write_col;
+
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
 
-template<int bits=4>
-void permute_B_rows_for_mixed_gemm_int4(int8_t*                    permuted_quantized_tensor,
-                                   const int8_t*              quantized_tensor,
-                                   const std::vector<size_t>& shape,
-                                   const int64_t              arch_version)
-{
+template <int bits = 4>
+void permute_B_rows_for_mixed_gemm_int4(int8_t* permuted_quantized_tensor,
+                                        const int8_t* quantized_tensor,
+                                        const std::vector<size_t>& shape,
+                                        const int64_t arch_version) {
+  // We only want to run this step for weight only quant.
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-    // We only want to run this step for weight only quant.
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const int BITS_PER_ELT = bits;               // 4
+  const int K = 16 / BITS_PER_ELT;             // 4
+  const int ELTS_PER_BYTE = 8 / BITS_PER_ELT;  // 2
+  const int ELTS_PER_REG = 32 / BITS_PER_ELT;  // 8
 
-    const int BITS_PER_ELT  = bits; //4
-    const int K             = 16 / BITS_PER_ELT; // 4
-    const int ELTS_PER_BYTE = 8 / BITS_PER_ELT; // 2
-    const int ELTS_PER_REG  = 32 / BITS_PER_ELT; // 8
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(permuted_quantized_tensor);
+  int MMA_SHAPE_N = 8;
+  int B_ROWS_PER_MMA = 8 * K;  // 32
+  const int elts_in_int32 = 32 / BITS_PER_ELT;
 
-    int       MMA_SHAPE_N    = 8;
-    int       B_ROWS_PER_MMA = 8 * K; // 32
-    const int elts_in_int32  = 32 / BITS_PER_ELT;
+  const int num_vec_cols = num_cols / elts_in_int32;
+  const std::vector<int> tile_col_map{0,  2,  16, 18, 1,  3,  17, 19, 4,  6, 20,
+                                      22, 5,  7,  21, 23, 8,  10, 24, 26, 9, 11,
+                                      25, 27, 12, 14, 28, 30, 13, 15, 29, 31};
 
-    const int num_vec_cols = num_cols / elts_in_int32;
-    const std::vector<int> tile_col_map{
-                                       0,2,16,18,
-                                       1,3,17,19,
-                                       4,6,20,22,
-                                       5,7,21,23,
-                                       8,10,24,26,
-                                       9,11,25,27,
-                                       12,14,28,30,
-                                       13,15,29,31};
+  // const std::vector<int> tile_col_map{
+  //                   0                   0,2,16,18,
+  //                   4                   1,3,17,19,
+  //                   8                   4,6,20,22,
+  //                   12                  5,7,21,23,
+  //                   16                  8,10,24,26,
+  //                   20                  9,11,25,27,
+  //                   24                  12,14,28,30,
+  //                   28                  13,15,29,31};
+  // std::vector<int> tile_col_map(32);
+  // for(int i=0;i<32;i++){
+  //     tile_col_map[i]=i;
+  // }
+  // // tile_col_map[1]=4;
+  // tile_col_map[0]=0;
+  // tile_col_map[4]=1;
+  // tile_col_map[1]=2;
+  // tile_col_map[5]=3;
+  // tile_col_map[8]=4;
+  // tile_col_map[12]=5;
+  // tile_col_map[9]=6;
+  // tile_col_map[13]=7;
+  // tile_col_map[16]=8;
+  // tile_col_map[20]=9;
+  // tile_col_map[17]=10;
+  // tile_col_map[21]=11;
+  // tile_col_map[24]=12;
+  // tile_col_map[28]=13;
+  // tile_col_map[25]=14;
+  // tile_col_map[29]=15;
 
-    // const std::vector<int> tile_col_map{
-    //                   0                   0,2,16,18,
-    //                   4                   1,3,17,19,
-    //                   8                   4,6,20,22,
-    //                   12                  5,7,21,23,
-    //                   16                  8,10,24,26,
-    //                   20                  9,11,25,27,
-    //                   24                  12,14,28,30,
-    //                   28                  13,15,29,31};
-    // std::vector<int> tile_col_map(32);
-    // for(int i=0;i<32;i++){
-    //     tile_col_map[i]=i;
-    // }
-    // // tile_col_map[1]=4;
-    // tile_col_map[0]=0;
-    // tile_col_map[4]=1;
-    // tile_col_map[1]=2;
-    // tile_col_map[5]=3;
-    // tile_col_map[8]=4;
-    // tile_col_map[12]=5;
-    // tile_col_map[9]=6;
-    // tile_col_map[13]=7;
-    // tile_col_map[16]=8;
-    // tile_col_map[20]=9;
-    // tile_col_map[17]=10;
-    // tile_col_map[21]=11;
-    // tile_col_map[24]=12;
-    // tile_col_map[28]=13;
-    // tile_col_map[25]=14;
-    // tile_col_map[29]=15;
+  // tile_col_map[4]=1;
+  // tile_col_map[4]=1;
+  // tile_col_map[4]=2;
 
-    // tile_col_map[4]=1;
-    // tile_col_map[4]=1;
-    // tile_col_map[4]=2;
-
-    // The code is written as below so it works for both int8 and packed int4.
-    for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
-        for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
-
-            for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
-                const int write_row = base_row + tile_row;
-                // const int tile_read_row =
-                //     8 * (((tile_row % ELTS_PER_REG) / 2)) + tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
-                // const int tile_read_row = std::distance(tile_col_map.begin(), std::find(tile_col_map.begin(),tile_col_map.end(), tile_row));
-                const int tile_read_row = tile_col_map[tile_row];
-                if(base_row == 0 && write_col == 0){
-                    std::cout<<" write_row:"<<tile_row<<" tile_read_row:"<<tile_read_row<<std::endl;
-                }
-                const int read_row = base_row + tile_read_row;
-                const int read_col = write_col;
-
-                const int64_t read_offset  = int64_t(read_row) * num_vec_cols + read_col;
-                const int64_t write_offset = int64_t(write_row) * num_vec_cols + write_col;
-
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
+  // The code is written as below so it works for both int8 and packed int4.
+  for (int base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
+    for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
+      for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
+        const int write_row = base_row + tile_row;
+        // const int tile_read_row =
+        //     8 * (((tile_row % ELTS_PER_REG) / 2)) + tile_row % 2 + 2 *
+        //     (tile_row / ELTS_PER_REG);
+        // const int tile_read_row = std::distance(tile_col_map.begin(),
+        // std::find(tile_col_map.begin(),tile_col_map.end(), tile_row));
+        const int tile_read_row = tile_col_map[tile_row];
+        if (base_row == 0 && write_col == 0) {
+          std::cout << " write_row:" << tile_row
+                    << " tile_read_row:" << tile_read_row << std::endl;
         }
+        const int read_row = base_row + tile_read_row;
+        const int read_col = write_col;
+
+        const int64_t read_offset = int64_t(read_row) * num_vec_cols + read_col;
+        const int64_t write_offset =
+            int64_t(write_row) * num_vec_cols + write_col;
+
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
 
+void interleave_column_major_tensor(int8_t* interleaved_quantized_tensor,
+                                    const int8_t* quantized_tensor,
+                                    const std::vector<size_t>& shape) {
+  // We only want to run this step for weight only quant.
+  std::cout << "### in interleave_column_major_tensor" << std::endl;
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-void interleave_column_major_tensor(int8_t*                    interleaved_quantized_tensor,
-                                    const int8_t*              quantized_tensor,
-                                    const std::vector<size_t>& shape)
-{
+  const size_t BITS_PER_ELT = 8;
+  const size_t elts_in_int32 = 32 / BITS_PER_ELT;
 
-    // We only want to run this step for weight only quant.
-    std::cout<<"### in interleave_column_major_tensor"<<std::endl;
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const size_t rows_per_tile = 64;
+  std::cout << "running interleave_column_major_tensor" << std::endl;
+  std::cout << "num_rows:" << num_rows << ","
+            << "num_cols:" << num_cols << ","
+            << "BITS_PER_ELT:" << BITS_PER_ELT << ","
+            << "elts_in_int32:" << elts_in_int32 << ","
+            << "rows_per_tile:" << rows_per_tile << std::endl;
 
-    const size_t BITS_PER_ELT  = 8;
-    const size_t elts_in_int32 = 32 / BITS_PER_ELT;
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
 
-    const size_t rows_per_tile = 64;
-    std::cout<<"running interleave_column_major_tensor"<<std::endl;
-    std::cout<<"num_rows:"<<num_rows<<","
-             <<"num_cols:"<<num_cols<<","
-             <<"BITS_PER_ELT:"<<BITS_PER_ELT<<","
-             <<"elts_in_int32:"<<elts_in_int32<<","
-             <<"rows_per_tile:"<<rows_per_tile<<std::endl;
+  const size_t num_vec_rows = num_rows / elts_in_int32;
+  const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
+  const size_t interleave = 2;
+  std::cout << "num_vec_rows:" << num_vec_rows << ","
+            << "vec_rows_per_tile:" << vec_rows_per_tile << ","
+            << "interleave:" << interleave << std::endl;
+  for (int read_col = 0; read_col < num_cols; ++read_col) {
+    const size_t write_col = read_col / interleave;
+    for (int base_vec_row = 0; base_vec_row < num_vec_rows;
+         base_vec_row += vec_rows_per_tile) {
+      for (int vec_read_row = base_vec_row;
+           vec_read_row <
+           std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
+           ++vec_read_row) {
+        const size_t vec_write_row =
+            interleave * base_vec_row +
+            vec_rows_per_tile * (read_col % interleave) +
+            vec_read_row % vec_rows_per_tile;
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
-
-
-    const size_t num_vec_rows      = num_rows / elts_in_int32;
-    const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
-    const size_t interleave        = 2;
-    std::cout<<"num_vec_rows:"<<num_vec_rows<<","
-             <<"vec_rows_per_tile:"<<vec_rows_per_tile<<","
-             <<"interleave:"<<interleave<<std::endl;
-    for (int read_col = 0; read_col < num_cols; ++read_col) {
-        const size_t write_col = read_col / interleave;
-        for (int base_vec_row = 0; base_vec_row < num_vec_rows; base_vec_row += vec_rows_per_tile) {
-            for (int vec_read_row = base_vec_row;
-                    vec_read_row < std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
-                    ++vec_read_row) {
-                const size_t vec_write_row = interleave * base_vec_row
-                                                + vec_rows_per_tile * (read_col % interleave)
-                                                + vec_read_row % vec_rows_per_tile;
-
-                const size_t read_offset = size_t(read_col) * num_vec_rows + vec_read_row;
-                const size_t write_offset = size_t(write_col) * num_vec_rows * interleave + vec_write_row;
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
-        }
+        const size_t read_offset =
+            size_t(read_col) * num_vec_rows + vec_read_row;
+        const size_t write_offset =
+            size_t(write_col) * num_vec_rows * interleave + vec_write_row;
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
 
+void interleave_column_major_tensor_int4(int8_t* interleaved_quantized_tensor,
+                                         const int8_t* quantized_tensor,
+                                         const std::vector<size_t>& shape) {
+  // We only want to run this step for weight only quant.
+  std::cout << "### in interleave_column_major_tensor" << std::endl;
+  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
+  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];
 
-void interleave_column_major_tensor_int4(int8_t*                    interleaved_quantized_tensor,
-                                    const int8_t*              quantized_tensor,
-                                    const std::vector<size_t>& shape)
-{
+  const size_t BITS_PER_ELT = 4;
+  const size_t elts_in_int32 = 32 / BITS_PER_ELT;
 
-    // We only want to run this step for weight only quant.
-    std::cout<<"### in interleave_column_major_tensor"<<std::endl;
-    const size_t num_rows    = shape.size() == 2 ? shape[0] : shape[1];
-    const size_t num_cols    = shape.size() == 2 ? shape[1] : shape[2];
+  const size_t rows_per_tile = 64;
+  std::cout << "running interleave_column_major_tensor" << std::endl;
+  std::cout << "num_rows:" << num_rows << ","
+            << "num_cols:" << num_cols << ","
+            << "BITS_PER_ELT:" << BITS_PER_ELT << ","
+            << "elts_in_int32:" << elts_in_int32 << ","
+            << "rows_per_tile:" << rows_per_tile << std::endl;
 
-    const size_t BITS_PER_ELT  = 4;
-    const size_t elts_in_int32 = 32 / BITS_PER_ELT;
+  const uint32_t* input_byte_ptr =
+      reinterpret_cast<const uint32_t*>(quantized_tensor);
+  uint32_t* output_byte_ptr =
+      reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
 
-    const size_t rows_per_tile = 64;
-    std::cout<<"running interleave_column_major_tensor"<<std::endl;
-    std::cout<<"num_rows:"<<num_rows<<","
-             <<"num_cols:"<<num_cols<<","
-             <<"BITS_PER_ELT:"<<BITS_PER_ELT<<","
-             <<"elts_in_int32:"<<elts_in_int32<<","
-             <<"rows_per_tile:"<<rows_per_tile<<std::endl;
+  const size_t num_vec_rows = num_rows / elts_in_int32;
+  const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
+  const size_t interleave = 4;
+  std::cout << "num_vec_rows:" << num_vec_rows << ","
+            << "vec_rows_per_tile:" << vec_rows_per_tile << ","
+            << "interleave:" << interleave << std::endl;
+  for (int read_col = 0; read_col < num_cols; ++read_col) {
+    const size_t write_col = read_col / interleave;
+    for (int base_vec_row = 0; base_vec_row < num_vec_rows;
+         base_vec_row += vec_rows_per_tile) {
+      for (int vec_read_row = base_vec_row;
+           vec_read_row <
+           std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
+           ++vec_read_row) {
+        const size_t vec_write_row =
+            interleave * base_vec_row +
+            vec_rows_per_tile * (read_col % interleave) +
+            vec_read_row % vec_rows_per_tile;
 
-    const uint32_t* input_byte_ptr  = reinterpret_cast<const uint32_t*>(quantized_tensor);
-    uint32_t*       output_byte_ptr = reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);
-
-
-    const size_t num_vec_rows      = num_rows / elts_in_int32;
-    const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
-    const size_t interleave        = 4;
-    std::cout<<"num_vec_rows:"<<num_vec_rows<<","
-             <<"vec_rows_per_tile:"<<vec_rows_per_tile<<","
-             <<"interleave:"<<interleave<<std::endl;
-    for (int read_col = 0; read_col < num_cols; ++read_col) {
-        const size_t write_col = read_col / interleave;
-        for (int base_vec_row = 0; base_vec_row < num_vec_rows; base_vec_row += vec_rows_per_tile) {
-            for (int vec_read_row = base_vec_row;
-                    vec_read_row < std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
-                    ++vec_read_row) {
-                const size_t vec_write_row = interleave * base_vec_row
-                                                + vec_rows_per_tile * (read_col % interleave)
-                                                + vec_read_row % vec_rows_per_tile;
-
-                const size_t read_offset = size_t(read_col) * num_vec_rows + vec_read_row;
-                const size_t write_offset = size_t(write_col) * num_vec_rows * interleave + vec_write_row;
-                // std::cout<<"read_offset:"<<read_offset
-                //          <<",write_offset:"<<write_offset<<std::endl;
-                output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
-            }
-        }
+        const size_t read_offset =
+            size_t(read_col) * num_vec_rows + vec_read_row;
+        const size_t write_offset =
+            size_t(write_col) * num_vec_rows * interleave + vec_write_row;
+        // std::cout<<"read_offset:"<<read_offset
+        //          <<",write_offset:"<<write_offset<<std::endl;
+        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
+      }
     }
+  }
 }
diff --git a/custom_ops/gpu_ops/dequant_int8.cu b/custom_ops/gpu_ops/dequant_int8.cu
index a9d84041a6..2b5f1b92d9 100644
--- a/custom_ops/gpu_ops/dequant_int8.cu
+++ b/custom_ops/gpu_ops/dequant_int8.cu
@@ -27,16 +27,16 @@ constexpr int DequantKernelVecSize = 4;
 
 template <typename data_t>
 inline HOSTDEVICE data_t roundWithTiesToEven(data_t x) {
-    data_t xLower = floor(x);
-    data_t xUpper = ceil(x);
-    // x is in interval [xl,xu]. Choose closest of two bounds, breaking ties to
-    // even.
-    data_t dLower = x - xLower;
-    data_t dUpper = xUpper - x;
-    return static_cast<data_t>(
-        (dLower == dUpper ? fmod(xLower, 2.0F) == 0.0F : dLower < dUpper)
-            ? xLower
-            : xUpper);
+  data_t xLower = floor(x);
+  data_t xUpper = ceil(x);
+  // x is in interval [xl,xu]. Choose closest of two bounds, breaking ties to
+  // even.
+  data_t dLower = x - xLower;
+  data_t dUpper = xUpper - x;
+  return static_cast<data_t>(
+      (dLower == dUpper ? fmod(xLower, 2.0F) == 0.0F : dLower < dUpper)
+          ? xLower
+          : xUpper);
 }
 
 template <typename data_t, int VecSize>
@@ -45,128 +45,125 @@ __global__ void DequantKernel(data_t *output,
                               const int m,  // batch size
                               const int n,  // hidden
                               const float *dequant_out_scale_data) {
-    int numel = m * n;
-    int stride = blockDim.x * gridDim.x * VecSize;
-    int idx = (blockIdx.x * blockDim.x + threadIdx.x) * VecSize;
-    int col_id = idx % n;
+  int numel = m * n;
+  int stride = blockDim.x * gridDim.x * VecSize;
+  int idx = (blockIdx.x * blockDim.x + threadIdx.x) * VecSize;
+  int col_id = idx % n;
 
-    AlignedVector<int32_t, VecSize> in_vec;
-    AlignedVector<float, VecSize> out_scale_vec;
-    AlignedVector<data_t, VecSize> out_vec;
+  AlignedVector<int32_t, VecSize> in_vec;
+  AlignedVector<float, VecSize> out_scale_vec;
+  AlignedVector<data_t, VecSize> out_vec;
 
-    for (; idx < numel; idx += stride) {
-        Load<int32_t, VecSize>(input + idx, &in_vec);
-        Load<float, VecSize>(dequant_out_scale_data + col_id, &out_scale_vec);
+  for (; idx < numel; idx += stride) {
+    Load<int32_t, VecSize>(input + idx, &in_vec);
+    Load<float, VecSize>(dequant_out_scale_data + col_id, &out_scale_vec);
 
 #pragma unroll
-        for (int i = 0; i < VecSize; ++i) {
-            out_vec[i] = static_cast<data_t>(static_cast<float>(in_vec[i]) *
-                                             out_scale_vec[i]);
-        }
-
-        Store<data_t, VecSize>(out_vec, output + idx);
+    for (int i = 0; i < VecSize; ++i) {
+      out_vec[i] =
+          static_cast<data_t>(static_cast<float>(in_vec[i]) * out_scale_vec[i]);
     }
+
+    Store<data_t, VecSize>(out_vec, output + idx);
+  }
 }
 
 template <paddle::DataType D>
 std::vector<paddle::Tensor> DispatchLaunchDequantInt8(
     const paddle::Tensor &input, const paddle::Tensor &scale) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    std::vector<int64_t> input_shape = input.shape();
+  std::vector<int64_t> input_shape = input.shape();
 
-    auto output = paddle::empty(input_shape, D, input.place());
-    int64_t m = input_shape[0];
-    int64_t n = input_shape[1];
+  auto output = paddle::empty(input_shape, D, input.place());
+  int64_t m = input_shape[0];
+  int64_t n = input_shape[1];
 
-    int64_t numel = m * n;
-    constexpr int64_t thread_per_block = 512;
-    int64_t block_per_grid =
-        (numel / DequantKernelVecSize + thread_per_block - 1) /
-        thread_per_block;
-    auto stream = input.stream();
+  int64_t numel = m * n;
+  constexpr int64_t thread_per_block = 512;
+  int64_t block_per_grid =
+      (numel / DequantKernelVecSize + thread_per_block - 1) / thread_per_block;
+  auto stream = input.stream();
 
-    DequantKernel<DataType_, DequantKernelVecSize>
-        <<<block_per_grid, thread_per_block, 0, stream>>>(
-            reinterpret_cast<DataType_ *>(output.data<data_t>()),
-            reinterpret_cast<const int32_t *>(input.data<int32_t>()),
-            m,
-            n,
-            reinterpret_cast<const float *>(scale.data<float>()));
+  DequantKernel<DataType_, DequantKernelVecSize>
+      <<<block_per_grid, thread_per_block, 0, stream>>>(
+          reinterpret_cast<DataType_ *>(output.data<data_t>()),
+          reinterpret_cast<const int32_t *>(input.data<int32_t>()),
+          m,
+          n,
+          reinterpret_cast<const float *>(scale.data<float>()));
 
-    return {output};
+  return {output};
 }
 
 std::vector<paddle::Tensor> LaunchDequantInt8(const paddle::Tensor &input,
                                               const paddle::Tensor &scale,
                                               std::string dtype) {
-    paddle::DataType data_type;
+  paddle::DataType data_type;
 
-    if (dtype == "float32")
-        data_type = paddle::DataType::FLOAT32;
-    else if (dtype == "bfloat16")
-        data_type = paddle::DataType::BFLOAT16;
-    else if (dtype == "float16")
-        data_type = paddle::DataType::FLOAT16;
-    else
-        PD_THROW(
-            "NOT supported data type. "
-            "Only bfloat16, float16 and float32 are supported. ");
+  if (dtype == "float32")
+    data_type = paddle::DataType::FLOAT32;
+  else if (dtype == "bfloat16")
+    data_type = paddle::DataType::BFLOAT16;
+  else if (dtype == "float16")
+    data_type = paddle::DataType::FLOAT16;
+  else
+    PD_THROW(
+        "NOT supported data type. "
+        "Only bfloat16, float16 and float32 are supported. ");
 
-    switch (data_type) {
-        case paddle::DataType::BFLOAT16:
-            return DispatchLaunchDequantInt8<paddle::DataType::BFLOAT16>(input,
-                                                                         scale);
-            break;
-        case paddle::DataType::FLOAT16:
-            return DispatchLaunchDequantInt8<paddle::DataType::FLOAT16>(input,
-                                                                        scale);
-            break;
-        case paddle::DataType::FLOAT32:
-            return DispatchLaunchDequantInt8<paddle::DataType::FLOAT32>(input,
-                                                                        scale);
-            break;
-        default:
-            break;
-    }
+  switch (data_type) {
+    case paddle::DataType::BFLOAT16:
+      return DispatchLaunchDequantInt8<paddle::DataType::BFLOAT16>(input,
+                                                                   scale);
+      break;
+    case paddle::DataType::FLOAT16:
+      return DispatchLaunchDequantInt8<paddle::DataType::FLOAT16>(input, scale);
+      break;
+    case paddle::DataType::FLOAT32:
+      return DispatchLaunchDequantInt8<paddle::DataType::FLOAT32>(input, scale);
+      break;
+    default:
+      break;
+  }
 }
 
 paddle::Tensor DequantInt8Func(const paddle::Tensor &input,
                                const paddle::Tensor &out_scale,
                                std::string dtype) {
-    return LaunchDequantInt8(input, out_scale, dtype)[0];
+  return LaunchDequantInt8(input, out_scale, dtype)[0];
 }
 
 std::vector<paddle::Tensor> DequantInt8(const paddle::Tensor &input,
                                         const paddle::Tensor &out_scale,
                                         std::string dtype) {
-    return LaunchDequantInt8(input, out_scale, dtype);
+  return LaunchDequantInt8(input, out_scale, dtype);
 }
 
 std::vector<std::vector<int64_t>> DequantInt8Shape(
     const std::vector<int64_t> &input_shape) {
-    return {input_shape};
+  return {input_shape};
 }
 
 std::vector<paddle::DataType> DequantInt8Dtype(
     const paddle::DataType &input_dtype,
     const paddle::DataType &out_scale_dtype,
     std::string dtype) {
-    paddle::DataType data_type;
-    if (dtype == "float32")
-        data_type = paddle::DataType::FLOAT32;
-    else if (dtype == "bfloat16")
-        data_type = paddle::DataType::BFLOAT16;
-    else if (dtype == "float16")
-        data_type = paddle::DataType::FLOAT16;
-    else
-        PD_THROW(
-            "NOT supported data type. "
-            "Only bfloat16, float16 and float32 are supported. ");
+  paddle::DataType data_type;
+  if (dtype == "float32")
+    data_type = paddle::DataType::FLOAT32;
+  else if (dtype == "bfloat16")
+    data_type = paddle::DataType::BFLOAT16;
+  else if (dtype == "float16")
+    data_type = paddle::DataType::FLOAT16;
+  else
+    PD_THROW(
+        "NOT supported data type. "
+        "Only bfloat16, float16 and float32 are supported. ");
 
-    return {data_type};
+  return {data_type};
 }
 
 PD_BUILD_STATIC_OP(dequant_int8)
diff --git a/custom_ops/gpu_ops/enforce_generation.cu b/custom_ops/gpu_ops/enforce_generation.cu
index 986101e824..7a1fd1f330 100644
--- a/custom_ops/gpu_ops/enforce_generation.cu
+++ b/custom_ops/gpu_ops/enforce_generation.cu
@@ -22,30 +22,30 @@ __global__ void apply_token_enforce_generation_scores_kernel(
     int logits_length,
     int logit_mask_length,
     int allowed_token_max_len) {
-    int bs = blockIdx.x;
-    int ti = threadIdx.x;
-    int32_t cur_allowed_token_num =
-        status_and_tokens[bs * allowed_token_max_len + 2];
+  int bs = blockIdx.x;
+  int ti = threadIdx.x;
+  int32_t cur_allowed_token_num =
+      status_and_tokens[bs * allowed_token_max_len + 2];
 #pragma unroll
-    if (cur_allowed_token_num > 0) {
-        for (int i = ti; i < logit_mask_length; i += blockDim.x) {
-            logit_mask[bs * logit_mask_length + i] = true;
-        }
-        __syncthreads();
+  if (cur_allowed_token_num > 0) {
+    for (int i = ti; i < logit_mask_length; i += blockDim.x) {
+      logit_mask[bs * logit_mask_length + i] = true;
+    }
+    __syncthreads();
 #pragma unroll
-        for (int i = ti; i < cur_allowed_token_num; i += blockDim.x) {
-            int idx = status_and_tokens[bs * allowed_token_max_len + i + 3];
-            logit_mask[bs * logit_mask_length + idx] = false;
-        }
-        __syncthreads();
+    for (int i = ti; i < cur_allowed_token_num; i += blockDim.x) {
+      int idx = status_and_tokens[bs * allowed_token_max_len + i + 3];
+      logit_mask[bs * logit_mask_length + idx] = false;
+    }
+    __syncthreads();
 
 #pragma unroll
-        for (int i = ti; i < logits_length; i += blockDim.x) {
-            if (logit_mask[bs * logits_length + i] == true) {
-                logits[bs * logits_length + i] = -1e10;
-            }
-        }
+    for (int i = ti; i < logits_length; i += blockDim.x) {
+      if (logit_mask[bs * logits_length + i] == true) {
+        logits[bs * logits_length + i] = -1e10;
+      }
     }
+  }
 }
 
 template <paddle::DataType D>
@@ -53,74 +53,71 @@ void token_enforce_generation_scores_kernel(
     const paddle::Tensor &logits,
     const paddle::Tensor &logit_mask,
     const paddle::Tensor &status_and_tokens) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
-    auto cu_stream = logits.stream();
-    std::vector<int64_t> logits_shape = logits.shape();
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+  auto cu_stream = logits.stream();
+  std::vector<int64_t> logits_shape = logits.shape();
 
-    std::vector<int64_t> allowed_token_shape = status_and_tokens.shape();
-    std::vector<int64_t> logit_mask_shape = logit_mask.shape();
-    int bs = logits_shape[0];
-    int logits_length = logits_shape[1];
-    int logit_mask_length = logit_mask_shape[1];
-    int allowed_token_max_len = allowed_token_shape[1];
-    int block_size = (logits_length + 32 - 1) / 32 * 32;
-    block_size = min(block_size, 512);
+  std::vector<int64_t> allowed_token_shape = status_and_tokens.shape();
+  std::vector<int64_t> logit_mask_shape = logit_mask.shape();
+  int bs = logits_shape[0];
+  int logits_length = logits_shape[1];
+  int logit_mask_length = logit_mask_shape[1];
+  int allowed_token_max_len = allowed_token_shape[1];
+  int block_size = (logits_length + 32 - 1) / 32 * 32;
+  block_size = min(block_size, 512);
 
-    // TODO(liuzichang): Reserved for multi-process
-    // int32_t con_gen_flag;
-    // printf("before for loop\n");
-    // for (;;) {
-    //   cudaMemcpy(reinterpret_cast<void*>(&con_gen_flag),
-    //   status_and_tokens.data<int32_t>(), sizeof(int32_t),
-    //   cudaMemcpyDeviceToHost); if (con_gen_flag == 1) {
-    //     break;
-    //   }
-    // }
-    // printf("finish for loop\n");
-    // printf("bs: %d, logits_length: %d, logit_mask_length: %d,
-    // allowed_token_max_len: %d, block_size: %d\n", bs, logits_length,
-    // logit_mask_length, allowed_token_max_len, block_size);
-    apply_token_enforce_generation_scores_kernel<<<bs,
-                                                   block_size,
-                                                   0,
-                                                   logits.stream()>>>(
-        reinterpret_cast<DataType_ *>(
-            const_cast<data_t *>(logits.data<data_t>())),
-        const_cast<bool *>(logit_mask.data<bool>()),
-        status_and_tokens.data<int32_t>(),
-        logits_length,
-        logit_mask_length,
-        allowed_token_max_len);
+  // TODO(liuzichang): Reserved for multi-process
+  // int32_t con_gen_flag;
+  // printf("before for loop\n");
+  // for (;;) {
+  //   cudaMemcpy(reinterpret_cast<void*>(&con_gen_flag),
+  //   status_and_tokens.data<int32_t>(), sizeof(int32_t),
+  //   cudaMemcpyDeviceToHost); if (con_gen_flag == 1) {
+  //     break;
+  //   }
+  // }
+  // printf("finish for loop\n");
+  // printf("bs: %d, logits_length: %d, logit_mask_length: %d,
+  // allowed_token_max_len: %d, block_size: %d\n", bs, logits_length,
+  // logit_mask_length, allowed_token_max_len, block_size);
+  apply_token_enforce_generation_scores_kernel<<<bs,
+                                                 block_size,
+                                                 0,
+                                                 logits.stream()>>>(
+      reinterpret_cast<DataType_ *>(
+          const_cast<data_t *>(logits.data<data_t>())),
+      const_cast<bool *>(logit_mask.data<bool>()),
+      status_and_tokens.data<int32_t>(),
+      logits_length,
+      logit_mask_length,
+      allowed_token_max_len);
 }
 
 void TokenEnforceGenerationScores(const paddle::Tensor &logits,
                                   const paddle::Tensor &logit_mask,
                                   const paddle::Tensor &status_and_tokens) {
-    switch (logits.type()) {
-        case paddle::DataType::BFLOAT16: {
-            return token_enforce_generation_scores_kernel<
-                paddle::DataType::BFLOAT16>(
-                logits, logit_mask, status_and_tokens);
-        }
-        case paddle::DataType::FLOAT16: {
-            return token_enforce_generation_scores_kernel<
-                paddle::DataType::FLOAT16>(
-                logits, logit_mask, status_and_tokens);
-        }
-        case paddle::DataType::FLOAT32: {
-            return token_enforce_generation_scores_kernel<
-                paddle::DataType::FLOAT32>(
-                logits, logit_mask, status_and_tokens);
-        }
-        default: {
-            PD_THROW(
-                "NOT supported data type. "
-                "Only float16, bfloat16 and float32 are supported. ");
-            break;
-        }
+  switch (logits.type()) {
+    case paddle::DataType::BFLOAT16: {
+      return token_enforce_generation_scores_kernel<paddle::DataType::BFLOAT16>(
+          logits, logit_mask, status_and_tokens);
     }
+    case paddle::DataType::FLOAT16: {
+      return token_enforce_generation_scores_kernel<paddle::DataType::FLOAT16>(
+          logits, logit_mask, status_and_tokens);
+    }
+    case paddle::DataType::FLOAT32: {
+      return token_enforce_generation_scores_kernel<paddle::DataType::FLOAT32>(
+          logits, logit_mask, status_and_tokens);
+    }
+    default: {
+      PD_THROW(
+          "NOT supported data type. "
+          "Only float16, bfloat16 and float32 are supported. ");
+      break;
+    }
+  }
 }
 
 void __global__ update_enf_gen_values_kernel(int32_t *status_and_tokens,
@@ -129,40 +126,39 @@ void __global__ update_enf_gen_values_kernel(int32_t *status_and_tokens,
                                              int32_t now_bs,
                                              int32_t status_and_tokens_max_len,
                                              int32_t next_tokens_len) {
-    int tid = threadIdx.x;
-    for (int32_t bs_idx = tid; bs_idx < now_bs; bs_idx += blockDim.x) {
-        bool stop_flag = stop_flags[bs_idx];
-        int32_t *cur_status_and_tokens =
-            status_and_tokens + bs_idx * status_and_tokens_max_len;
-        bool is_first = (cur_status_and_tokens[1] == 2);
-        if (!stop_flag) {
-            cur_status_and_tokens[1] = 1;
-            cur_status_and_tokens[2] =
-                static_cast<int32_t>(next_tokens[bs_idx]);
-        } else if (!is_first) {  // stop_flag && not first
-            cur_status_and_tokens[1] = 0;
-        }
-    }
-    __syncthreads();
-    if (tid == 0) {
-        status_and_tokens[0] = 2;
+  int tid = threadIdx.x;
+  for (int32_t bs_idx = tid; bs_idx < now_bs; bs_idx += blockDim.x) {
+    bool stop_flag = stop_flags[bs_idx];
+    int32_t *cur_status_and_tokens =
+        status_and_tokens + bs_idx * status_and_tokens_max_len;
+    bool is_first = (cur_status_and_tokens[1] == 2);
+    if (!stop_flag) {
+      cur_status_and_tokens[1] = 1;
+      cur_status_and_tokens[2] = static_cast<int32_t>(next_tokens[bs_idx]);
+    } else if (!is_first) {  // stop_flag && not first
+      cur_status_and_tokens[1] = 0;
     }
+  }
+  __syncthreads();
+  if (tid == 0) {
+    status_and_tokens[0] = 2;
+  }
 }
 
 void UpdateEnfGenValues(const paddle::Tensor &status_and_tokens,
                         const paddle::Tensor &stop_flags,
                         const paddle::Tensor &next_tokens) {
-    const int bsz = next_tokens.shape()[0];
-    const int status_and_tokens_max_len = status_and_tokens.shape()[1];
-    const int next_tokens_len = next_tokens.shape()[0];
+  const int bsz = next_tokens.shape()[0];
+  const int status_and_tokens_max_len = status_and_tokens.shape()[1];
+  const int next_tokens_len = next_tokens.shape()[0];
 
-    update_enf_gen_values_kernel<<<1, 1024, 0, next_tokens.stream()>>>(
-        const_cast<int32_t *>(status_and_tokens.data<int32_t>()),
-        stop_flags.data<bool>(),
-        next_tokens.data<int64_t>(),
-        bsz,
-        status_and_tokens_max_len,
-        next_tokens_len);
+  update_enf_gen_values_kernel<<<1, 1024, 0, next_tokens.stream()>>>(
+      const_cast<int32_t *>(status_and_tokens.data<int32_t>()),
+      stop_flags.data<bool>(),
+      next_tokens.data<int64_t>(),
+      bsz,
+      status_and_tokens_max_len,
+      next_tokens_len);
 }
 
 PD_BUILD_STATIC_OP(get_enf_gen_scores)
diff --git a/custom_ops/gpu_ops/env.h b/custom_ops/gpu_ops/env.h
index 5e8eee339e..eed65bea2c 100644
--- a/custom_ops/gpu_ops/env.h
+++ b/custom_ops/gpu_ops/env.h
@@ -15,59 +15,83 @@
 #pragma once
 
 inline uint32_t get_decoder_block_shape_q() {
-    static const char* decoder_block_shape_q_env = std::getenv("FLAGS_dec_block_shape_q");
-    static const uint32_t decoder_block_shape_q =
-            decoder_block_shape_q_env == nullptr ? 16 : std::stoi(std::string(decoder_block_shape_q_env));
-    return decoder_block_shape_q;
+  static const char* decoder_block_shape_q_env =
+      std::getenv("FLAGS_dec_block_shape_q");
+  static const uint32_t decoder_block_shape_q =
+      decoder_block_shape_q_env == nullptr
+          ? 16
+          : std::stoi(std::string(decoder_block_shape_q_env));
+  return decoder_block_shape_q;
 }
 
 inline uint32_t get_encoder_block_shape_q() {
-    static const char* encoder_block_shape_q_env = std::getenv("FLAGS_enc_block_shape_q");
-    static const uint32_t encoder_block_shape_q =
-            encoder_block_shape_q_env == nullptr ? 64 : std::stoi(std::string(encoder_block_shape_q_env));
-    return encoder_block_shape_q;
+  static const char* encoder_block_shape_q_env =
+      std::getenv("FLAGS_enc_block_shape_q");
+  static const uint32_t encoder_block_shape_q =
+      encoder_block_shape_q_env == nullptr
+          ? 64
+          : std::stoi(std::string(encoder_block_shape_q_env));
+  return encoder_block_shape_q;
 }
 
 inline uint32_t get_max_partition_size(int bsz) {
-    static const char* max_partition_size_env = std::getenv("FLAGS_cascade_attention_max_partition_size");
-    static const uint32_t max_partition_size =
-            max_partition_size_env == nullptr ? 32768 : std::stoul(std::string(max_partition_size_env));
-    return max_partition_size;
+  static const char* max_partition_size_env =
+      std::getenv("FLAGS_cascade_attention_max_partition_size");
+  static const uint32_t max_partition_size =
+      max_partition_size_env == nullptr
+          ? 32768
+          : std::stoul(std::string(max_partition_size_env));
+  return max_partition_size;
 }
 
 inline uint32_t get_cascade_attention_deal_each_time() {
-    static const char* cascade_attention_deal_each_time_env = std::getenv("FLAGS_cascade_attention_deal_each_time");
-    static const uint32_t cascade_attention_deal_each_time =
-            cascade_attention_deal_each_time_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_deal_each_time_env));
-    return (cascade_attention_deal_each_time != 0 ? cascade_attention_deal_each_time : 32);
+  static const char* cascade_attention_deal_each_time_env =
+      std::getenv("FLAGS_cascade_attention_deal_each_time");
+  static const uint32_t cascade_attention_deal_each_time =
+      cascade_attention_deal_each_time_env == nullptr
+          ? 0
+          : std::stoul(std::string(cascade_attention_deal_each_time_env));
+  return (cascade_attention_deal_each_time != 0
+              ? cascade_attention_deal_each_time
+              : 32);
 }
 
 inline uint32_t get_cascade_attention_num_stages() {
-    static const char* cascade_attention_num_stages_env = std::getenv("FLAGS_cascade_attention_num_stages");
-    static const uint32_t cascade_attention_num_stages =
-            cascade_attention_num_stages_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_num_stages_env));
-    return cascade_attention_num_stages != 0 ? cascade_attention_num_stages : 2;
+  static const char* cascade_attention_num_stages_env =
+      std::getenv("FLAGS_cascade_attention_num_stages");
+  static const uint32_t cascade_attention_num_stages =
+      cascade_attention_num_stages_env == nullptr
+          ? 0
+          : std::stoul(std::string(cascade_attention_num_stages_env));
+  return cascade_attention_num_stages != 0 ? cascade_attention_num_stages : 2;
 }
 
 inline uint32_t get_cascade_attention_num_threads() {
-    static const char* cascade_attention_num_threads_env = std::getenv("FLAGS_cascade_attention_num_threads");
-    static const uint32_t cascade_attention_num_threads =
-            cascade_attention_num_threads_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_num_threads_env));
-    return cascade_attention_num_threads != 0 ? cascade_attention_num_threads : 128;
+  static const char* cascade_attention_num_threads_env =
+      std::getenv("FLAGS_cascade_attention_num_threads");
+  static const uint32_t cascade_attention_num_threads =
+      cascade_attention_num_threads_env == nullptr
+          ? 0
+          : std::stoul(std::string(cascade_attention_num_threads_env));
+  return cascade_attention_num_threads != 0 ? cascade_attention_num_threads
+                                            : 128;
 }
 
 inline bool get_mla_use_tensorcore() {
-    static const char* mla_use_tensorcore_env = std::getenv("FLAGS_mla_use_tensorcore");
-    static const uint32_t mla_use_tensorcore =
-            mla_use_tensorcore_env == nullptr ? 0 : std::stoul(std::string(mla_use_tensorcore_env));
-    return mla_use_tensorcore != 0 ? true : false;
+  static const char* mla_use_tensorcore_env =
+      std::getenv("FLAGS_mla_use_tensorcore");
+  static const uint32_t mla_use_tensorcore =
+      mla_use_tensorcore_env == nullptr
+          ? 0
+          : std::stoul(std::string(mla_use_tensorcore_env));
+  return mla_use_tensorcore != 0 ? true : false;
 }
 inline int get_mla_dec_chunk_size(int bsz) {
-    static const char* mla_dec_chunk_size_env =
-        std::getenv("FLAGS_mla_dec_chunk_size");
-    static const int mla_dec_chunk_size =
-        mla_dec_chunk_size_env == nullptr
-            ? -1
-            : std::stoi(std::string(mla_dec_chunk_size_env));
-    return bsz > 1 ? mla_dec_chunk_size : 64;
+  static const char* mla_dec_chunk_size_env =
+      std::getenv("FLAGS_mla_dec_chunk_size");
+  static const int mla_dec_chunk_size =
+      mla_dec_chunk_size_env == nullptr
+          ? -1
+          : std::stoi(std::string(mla_dec_chunk_size_env));
+  return bsz > 1 ? mla_dec_chunk_size : 64;
 }
diff --git a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_common.h b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_common.h
index 6ec4103e11..e6277bf436 100644
--- a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_common.h
+++ b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_common.h
@@ -26,7 +26,6 @@
 #include "paddle/phi/core/allocator.h"
 #include "paddle/common/flags.h"
 
-
 typedef struct {
   const void *A;
   const void *B;
@@ -47,8 +46,8 @@ typedef struct {
   std::vector<int64_t> &bias_dims;
   std::string &fuse_gemm_config;
   int split_k = 1;
-  const float* scalar_scale = nullptr;
-  const float* channel_scale = nullptr;
+  const float *scalar_scale = nullptr;
+  const float *channel_scale = nullptr;
   const void *A_scale = nullptr;
   const void *B_scale = nullptr;
 } GemmEpilogueAllParams;
diff --git a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu
index 22cdaa4c8a..7434670310 100644
--- a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu
+++ b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu
@@ -29,143 +29,132 @@ std::vector<paddle::Tensor> cutlass_fp8_fp8_fp8_dual_gemm(
     float scale1,     // only support per-tensor quantization
     float scale_out,  // only support per-tensor quantization
     std::string activation_type) {
-    paddle::Tensor out;
-    void* out_ptr = nullptr;
-    const void* x_ptr = nullptr;
-    const void* y0_ptr = nullptr;
-    const void* y1_ptr = nullptr;
+  paddle::Tensor out;
+  void* out_ptr = nullptr;
+  const void* x_ptr = nullptr;
+  const void* y0_ptr = nullptr;
+  const void* y1_ptr = nullptr;
 
-    auto place = x.place();
-    cudaStream_t stream = x.stream();
-    int64_t device_id = place.GetDeviceId();
-    int sm_version = GetGPUComputeCapability(device_id);
+  auto place = x.place();
+  cudaStream_t stream = x.stream();
+  int64_t device_id = place.GetDeviceId();
+  int sm_version = GetGPUComputeCapability(device_id);
 
-    int rank = x.dims().size();
-    int M = 0;
-    int K = 0;
-    int N = 0;
-    int ldd = 0;
+  int rank = x.dims().size();
+  int M = 0;
+  int K = 0;
+  int N = 0;
+  int ldd = 0;
 
-    int lda = x.dims()[rank - 1];
-    int ldb = y0.dims()[rank - 1];
-    if (!trans_x) {
-        M = x.dims()[rank - 2];
-        K = x.dims()[rank - 1];
+  int lda = x.dims()[rank - 1];
+  int ldb = y0.dims()[rank - 1];
+  if (!trans_x) {
+    M = x.dims()[rank - 2];
+    K = x.dims()[rank - 1];
 
+  } else {
+    M = x.dims()[rank - 1];
+    K = x.dims()[rank - 2];
+  }
+
+  if (!trans_y) {
+    N = y0.dims()[rank - 1];
+    ldd = y0.dims()[rank - 1];
+  } else {
+    N = y0.dims()[rank - 2];
+    ldd = y0.dims()[rank - 2];
+  }
+
+  int batch_count = 1;
+  for (size_t i = 0; i < rank - 2; ++i) {
+    batch_count *= x.dims()[i];
+  }
+
+  std::string input_dtype = "";
+  std::string output_dtype = "";
+  std::vector<int64_t> out_shape = x.shape();
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
+
+  if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
+    input_dtype = "float8_e4m3fn";
+    output_dtype = "float8_e4m3fn";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
+    y0_ptr =
+        reinterpret_cast<const void*>(y0.data<phi::dtype::float8_e4m3fn>());
+    y1_ptr =
+        reinterpret_cast<const void*>(y1.data<phi::dtype::float8_e4m3fn>());
+    out = paddle::empty(out_shape, paddle::DataType::FLOAT8_E4M3FN, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float8_e4m3fn>());
+  } else if (x.dtype() == phi::DataType::FLOAT8_E5M2) {
+    input_dtype = "float8_e5m2";
+    output_dtype = "float8_e5m2";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e5m2>());
+    y0_ptr = reinterpret_cast<const void*>(y0.data<phi::dtype::float8_e5m2>());
+    y1_ptr = reinterpret_cast<const void*>(y1.data<phi::dtype::float8_e5m2>());
+    out = paddle::empty(out_shape, paddle::DataType::FLOAT8_E5M2, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float8_e5m2>());
+  } else {
+    PADDLE_THROW(phi::errors::Fatal(
+        "fp8_fp8_fp8_dual_gemm_fused only support e4m3 and e5m2 input"));
+  }
+
+  std::string isbias = "false";
+  std::string bias_dtype = "float16";
+  void* bias_data0 = nullptr;
+  void* bias_data1 = nullptr;
+  std::vector<int64_t> bias_dims0{};
+  std::vector<int64_t> bias_dims1{};
+  if (bias0 && bias1) {
+    isbias = "true";
+    bias_dims0 = common::vectorize(bias0.get().dims());
+    bias_dims1 = common::vectorize(bias1.get().dims());
+    if (bias0.get().dtype() == phi::DataType::FLOAT16) {
+      bias_dtype = "float16";
+      bias_data0 = reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
+          bias0.get().data<phi::dtype::float16>()));
+      bias_data1 = reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
+          bias1.get().data<phi::dtype::float16>()));
     } else {
-        M = x.dims()[rank - 1];
-        K = x.dims()[rank - 2];
+      bias_dtype = "bfloat16";
+      bias_data0 = reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
+          bias0.get().data<phi::dtype::bfloat16>()));
+      bias_data1 = reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
+          bias1.get().data<phi::dtype::bfloat16>()));
     }
+  }
 
-    if (!trans_y) {
-        N = y0.dims()[rank - 1];
-        ldd = y0.dims()[rank - 1];
-    } else {
-        N = y0.dims()[rank - 2];
-        ldd = y0.dims()[rank - 2];
-    }
+  paddle::Tensor out0;
+  paddle::Tensor out1;
+  void* out0_ptr = nullptr;
+  void* out1_ptr = nullptr;
+  if (bias_dtype == "float16") {
+    out0 = paddle::empty(out_shape, phi::DataType::FLOAT16, x.place());
+    out0_ptr = reinterpret_cast<void*>(out0.data<phi::dtype::float16>());
+    out1 = paddle::empty(out_shape, phi::DataType::FLOAT16, x.place());
+    out1_ptr = reinterpret_cast<void*>(out1.data<phi::dtype::float16>());
+  } else {
+    out0 = paddle::empty(out_shape, phi::DataType::BFLOAT16, x.place());
+    out0_ptr = reinterpret_cast<void*>(out0.data<phi::dtype::bfloat16>());
+    out1 = paddle::empty(out_shape, phi::DataType::BFLOAT16, x.place());
+    out1_ptr = reinterpret_cast<void*>(out1.data<phi::dtype::bfloat16>());
+  }
 
-    int batch_count = 1;
-    for (size_t i = 0; i < rank - 2; ++i) {
-        batch_count *= x.dims()[i];
-    }
+  std::string act = (activation_type == "") ? "swiglu" : activation_type;
 
-    std::string input_dtype = "";
-    std::string output_dtype = "";
-    std::vector<int64_t> out_shape = x.shape();
-    out_shape[rank - 1] = N;
-    out_shape[rank - 2] = M;
+  std::string fuse_gemm_config = input_dtype + "_" + output_dtype + "_" +
+                                 bias_dtype + "_" + isbias + "_" + act;
 
-    if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
-        input_dtype = "float8_e4m3fn";
-        output_dtype = "float8_e4m3fn";
-        x_ptr =
-            reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
-        y0_ptr =
-            reinterpret_cast<const void*>(y0.data<phi::dtype::float8_e4m3fn>());
-        y1_ptr =
-            reinterpret_cast<const void*>(y1.data<phi::dtype::float8_e4m3fn>());
-        out = paddle::empty(
-            out_shape, paddle::DataType::FLOAT8_E4M3FN, x.place());
-        out_ptr =
-            reinterpret_cast<void*>(out.data<phi::dtype::float8_e4m3fn>());
-    } else if (x.dtype() == phi::DataType::FLOAT8_E5M2) {
-        input_dtype = "float8_e5m2";
-        output_dtype = "float8_e5m2";
-        x_ptr =
-            reinterpret_cast<const void*>(x.data<phi::dtype::float8_e5m2>());
-        y0_ptr =
-            reinterpret_cast<const void*>(y0.data<phi::dtype::float8_e5m2>());
-        y1_ptr =
-            reinterpret_cast<const void*>(y1.data<phi::dtype::float8_e5m2>());
-        out =
-            paddle::empty(out_shape, paddle::DataType::FLOAT8_E5M2, x.place());
-        out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float8_e5m2>());
-    } else {
-        PADDLE_THROW(phi::errors::Fatal(
-            "fp8_fp8_fp8_dual_gemm_fused only support e4m3 and e5m2 input"));
-    }
+  DualGemmEpilogueAllParams params = {
+      x_ptr,      y0_ptr,     y1_ptr,     out0_ptr,
+      out1_ptr,   out_ptr,    scale0,     scale1,
+      scale_out,  M,          N,          K,
+      lda,        ldb,        ldd,        batch_count,
+      place,      stream,     sm_version, bias_data0,
+      bias_data1, bias_dims0, bias_dims1, fuse_gemm_config};
 
-    std::string isbias = "false";
-    std::string bias_dtype = "float16";
-    void* bias_data0 = nullptr;
-    void* bias_data1 = nullptr;
-    std::vector<int64_t> bias_dims0{};
-    std::vector<int64_t> bias_dims1{};
-    if (bias0 && bias1) {
-        isbias = "true";
-        bias_dims0 = common::vectorize(bias0.get().dims());
-        bias_dims1 = common::vectorize(bias1.get().dims());
-        if (bias0.get().dtype() == phi::DataType::FLOAT16) {
-            bias_dtype = "float16";
-            bias_data0 =
-                reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
-                    bias0.get().data<phi::dtype::float16>()));
-            bias_data1 =
-                reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
-                    bias1.get().data<phi::dtype::float16>()));
-        } else {
-            bias_dtype = "bfloat16";
-            bias_data0 =
-                reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
-                    bias0.get().data<phi::dtype::bfloat16>()));
-            bias_data1 =
-                reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
-                    bias1.get().data<phi::dtype::bfloat16>()));
-        }
-    }
-
-    paddle::Tensor out0;
-    paddle::Tensor out1;
-    void* out0_ptr = nullptr;
-    void* out1_ptr = nullptr;
-    if (bias_dtype == "float16") {
-        out0 = paddle::empty(out_shape, phi::DataType::FLOAT16, x.place());
-        out0_ptr = reinterpret_cast<void*>(out0.data<phi::dtype::float16>());
-        out1 = paddle::empty(out_shape, phi::DataType::FLOAT16, x.place());
-        out1_ptr = reinterpret_cast<void*>(out1.data<phi::dtype::float16>());
-    } else {
-        out0 = paddle::empty(out_shape, phi::DataType::BFLOAT16, x.place());
-        out0_ptr = reinterpret_cast<void*>(out0.data<phi::dtype::bfloat16>());
-        out1 = paddle::empty(out_shape, phi::DataType::BFLOAT16, x.place());
-        out1_ptr = reinterpret_cast<void*>(out1.data<phi::dtype::bfloat16>());
-    }
-
-    std::string act = (activation_type == "") ? "swiglu" : activation_type;
-
-    std::string fuse_gemm_config = input_dtype + "_" + output_dtype + "_" +
-                                   bias_dtype + "_" + isbias + "_" + act;
-
-    DualGemmEpilogueAllParams params = {
-        x_ptr,      y0_ptr,     y1_ptr,     out0_ptr,
-        out1_ptr,   out_ptr,    scale0,     scale1,
-        scale_out,  M,          N,          K,
-        lda,        ldb,        ldd,        batch_count,
-        place,      stream,     sm_version, bias_data0,
-        bias_data1, bias_dims0, bias_dims1, fuse_gemm_config};
-
-    fp8_fp8_dual_gemm_scale_bias_act(params);
-    return {out};
+  fp8_fp8_dual_gemm_scale_bias_act(params);
+  return {out};
 }
 
 std::vector<std::vector<int64_t>> CutlassFp8Fp8Fp8DualGemmFusedInferShape(
@@ -176,41 +165,41 @@ std::vector<std::vector<int64_t>> CutlassFp8Fp8Fp8DualGemmFusedInferShape(
     const paddle::optional<std::vector<int64_t>>& bias1_shape,
     bool trans_x,
     bool trans_y) {
-    if (x_shape.size() != y0_shape.size()) {
-        PD_THROW(
-            "The rank of input X and Y0 should be equal, but received X's rank "
-            "is %d, Y0's rank is %d",
-            x_shape.size(),
-            y0_shape.size());
-    }
+  if (x_shape.size() != y0_shape.size()) {
+    PD_THROW(
+        "The rank of input X and Y0 should be equal, but received X's rank "
+        "is %d, Y0's rank is %d",
+        x_shape.size(),
+        y0_shape.size());
+  }
 
-    if (y0_shape.size() != y1_shape.size()) {
-        PD_THROW(
-            "The rank of input Y0 and Y1 should be equal, but received Y0's "
-            "rank is %d, Y1's rank is %d.",
-            y0_shape.size(),
-            y1_shape.size());
-    }
+  if (y0_shape.size() != y1_shape.size()) {
+    PD_THROW(
+        "The rank of input Y0 and Y1 should be equal, but received Y0's "
+        "rank is %d, Y1's rank is %d.",
+        y0_shape.size(),
+        y1_shape.size());
+  }
 
-    int rank = x_shape.size();
-    int M = 0;
-    int N = 0;
+  int rank = x_shape.size();
+  int M = 0;
+  int N = 0;
 
-    if (!trans_x) {
-        M = x_shape[rank - 2];
+  if (!trans_x) {
+    M = x_shape[rank - 2];
 
-    } else {
-        M = x_shape[rank - 1];
-    }
-    if (!trans_y) {
-        N = y0_shape[rank - 1];
-    } else {
-        N = y0_shape[rank - 2];
-    }
-    std::vector<int64_t> out_shape = x_shape;
-    out_shape[rank - 1] = N;
-    out_shape[rank - 2] = M;
-    return {out_shape};
+  } else {
+    M = x_shape[rank - 1];
+  }
+  if (!trans_y) {
+    N = y0_shape[rank - 1];
+  } else {
+    N = y0_shape[rank - 2];
+  }
+  std::vector<int64_t> out_shape = x_shape;
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
+  return {out_shape};
 }
 
 std::vector<paddle::DataType> CutlassFp8Fp8Fp8DualGemmFusedInferDtype(
@@ -219,33 +208,33 @@ std::vector<paddle::DataType> CutlassFp8Fp8Fp8DualGemmFusedInferDtype(
     const paddle::DataType& y1_type,
     const paddle::optional<paddle::DataType>& bias0_type,
     const paddle::optional<paddle::DataType>& bias1_type) {
-    if (x_type != y0_type) {
-        PD_THROW(
-            "The type of input X and Y0 should be equal, but received X's type "
-            "is %s, Y0's type is %s.",
-            x_type,
-            y0_type);
-    }
+  if (x_type != y0_type) {
+    PD_THROW(
+        "The type of input X and Y0 should be equal, but received X's type "
+        "is %s, Y0's type is %s.",
+        x_type,
+        y0_type);
+  }
 
-    if (y0_type != y1_type) {
-        PD_THROW(
-            "The type of input Y0 and Y1 should be equal, but received Y0's "
-            "type is %s, Y1's type is %s.",
-            y0_type,
-            y1_type);
-    }
+  if (y0_type != y1_type) {
+    PD_THROW(
+        "The type of input Y0 and Y1 should be equal, but received Y0's "
+        "type is %s, Y1's type is %s.",
+        y0_type,
+        y1_type);
+  }
 
-    if (bias0_type != bias1_type) {
-        PD_THROW(
-            "The type of bias0 and bias1 should be equal, but received bias0's "
-            "type is %s, bias1's type is %s.",
-            bias0_type,
-            bias1_type);
-    }
+  if (bias0_type != bias1_type) {
+    PD_THROW(
+        "The type of bias0 and bias1 should be equal, but received bias0's "
+        "type is %s, bias1's type is %s.",
+        bias0_type,
+        bias1_type);
+  }
 
-    paddle::DataType data_type;
-    data_type = paddle::DataType::FLOAT8_E4M3FN;
-    return {data_type};
+  paddle::DataType data_type;
+  data_type = paddle::DataType::FLOAT8_E4M3FN;
+  return {data_type};
 }
 
 PD_BUILD_STATIC_OP(cutlass_fp8_fp8_fp8_dual_gemm_fused)
diff --git a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu
index d4bb375dc4..f713cda812 100644
--- a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu
+++ b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu
@@ -14,15 +14,20 @@
 
 #include <iostream>
 
-#include "fp8_common.h" // NOLINT
+#include "fp8_common.h"  // NOLINT
 #include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
 #include "fp8_gemm_fused/fuse_block_gemm_act_template_3x.h"
 
 std::vector<paddle::Tensor> cutlass_fp8_fp8_half_block_gemm_fused(
-    const paddle::Tensor &x, const paddle::Tensor &y,
-    const paddle::Tensor &x_scale, const paddle::Tensor &y_scale,
-    const paddle::optional<paddle::Tensor> &bias, bool trans_x, bool trans_y,
-    std::string output_dtype, std::string activation_type) {
+    const paddle::Tensor &x,
+    const paddle::Tensor &y,
+    const paddle::Tensor &x_scale,
+    const paddle::Tensor &y_scale,
+    const paddle::optional<paddle::Tensor> &bias,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype,
+    std::string activation_type) {
   paddle::Tensor out;
   void *out_ptr = nullptr;
   const void *x_ptr = nullptr;
@@ -122,7 +127,7 @@ std::vector<paddle::Tensor> cutlass_fp8_fp8_half_block_gemm_fused(
                                   K,           lda,         ldb,
                                   ldd,         batch_count, place,
                                   stream,      sm_version,
-                                  0.01, // for leaky_relu
+                                  0.01,  // for leaky_relu
                                   bias_data,   bias_dims,   fuse_gemm_config,
                                   0,           nullptr,     nullptr,
                                   x_scale_ptr, y_scale_ptr};
@@ -133,26 +138,34 @@ std::vector<paddle::Tensor> cutlass_fp8_fp8_half_block_gemm_fused(
 }
 
 std::vector<std::vector<int64_t>> CutlassFp8Fp8HalfBlockGemmFusedInferShape(
-    const std::vector<int64_t> &x_shape, const std::vector<int64_t> &y_shape,
+    const std::vector<int64_t> &x_shape,
+    const std::vector<int64_t> &y_shape,
     const std::vector<int64_t> &x_scale_shape,
     const std::vector<int64_t> &y_scale_shape,
-    const paddle::optional<std::vector<int64_t>> &bias_shape, bool trans_x,
+    const paddle::optional<std::vector<int64_t>> &bias_shape,
+    bool trans_x,
     bool trans_y) {
-  PADDLE_ENFORCE_EQ(x_shape.size(), y_shape.size(),
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    y_shape.size(),
                     phi::errors::InvalidArgument(
                         "The rank of input X and Y should be equal, but "
                         "received X's rank is %d, Y's rank is %d.",
-                        x_shape.size(), y_shape.size()));
-  PADDLE_ENFORCE_EQ(x_shape.size(), x_scale_shape.size(),
+                        x_shape.size(),
+                        y_shape.size()));
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    x_scale_shape.size(),
                     phi::errors::InvalidArgument(
                         "The rank of input X and X_scale should be equal, but "
                         "received X's rank is %d, X_scale's rank is %d.",
-                        x_shape.size(), x_scale_shape.size()));
-  PADDLE_ENFORCE_EQ(y_shape.size(), y_scale_shape.size(),
+                        x_shape.size(),
+                        x_scale_shape.size()));
+  PADDLE_ENFORCE_EQ(y_shape.size(),
+                    y_scale_shape.size(),
                     phi::errors::InvalidArgument(
                         "The rank of input Y and Y_scale should be equal, but "
                         "received Y's rank is %d, Y_scale's rank is %d.",
-                        y_shape.size(), y_scale_shape.size()));
+                        y_shape.size(),
+                        y_scale_shape.size()));
   int rank = x_shape.size();
   int M = 0;
   int N = 0;
@@ -186,25 +199,32 @@ std::vector<std::vector<int64_t>> CutlassFp8Fp8HalfBlockGemmFusedInferShape(
 }
 
 std::vector<paddle::DataType> CutlassFp8Fp8HalfBlockGemmFusedInferDtype(
-    const paddle::DataType &x_type, const paddle::DataType &y_type,
-    const paddle::DataType &x_scale_type, const paddle::DataType &y_scale_type,
-    const paddle::optional<paddle::DataType> &bias_type, bool trans_x,
-    bool trans_y, std::string output_dtype) {
+    const paddle::DataType &x_type,
+    const paddle::DataType &y_type,
+    const paddle::DataType &x_scale_type,
+    const paddle::DataType &y_scale_type,
+    const paddle::optional<paddle::DataType> &bias_type,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype) {
   paddle::DataType data_type;
   if (output_dtype == "bfloat16")
     data_type = paddle::DataType::BFLOAT16;
   else if (output_dtype == "float16")
     data_type = paddle::DataType::FLOAT16;
   else
-    PD_THROW("cutlass_fp8_fp8_half_block_gemm only support bfloat16 and "
-             "float16 output");
+    PD_THROW(
+        "cutlass_fp8_fp8_half_block_gemm only support bfloat16 and "
+        "float16 output");
   return {data_type};
 }
 
 PD_BUILD_STATIC_OP(cutlass_fp8_fp8_half_block_gemm_fused)
     .Inputs({"x", "y", "x_sacle", "y_scale", paddle::Optional("bias")})
-    .Attrs({"transpose_x: bool", "transpose_y: bool",
-            "output_dtype: std::string", "act: std::string"})
+    .Attrs({"transpose_x: bool",
+            "transpose_y: bool",
+            "output_dtype: std::string",
+            "act: std::string"})
     .Outputs({"out"})
     .SetKernelFn(PD_KERNEL(cutlass_fp8_fp8_half_block_gemm_fused))
     .SetInferShapeFn(PD_INFER_SHAPE(CutlassFp8Fp8HalfBlockGemmFusedInferShape))
diff --git a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu
index 3e1ce299a3..29cf89d7e8 100644
--- a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu
+++ b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu
@@ -116,9 +116,10 @@ __global__ void cudaCoreGemm(InputType const* __restrict__ act,
     }
 
     if constexpr (UseBias) {
-        output[mid * n + nid] = static_cast<OutputType>(val * alpha + (float)*(bias+nid)) ;
+      output[mid * n + nid] =
+          static_cast<OutputType>(val * alpha + (float)*(bias + nid));
     } else {
-        output[mid * n + nid] = static_cast<OutputType>(val * alpha);
+      output[mid * n + nid] = static_cast<OutputType>(val * alpha);
     }
   }
 }
@@ -129,12 +130,13 @@ template <typename InputType,
           int32_t TILE_N,
           int32_t BLOCK_SIZE>
 void cudaCoreGemmKernel(GemmParams const& params) {
-    dim3 block(BLOCK_SIZE);
-    dim3 grid(params.m / TILE_M, params.n / TILE_N);
-    // std::cout << "m" << params.m << " n" << params.n <<  " k " << params.k << std::endl;
+  dim3 block(BLOCK_SIZE);
+  dim3 grid(params.m / TILE_M, params.n / TILE_N);
+  // std::cout << "m" << params.m << " n" << params.n <<  " k " << params.k <<
+  // std::endl;
 
-    if (params.bias != nullptr) {
-        cudaCoreGemm<InputType, OutputType, TILE_M, TILE_N, BLOCK_SIZE, true>
+  if (params.bias != nullptr) {
+    cudaCoreGemm<InputType, OutputType, TILE_M, TILE_N, BLOCK_SIZE, true>
         <<<grid, block, 0, params.stream>>>(
             reinterpret_cast<InputType const*>(params.act),
             reinterpret_cast<InputType const*>(params.weight),
@@ -144,8 +146,8 @@ void cudaCoreGemmKernel(GemmParams const& params) {
             params.n,
             params.k,
             params.alpha);
-    } else {
-        cudaCoreGemm<InputType, OutputType, TILE_M, TILE_N, BLOCK_SIZE, false>
+  } else {
+    cudaCoreGemm<InputType, OutputType, TILE_M, TILE_N, BLOCK_SIZE, false>
         <<<grid, block, 0, params.stream>>>(
             reinterpret_cast<InputType const*>(params.act),
             reinterpret_cast<InputType const*>(params.weight),
@@ -155,7 +157,7 @@ void cudaCoreGemmKernel(GemmParams const& params) {
             params.n,
             params.k,
             params.alpha);
-    }
+  }
 }
 
 template <typename InputType,
@@ -185,7 +187,9 @@ bool cuda_core_gemm_launcher(GemmParams const& params) {
   return cudaCoreGemmTemplateCaller<InputType, OutputType, 1, 2, 256>(params);
 }
 
-template bool cuda_core_gemm_launcher<__nv_fp8_e4m3, __nv_bfloat16>(GemmParams const&);
+template bool cuda_core_gemm_launcher<__nv_fp8_e4m3, __nv_bfloat16>(
+    GemmParams const&);
 template bool cuda_core_gemm_launcher<__nv_fp8_e4m3, half>(GemmParams const&);
-template bool cuda_core_gemm_launcher<__nv_fp8_e5m2, __nv_bfloat16>(GemmParams const&);
+template bool cuda_core_gemm_launcher<__nv_fp8_e5m2, __nv_bfloat16>(
+    GemmParams const&);
 template bool cuda_core_gemm_launcher<__nv_fp8_e5m2, half>(GemmParams const&);
diff --git a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.h b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.h
index 31eab1943d..e033375011 100644
--- a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.h
+++ b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.h
@@ -17,20 +17,22 @@
 #include "fp8_common.h"  // NOLINT
 
 typedef struct {
-    void const* act;
-    void const* weight;
-    void const* bias;
-    void* output;
-    int32_t m, n, k;
-    float alpha;
-    cudaStream_t stream;
+  void const* act;
+  void const* weight;
+  void const* bias;
+  void* output;
+  int32_t m, n, k;
+  float alpha;
+  cudaStream_t stream;
 } GemmParams;
 
 inline bool enable_cuda_core_fp8_gemm() {
-    static const char* enable_cuda_core_fp8_env = std::getenv("FLAGS_cuda_core_fp8_gemm");
-    static const bool enable_cuda_core_fp8_gemm =
-            enable_cuda_core_fp8_env != nullptr && std::string(enable_cuda_core_fp8_env) == "1";
-    return enable_cuda_core_fp8_gemm;
+  static const char* enable_cuda_core_fp8_env =
+      std::getenv("FLAGS_cuda_core_fp8_gemm");
+  static const bool enable_cuda_core_fp8_gemm =
+      enable_cuda_core_fp8_env != nullptr &&
+      std::string(enable_cuda_core_fp8_env) == "1";
+  return enable_cuda_core_fp8_gemm;
 }
 
 template <typename InputType, typename OutputType>
diff --git a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu
index c62b7effa1..660fc381b1 100644
--- a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu
+++ b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu
@@ -18,7 +18,6 @@
 #include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
 #include "fp8_fp8_half_cuda_core_gemm.h"
 
-
 paddle::Tensor cutlass_fp8_fp8_half_gemm_func(
     const paddle::Tensor& x,
     const paddle::Tensor& y,
@@ -28,153 +27,147 @@ paddle::Tensor cutlass_fp8_fp8_half_gemm_func(
     float scale,  // only support per-tensor quantization
     std::string output_dtype,
     std::string activation_type) {
-    paddle::Tensor out;
-    void* out_ptr = nullptr;
-    const void* x_ptr = nullptr;
-    const void* y_ptr = nullptr;
+  paddle::Tensor out;
+  void* out_ptr = nullptr;
+  const void* x_ptr = nullptr;
+  const void* y_ptr = nullptr;
 
-    auto place = x.place();
-    cudaStream_t stream = x.stream();
-    int64_t device_id = place.GetDeviceId();
-    int sm_version = GetGPUComputeCapability(device_id);
+  auto place = x.place();
+  cudaStream_t stream = x.stream();
+  int64_t device_id = place.GetDeviceId();
+  int sm_version = GetGPUComputeCapability(device_id);
 
-    int rank = x.dims().size();
-    int M = 0;
-    int K = 0;
-    int N = 0;
-    int ldd = 0;
+  int rank = x.dims().size();
+  int M = 0;
+  int K = 0;
+  int N = 0;
+  int ldd = 0;
 
-    int lda = x.dims()[rank - 1];
-    int ldb = y.dims()[rank - 1];
+  int lda = x.dims()[rank - 1];
+  int ldb = y.dims()[rank - 1];
 
-    if (!trans_x) {
-        M = x.dims()[rank - 2];
-        K = x.dims()[rank - 1];
+  if (!trans_x) {
+    M = x.dims()[rank - 2];
+    K = x.dims()[rank - 1];
 
-    } else {
-        M = x.dims()[rank - 1];
-        K = x.dims()[rank - 2];
-    }
-    if (!trans_y) {
-        N = y.dims()[rank - 1];
-        ldd = y.dims()[rank - 1];
-    } else {
-        N = y.dims()[rank - 2];
-        ldd = y.dims()[rank - 2];
-    }
+  } else {
+    M = x.dims()[rank - 1];
+    K = x.dims()[rank - 2];
+  }
+  if (!trans_y) {
+    N = y.dims()[rank - 1];
+    ldd = y.dims()[rank - 1];
+  } else {
+    N = y.dims()[rank - 2];
+    ldd = y.dims()[rank - 2];
+  }
 
-    int batch_count = 1;
-    for (size_t i = 0; i < rank - 2; ++i) {
-        batch_count *= x.dims()[i];
-    }
+  int batch_count = 1;
+  for (size_t i = 0; i < rank - 2; ++i) {
+    batch_count *= x.dims()[i];
+  }
 
-    std::string input_dtype = "";
-    if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
-        input_dtype = "float8_e4m3fn";
-        x_ptr =
-            reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
-        y_ptr =
-            reinterpret_cast<const void*>(y.data<phi::dtype::float8_e4m3fn>());
-    } else if (x.dtype() == phi::DataType::FLOAT8_E5M2) {
-        input_dtype = "float8_e5m2";
-        x_ptr =
-            reinterpret_cast<const void*>(x.data<phi::dtype::float8_e5m2>());
-        y_ptr =
-            reinterpret_cast<const void*>(y.data<phi::dtype::float8_e5m2>());
-    } else {
-        PADDLE_THROW(phi::errors::Fatal(
-            "fp8_fp8_half_gemm_fused only support e4m3 and e5m2 input"));
-    }
+  std::string input_dtype = "";
+  if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
+    input_dtype = "float8_e4m3fn";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
+    y_ptr = reinterpret_cast<const void*>(y.data<phi::dtype::float8_e4m3fn>());
+  } else if (x.dtype() == phi::DataType::FLOAT8_E5M2) {
+    input_dtype = "float8_e5m2";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e5m2>());
+    y_ptr = reinterpret_cast<const void*>(y.data<phi::dtype::float8_e5m2>());
+  } else {
+    PADDLE_THROW(phi::errors::Fatal(
+        "fp8_fp8_half_gemm_fused only support e4m3 and e5m2 input"));
+  }
 
-    std::vector<int64_t> out_shape = x.shape();
-    out_shape[rank - 1] = N;
-    out_shape[rank - 2] = M;
+  std::vector<int64_t> out_shape = x.shape();
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
 
+  if (output_dtype == "bfloat16") {
+    out = paddle::empty(out_shape, paddle::DataType::BFLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::bfloat16>());
+  } else if (output_dtype == "float16") {
+    out = paddle::empty(out_shape, paddle::DataType::FLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float16>());
+  } else {
+    PADDLE_THROW(
+        phi::errors::Fatal("fp8_fp8_half_gemm_fused only support bfloat16 "
+                           "and float16 output"));
+  }
+
+  std::string isbias = bias ? "true" : "false";
+  std::string act = (activation_type == "" || activation_type == "identity")
+                        ? "noact"
+                        : activation_type;
+
+  std::string fuse_gemm_config =
+      input_dtype + "_" + output_dtype + "_" + isbias + "_" + act;
+
+  void* bias_data = nullptr;
+  std::vector<int64_t> bias_dims{};
+  if (bias) {
+    bias_dims = common::vectorize(bias.get().dims());
     if (output_dtype == "bfloat16") {
-        out = paddle::empty(out_shape, paddle::DataType::BFLOAT16, x.place());
-        out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::bfloat16>());
-    } else if (output_dtype == "float16") {
-        out = paddle::empty(out_shape, paddle::DataType::FLOAT16, x.place());
-        out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float16>());
+      bias_data = reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
+          bias.get().data<phi::dtype::bfloat16>()));
     } else {
-        PADDLE_THROW(
-            phi::errors::Fatal("fp8_fp8_half_gemm_fused only support bfloat16 "
-                               "and float16 output"));
+      bias_data = reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
+          bias.get().data<phi::dtype::float16>()));
     }
+  }
 
-    std::string isbias = bias ? "true" : "false";
-    std::string act = (activation_type == "" || activation_type == "identity")
-                          ? "noact"
-                          : activation_type;
+  if (M <= 4 && trans_y && !trans_x && act == "noact" &&
+      enable_cuda_core_fp8_gemm()) {
+    GemmParams params = {
+        x_ptr,
+        y_ptr,
+        bias_data,
+        out_ptr,
+        M,
+        N,
+        K,
+        scale,
+        stream,
+    };
 
-    std::string fuse_gemm_config =
-        input_dtype + "_" + output_dtype + "_" + isbias + "_" + act;
+    if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
+      if (output_dtype == "bfloat16") {
+        cuda_core_gemm_launcher<__nv_fp8_e4m3, __nv_bfloat16>(params);
 
-    void* bias_data = nullptr;
-    std::vector<int64_t> bias_dims{};
-    if (bias) {
-        bias_dims = common::vectorize(bias.get().dims());
-        if (output_dtype == "bfloat16") {
-            bias_data =
-                reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
-                    bias.get().data<phi::dtype::bfloat16>()));
-        } else {
-            bias_data =
-                reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
-                    bias.get().data<phi::dtype::float16>()));
-        }
-    }
-
-    if (M <=4 && trans_y && !trans_x && act == "noact" && enable_cuda_core_fp8_gemm()) {
-        GemmParams params = {
-            x_ptr,
-            y_ptr,
-            bias_data,
-            out_ptr,
-            M,
-            N,
-            K,
-            scale,
-            stream,
-        };
-
-        if (x.dtype() == phi::DataType::FLOAT8_E4M3FN)
-        {
-            if(output_dtype == "bfloat16") {
-                cuda_core_gemm_launcher<__nv_fp8_e4m3, __nv_bfloat16>(params);
-
-            } else {
-                cuda_core_gemm_launcher<__nv_fp8_e4m3, half>(params);
-            }
-        } else {
-            if(output_dtype == "bfloat16") {
-                cuda_core_gemm_launcher<__nv_fp8_e5m2, __nv_bfloat16>(params);
-            } else {
-                cuda_core_gemm_launcher<__nv_fp8_e5m2, half>(params);
-            }
-        }
+      } else {
+        cuda_core_gemm_launcher<__nv_fp8_e4m3, half>(params);
+      }
     } else {
-        GemmEpilogueAllParams params = {x_ptr,
-                                        y_ptr,
-                                        out_ptr,
-                                        scale,
-                                        M,
-                                        N,
-                                        K,
-                                        lda,
-                                        ldb,
-                                        ldd,
-                                        batch_count,
-                                        place,
-                                        stream,
-                                        sm_version,
-                                        0.01,  // for leaky_relu
-                                        bias_data,
-                                        bias_dims,
-                                        fuse_gemm_config};
-        fp8_fp8_gemm_scale_bias_act(params);
+      if (output_dtype == "bfloat16") {
+        cuda_core_gemm_launcher<__nv_fp8_e5m2, __nv_bfloat16>(params);
+      } else {
+        cuda_core_gemm_launcher<__nv_fp8_e5m2, half>(params);
+      }
     }
-    return out;
+  } else {
+    GemmEpilogueAllParams params = {x_ptr,
+                                    y_ptr,
+                                    out_ptr,
+                                    scale,
+                                    M,
+                                    N,
+                                    K,
+                                    lda,
+                                    ldb,
+                                    ldd,
+                                    batch_count,
+                                    place,
+                                    stream,
+                                    sm_version,
+                                    0.01,  // for leaky_relu
+                                    bias_data,
+                                    bias_dims,
+                                    fuse_gemm_config};
+    fp8_fp8_gemm_scale_bias_act(params);
+  }
+  return out;
 }
 
 std::vector<paddle::Tensor> cutlass_fp8_fp8_half_gemm(
@@ -186,9 +179,8 @@ std::vector<paddle::Tensor> cutlass_fp8_fp8_half_gemm(
     float scale,  // only support per-tensor quantization
     std::string output_dtype,
     std::string activation_type) {
-    return {cutlass_fp8_fp8_half_gemm_func(
-            x, y, bias, trans_x, trans_y, scale,
-            output_dtype, activation_type)};
+  return {cutlass_fp8_fp8_half_gemm_func(
+      x, y, bias, trans_x, trans_y, scale, output_dtype, activation_type)};
 }
 
 std::vector<std::vector<int64_t>> CutlassFp8Fp8HalfGemmFusedInferShape(
@@ -197,33 +189,33 @@ std::vector<std::vector<int64_t>> CutlassFp8Fp8HalfGemmFusedInferShape(
     const paddle::optional<std::vector<int64_t>>& bias_shape,
     bool trans_x,
     bool trans_y) {
-    PADDLE_ENFORCE_EQ(x_shape.size(),
-                      y_shape.size(),
-                      phi::errors::InvalidArgument(
-                          "The rank of input X and Y should be equal, but "
-                          "received X's rank is %d, Y's rank is %d.",
-                          x_shape.size(),
-                          y_shape.size()));
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    y_shape.size(),
+                    phi::errors::InvalidArgument(
+                        "The rank of input X and Y should be equal, but "
+                        "received X's rank is %d, Y's rank is %d.",
+                        x_shape.size(),
+                        y_shape.size()));
 
-    int rank = x_shape.size();
-    int M = 0;
-    int N = 0;
+  int rank = x_shape.size();
+  int M = 0;
+  int N = 0;
 
-    if (!trans_x) {
-        M = x_shape[rank - 2];
+  if (!trans_x) {
+    M = x_shape[rank - 2];
 
-    } else {
-        M = x_shape[rank - 1];
-    }
-    if (!trans_y) {
-        N = y_shape[rank - 1];
-    } else {
-        N = y_shape[rank - 2];
-    }
-    std::vector<int64_t> out_shape = x_shape;
-    out_shape[rank - 1] = N;
-    out_shape[rank - 2] = M;
-    return {out_shape};
+  } else {
+    M = x_shape[rank - 1];
+  }
+  if (!trans_y) {
+    N = y_shape[rank - 1];
+  } else {
+    N = y_shape[rank - 2];
+  }
+  std::vector<int64_t> out_shape = x_shape;
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
+  return {out_shape};
 }
 
 std::vector<paddle::DataType> CutlassFp8Fp8HalfGemmFusedInferDtype(
@@ -234,15 +226,15 @@ std::vector<paddle::DataType> CutlassFp8Fp8HalfGemmFusedInferDtype(
     bool trans_y,
     float scale,  // only support per-tensor quantization
     std::string output_dtype) {
-    paddle::DataType data_type;
-    if (output_dtype == "bfloat16")
-        data_type = paddle::DataType::BFLOAT16;
-    else if (output_dtype == "float16")
-        data_type = paddle::DataType::FLOAT16;
-    else
-        PD_THROW(
-            "fp8_fp8_half_gemm_fused only support bfloat16 and float16 output");
-    return {data_type};
+  paddle::DataType data_type;
+  if (output_dtype == "bfloat16")
+    data_type = paddle::DataType::BFLOAT16;
+  else if (output_dtype == "float16")
+    data_type = paddle::DataType::FLOAT16;
+  else
+    PD_THROW(
+        "fp8_fp8_half_gemm_fused only support bfloat16 and float16 output");
+  return {data_type};
 }
 
 PD_BUILD_STATIC_OP(cutlass_fp8_fp8_half_gemm_fused)
diff --git a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/per_channel_fp8_fp8_half_gemm.cu b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/per_channel_fp8_fp8_half_gemm.cu
index 47138cfe6a..b7d7105b2c 100644
--- a/custom_ops/gpu_ops/fp8_gemm_with_cutlass/per_channel_fp8_fp8_half_gemm.cu
+++ b/custom_ops/gpu_ops/fp8_gemm_with_cutlass/per_channel_fp8_fp8_half_gemm.cu
@@ -27,126 +27,120 @@ std::vector<paddle::Tensor> cutlass_per_channel_fp8_gemm_fused(
     bool trans_x,
     bool trans_y,
     std::string output_dtype) {
-    auto scalar_scale_ptr = scalar_scale.data<float>();
-    auto channel_scale_ptr = channel_scale.data<float>();
-    paddle::Tensor out;
-    void* out_ptr = nullptr;
-    const void* x_ptr = nullptr;
-    const void* y_ptr = nullptr;
+  auto scalar_scale_ptr = scalar_scale.data<float>();
+  auto channel_scale_ptr = channel_scale.data<float>();
+  paddle::Tensor out;
+  void* out_ptr = nullptr;
+  const void* x_ptr = nullptr;
+  const void* y_ptr = nullptr;
 
-    auto place = x.place();
-    cudaStream_t stream = x.stream();
-    int64_t device_id = place.GetDeviceId();
-    int sm_version = GetGPUComputeCapability(device_id);
+  auto place = x.place();
+  cudaStream_t stream = x.stream();
+  int64_t device_id = place.GetDeviceId();
+  int sm_version = GetGPUComputeCapability(device_id);
 
-    int rank = x.dims().size();
-    int M = 0;
-    int K = 0;
-    int N = 0;
-    int ldd = 0;
+  int rank = x.dims().size();
+  int M = 0;
+  int K = 0;
+  int N = 0;
+  int ldd = 0;
 
-    int lda = x.dims()[rank - 1];
-    int ldb = y.dims()[rank - 1];
+  int lda = x.dims()[rank - 1];
+  int ldb = y.dims()[rank - 1];
 
-    if (!trans_x) {
-        M = x.dims()[rank - 2];
-        K = x.dims()[rank - 1];
+  if (!trans_x) {
+    M = x.dims()[rank - 2];
+    K = x.dims()[rank - 1];
 
-    } else {
-        M = x.dims()[rank - 1];
-        K = x.dims()[rank - 2];
-    }
-    if (!trans_y) {
-        N = y.dims()[rank - 1];
-        ldd = y.dims()[rank - 1];
-    } else {
-        N = y.dims()[rank - 2];
-        ldd = y.dims()[rank - 2];
-    }
+  } else {
+    M = x.dims()[rank - 1];
+    K = x.dims()[rank - 2];
+  }
+  if (!trans_y) {
+    N = y.dims()[rank - 1];
+    ldd = y.dims()[rank - 1];
+  } else {
+    N = y.dims()[rank - 2];
+    ldd = y.dims()[rank - 2];
+  }
 
-    int batch_count = 1;
-    for (size_t i = 0; i < rank - 2; ++i) {
-        batch_count *= x.dims()[i];
-    }
+  int batch_count = 1;
+  for (size_t i = 0; i < rank - 2; ++i) {
+    batch_count *= x.dims()[i];
+  }
 
-    std::string input_dtype = "";
-    if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
-        input_dtype = "float8_e4m3fn";
-        x_ptr =
-            reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
-        y_ptr =
-            reinterpret_cast<const void*>(y.data<phi::dtype::float8_e4m3fn>());
-    } else if (x.dtype() == phi::DataType::FLOAT8_E5M2) {
-        input_dtype = "float8_e5m2";
-        x_ptr =
-            reinterpret_cast<const void*>(x.data<phi::dtype::float8_e5m2>());
-        y_ptr =
-            reinterpret_cast<const void*>(y.data<phi::dtype::float8_e5m2>());
-    } else {
-        PADDLE_THROW(
-            phi::errors::Fatal("cutlass_per_channel_fp8_gemm_fused only "
-                               "support e4m3 and e5m2 input"));
-    }
+  std::string input_dtype = "";
+  if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
+    input_dtype = "float8_e4m3fn";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
+    y_ptr = reinterpret_cast<const void*>(y.data<phi::dtype::float8_e4m3fn>());
+  } else if (x.dtype() == phi::DataType::FLOAT8_E5M2) {
+    input_dtype = "float8_e5m2";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e5m2>());
+    y_ptr = reinterpret_cast<const void*>(y.data<phi::dtype::float8_e5m2>());
+  } else {
+    PADDLE_THROW(
+        phi::errors::Fatal("cutlass_per_channel_fp8_gemm_fused only "
+                           "support e4m3 and e5m2 input"));
+  }
 
-    std::vector<int64_t> out_shape = x.shape();
-    out_shape[rank - 1] = N;
-    out_shape[rank - 2] = M;
+  std::vector<int64_t> out_shape = x.shape();
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
 
+  if (output_dtype == "bfloat16") {
+    out = paddle::empty(out_shape, paddle::DataType::BFLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::bfloat16>());
+  } else if (output_dtype == "float16") {
+    out = paddle::empty(out_shape, paddle::DataType::FLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float16>());
+  } else {
+    PADDLE_THROW(phi::errors::Fatal(
+        "cutlass_per_channel_fp8_gemm_fused only support bfloat16 "
+        "and float16 output"));
+  }
+
+  std::string isbias = bias ? "true" : "false";
+
+  std::string fuse_gemm_config =
+      input_dtype + "_" + output_dtype + "_" + isbias;
+
+  void* bias_data = nullptr;
+  std::vector<int64_t> bias_dims{};
+  if (bias) {
+    bias_dims = common::vectorize(bias.get().dims());
     if (output_dtype == "bfloat16") {
-        out = paddle::empty(out_shape, paddle::DataType::BFLOAT16, x.place());
-        out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::bfloat16>());
-    } else if (output_dtype == "float16") {
-        out = paddle::empty(out_shape, paddle::DataType::FLOAT16, x.place());
-        out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float16>());
+      bias_data = reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
+          bias.get().data<phi::dtype::bfloat16>()));
     } else {
-        PADDLE_THROW(phi::errors::Fatal(
-            "cutlass_per_channel_fp8_gemm_fused only support bfloat16 "
-            "and float16 output"));
+      bias_data = reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
+          bias.get().data<phi::dtype::float16>()));
     }
+  }
 
-    std::string isbias = bias ? "true" : "false";
-
-    std::string fuse_gemm_config =
-        input_dtype + "_" + output_dtype + "_" + isbias;
-
-    void* bias_data = nullptr;
-    std::vector<int64_t> bias_dims{};
-    if (bias) {
-        bias_dims = common::vectorize(bias.get().dims());
-        if (output_dtype == "bfloat16") {
-            bias_data =
-                reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
-                    bias.get().data<phi::dtype::bfloat16>()));
-        } else {
-            bias_data =
-                reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
-                    bias.get().data<phi::dtype::float16>()));
-        }
-    }
-
-    GemmEpilogueAllParams params = {x_ptr,
-                                    y_ptr,
-                                    out_ptr,
-                                    0,
-                                    M,
-                                    N,
-                                    K,
-                                    lda,
-                                    ldb,
-                                    ldd,
-                                    batch_count,
-                                    place,
-                                    stream,
-                                    sm_version,
-                                    0.01,  // for leaky_relu
-                                    bias_data,
-                                    bias_dims,
-                                    fuse_gemm_config,
-                                    1,
-                                    scalar_scale_ptr,
-                                    channel_scale_ptr};
-    fp8_visitor_gemm_fused(params);
-    return {out};
+  GemmEpilogueAllParams params = {x_ptr,
+                                  y_ptr,
+                                  out_ptr,
+                                  0,
+                                  M,
+                                  N,
+                                  K,
+                                  lda,
+                                  ldb,
+                                  ldd,
+                                  batch_count,
+                                  place,
+                                  stream,
+                                  sm_version,
+                                  0.01,  // for leaky_relu
+                                  bias_data,
+                                  bias_dims,
+                                  fuse_gemm_config,
+                                  1,
+                                  scalar_scale_ptr,
+                                  channel_scale_ptr};
+  fp8_visitor_gemm_fused(params);
+  return {out};
 }
 
 std::vector<std::vector<int64_t>> CutlassPerChannelFp8GemmFusedInferShape(
@@ -158,33 +152,33 @@ std::vector<std::vector<int64_t>> CutlassPerChannelFp8GemmFusedInferShape(
     bool trans_x,
     bool trans_y,
     std::string output_dtype) {
-    PADDLE_ENFORCE_EQ(x_shape.size(),
-                      y_shape.size(),
-                      phi::errors::InvalidArgument(
-                          "The rank of input X and Y should be equal, but "
-                          "received X's rank is %d, Y's rank is %d.",
-                          x_shape.size(),
-                          y_shape.size()));
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    y_shape.size(),
+                    phi::errors::InvalidArgument(
+                        "The rank of input X and Y should be equal, but "
+                        "received X's rank is %d, Y's rank is %d.",
+                        x_shape.size(),
+                        y_shape.size()));
 
-    int rank = x_shape.size();
-    int M = 0;
-    int N = 0;
+  int rank = x_shape.size();
+  int M = 0;
+  int N = 0;
 
-    if (!trans_x) {
-        M = x_shape[rank - 2];
+  if (!trans_x) {
+    M = x_shape[rank - 2];
 
-    } else {
-        M = x_shape[rank - 1];
-    }
-    if (!trans_y) {
-        N = y_shape[rank - 1];
-    } else {
-        N = y_shape[rank - 2];
-    }
-    std::vector<int64_t> out_shape = x_shape;
-    out_shape[rank - 1] = N;
-    out_shape[rank - 2] = M;
-    return {out_shape};
+  } else {
+    M = x_shape[rank - 1];
+  }
+  if (!trans_y) {
+    N = y_shape[rank - 1];
+  } else {
+    N = y_shape[rank - 2];
+  }
+  std::vector<int64_t> out_shape = x_shape;
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
+  return {out_shape};
 }
 
 std::vector<paddle::DataType> CutlassPerChannelFp8GemmFusedInferDtype(
@@ -196,16 +190,16 @@ std::vector<paddle::DataType> CutlassPerChannelFp8GemmFusedInferDtype(
     bool trans_x,
     bool trans_y,
     std::string output_dtype) {
-    paddle::DataType data_type;
-    if (output_dtype == "bfloat16")
-        data_type = paddle::DataType::BFLOAT16;
-    else if (output_dtype == "float16")
-        data_type = paddle::DataType::FLOAT16;
-    else
-        PD_THROW(
-            "cutlass_per_channel_fp8_gemm_fused only support bfloat16 and "
-            "float16 output");
-    return {data_type};
+  paddle::DataType data_type;
+  if (output_dtype == "bfloat16")
+    data_type = paddle::DataType::BFLOAT16;
+  else if (output_dtype == "float16")
+    data_type = paddle::DataType::FLOAT16;
+  else
+    PD_THROW(
+        "cutlass_per_channel_fp8_gemm_fused only support bfloat16 and "
+        "float16 output");
+  return {data_type};
 }
 
 PD_BUILD_STATIC_OP(per_channel_fp8_fp8_half_gemm_fused)
diff --git a/custom_ops/gpu_ops/fused_get_rotary_embedding.cu b/custom_ops/gpu_ops/fused_get_rotary_embedding.cu
index 1942c561ae..6d5f50f832 100644
--- a/custom_ops/gpu_ops/fused_get_rotary_embedding.cu
+++ b/custom_ops/gpu_ops/fused_get_rotary_embedding.cu
@@ -32,65 +32,65 @@ constexpr int kNumWaves = 16;
 
 #ifdef PADDLE_WITH_HIP
 inline hipError_t GetNumBlocks(int64_t n, int *num_blocks) {
-    int dev;
-    {
-        hipError_t err = hipGetDevice(&dev);
-        if (err != hipSuccess) {
-            return err;
-        }
+  int dev;
+  {
+    hipError_t err = hipGetDevice(&dev);
+    if (err != hipSuccess) {
+      return err;
     }
-    int sm_count;
-    {
-        hipError_t err = hipDeviceGetAttribute(
-            &sm_count, hipDeviceAttributeMultiprocessorCount, dev);
-        if (err != hipSuccess) {
-            return err;
-        }
+  }
+  int sm_count;
+  {
+    hipError_t err = hipDeviceGetAttribute(
+        &sm_count, hipDeviceAttributeMultiprocessorCount, dev);
+    if (err != hipSuccess) {
+      return err;
     }
-    int tpm;
-    {
-        hipError_t err = hipDeviceGetAttribute(
-            &tpm, hipDeviceAttributeMaxThreadsPerMultiProcessor, dev);
-        if (err != hipSuccess) {
-            return err;
-        }
+  }
+  int tpm;
+  {
+    hipError_t err = hipDeviceGetAttribute(
+        &tpm, hipDeviceAttributeMaxThreadsPerMultiProcessor, dev);
+    if (err != hipSuccess) {
+      return err;
     }
-    *num_blocks = std::max<int>(
-        1,
-        std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
-                          sm_count * tpm / kBlockSize * kNumWaves));
-    return hipSuccess;
+  }
+  *num_blocks =
+      std::max<int>(1,
+                    std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
+                                      sm_count * tpm / kBlockSize * kNumWaves));
+  return hipSuccess;
 }
 #else
 inline cudaError_t GetNumBlocks(int64_t n, int *num_blocks) {
-    int dev;
-    {
-        cudaError_t err = cudaGetDevice(&dev);
-        if (err != cudaSuccess) {
-            return err;
-        }
+  int dev;
+  {
+    cudaError_t err = cudaGetDevice(&dev);
+    if (err != cudaSuccess) {
+      return err;
     }
-    int sm_count;
-    {
-        cudaError_t err = cudaDeviceGetAttribute(
-            &sm_count, cudaDevAttrMultiProcessorCount, dev);
-        if (err != cudaSuccess) {
-            return err;
-        }
+  }
+  int sm_count;
+  {
+    cudaError_t err =
+        cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
+    if (err != cudaSuccess) {
+      return err;
     }
-    int tpm;
-    {
-        cudaError_t err = cudaDeviceGetAttribute(
-            &tpm, cudaDevAttrMaxThreadsPerMultiProcessor, dev);
-        if (err != cudaSuccess) {
-            return err;
-        }
+  }
+  int tpm;
+  {
+    cudaError_t err = cudaDeviceGetAttribute(
+        &tpm, cudaDevAttrMaxThreadsPerMultiProcessor, dev);
+    if (err != cudaSuccess) {
+      return err;
     }
-    *num_blocks = std::max<int>(
-        1,
-        std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
-                          sm_count * tpm / kBlockSize * kNumWaves));
-    return cudaSuccess;
+  }
+  *num_blocks =
+      std::max<int>(1,
+                    std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
+                                      sm_count * tpm / kBlockSize * kNumWaves));
+  return cudaSuccess;
 }
 #endif
 template <paddle::DataType D>
@@ -98,16 +98,16 @@ class PDTraits;
 
 template <>
 class PDTraits<paddle::DataType::FLOAT32> {
-    public:
-    typedef float DataType;
-    typedef float data_t;
+ public:
+  typedef float DataType;
+  typedef float data_t;
 };
 
 template <>
 class PDTraits<paddle::DataType::FLOAT16> {
-    public:
-    typedef half DataType;
-    typedef paddle::float16 data_t;
+ public:
+  typedef half DataType;
+  typedef paddle::float16 data_t;
 };
 
 /*
@@ -116,8 +116,8 @@ Position_ids: bsz, max_seq_length
 
 template <typename T, int N>
 struct GetPackType {
-    using type =
-        typename std::aligned_storage<N * sizeof(T), N * sizeof(T)>::type;
+  using type =
+      typename std::aligned_storage<N * sizeof(T), N * sizeof(T)>::type;
 };
 
 template <typename T, int N>
@@ -125,12 +125,12 @@ using PackType = typename GetPackType<T, N>::type;
 
 template <typename T, int N>
 union Pack {
-    static_assert(sizeof(PackType<T, N>) == sizeof(T) * N, "");
-    __device__ Pack() {
-        // do nothing
-    }
-    PackType<T, N> storage;
-    T elem[N];
+  static_assert(sizeof(PackType<T, N>) == sizeof(T) * N, "");
+  __device__ Pack() {
+    // do nothing
+  }
+  PackType<T, N> storage;
+  T elem[N];
 };
 
 __global__ __launch_bounds__(kBlockSize) void fused_get_rotary_embedding(
@@ -143,106 +143,105 @@ __global__ __launch_bounds__(kBlockSize) void fused_get_rotary_embedding(
     const float inv_head_dim,
     const int32_t elem_cnt,
     float *rope_embedding) {
-    /*
-    In Naive implementation, it will stacks [freqs, freqs]
-    And actually, each threads can process 1 values, and store continuous 2 same
-    values. So here We construct a Pack to store 2 values.
-    */
-    constexpr int PackSize = 2;
-    Pack<float, PackSize> SinStorePack{};
-    Pack<float, PackSize> CosStorePack{};
+  /*
+  In Naive implementation, it will stacks [freqs, freqs]
+  And actually, each threads can process 1 values, and store continuous 2 same
+  values. So here We construct a Pack to store 2 values.
+  */
+  constexpr int PackSize = 2;
+  Pack<float, PackSize> SinStorePack{};
+  Pack<float, PackSize> CosStorePack{};
 
-    const int half_head_dim = head_dim / PackSize;
-    const int32_t global_thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
-    for (int idx = global_thread_idx, step = blockDim.x * gridDim.x;
-         idx < elem_cnt;
-         idx += step) {
-        const int32_t bsz_seq_idx = idx / half_head_dim;
-        const int32_t bsz_idx = bsz_seq_idx / max_seq_length;
-        const int32_t seq_idx = bsz_seq_idx % max_seq_length;
-        const int64_t position_offset =
-            bsz_idx * max_position_seq_length + seq_idx + prompt_num;
-        const int32_t half_head_idx = (idx % half_head_dim) * PackSize;
-        const float exponent_factor =
-            -static_cast<float>(half_head_idx) *
-            inv_head_dim;  // * inv_head_dim equals to / head_dim.
-        const float inv_freq_val = powf(10000.0f, exponent_factor);
-        const float freqs_val = position_ids[position_offset] * inv_freq_val;
-        const float cos_embedding_val = cos(freqs_val);
-        const float sin_embedding_val = sin(freqs_val);
+  const int half_head_dim = head_dim / PackSize;
+  const int32_t global_thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int idx = global_thread_idx, step = blockDim.x * gridDim.x;
+       idx < elem_cnt;
+       idx += step) {
+    const int32_t bsz_seq_idx = idx / half_head_dim;
+    const int32_t bsz_idx = bsz_seq_idx / max_seq_length;
+    const int32_t seq_idx = bsz_seq_idx % max_seq_length;
+    const int64_t position_offset =
+        bsz_idx * max_position_seq_length + seq_idx + prompt_num;
+    const int32_t half_head_idx = (idx % half_head_dim) * PackSize;
+    const float exponent_factor =
+        -static_cast<float>(half_head_idx) *
+        inv_head_dim;  // * inv_head_dim equals to / head_dim.
+    const float inv_freq_val = powf(10000.0f, exponent_factor);
+    const float freqs_val = position_ids[position_offset] * inv_freq_val;
+    const float cos_embedding_val = cos(freqs_val);
+    const float sin_embedding_val = sin(freqs_val);
 
 /*
 Since After stack, the continuous 2 elements value is same.
 So here each threads store 2 computed embedding value.
 */
 #pragma unroll
-        for (int unroll_idx = 0; unroll_idx < PackSize; unroll_idx++) {
-            CosStorePack.elem[unroll_idx] = cos_embedding_val;
-            SinStorePack.elem[unroll_idx] = sin_embedding_val;
-        }
-
-        const int32_t cos_offset = bsz_seq_idx * head_dim + half_head_idx;
-        const int32_t sin_offset = bsz * max_seq_length * head_dim + cos_offset;
-        *(reinterpret_cast<PackType<float, PackSize> *>(
-            rope_embedding + cos_offset)) = CosStorePack.storage;
-        *(reinterpret_cast<PackType<float, PackSize> *>(
-            rope_embedding + sin_offset)) = SinStorePack.storage;
+    for (int unroll_idx = 0; unroll_idx < PackSize; unroll_idx++) {
+      CosStorePack.elem[unroll_idx] = cos_embedding_val;
+      SinStorePack.elem[unroll_idx] = sin_embedding_val;
     }
+
+    const int32_t cos_offset = bsz_seq_idx * head_dim + half_head_idx;
+    const int32_t sin_offset = bsz * max_seq_length * head_dim + cos_offset;
+    *(reinterpret_cast<PackType<float, PackSize> *>(
+        rope_embedding + cos_offset)) = CosStorePack.storage;
+    *(reinterpret_cast<PackType<float, PackSize> *>(
+        rope_embedding + sin_offset)) = SinStorePack.storage;
+  }
 }
 
 std::vector<paddle::Tensor> GetRoPE(const paddle::Tensor &input_ids,
                                     const paddle::Tensor &position_ids,
                                     const paddle::Tensor &head_dim_shape_tensor,
                                     int prompt_num) {
-    const int64_t batch_size = input_ids.shape()[0];
-    const int64_t max_seq_length = input_ids.shape()[1];
-    const int64_t max_position_seq_length = position_ids.shape()[1];
-    const int64_t head_dim = head_dim_shape_tensor.shape()[0];
-    const float inv_head_dim = 1.0f / static_cast<float>(head_dim);
+  const int64_t batch_size = input_ids.shape()[0];
+  const int64_t max_seq_length = input_ids.shape()[1];
+  const int64_t max_position_seq_length = position_ids.shape()[1];
+  const int64_t head_dim = head_dim_shape_tensor.shape()[0];
+  const float inv_head_dim = 1.0f / static_cast<float>(head_dim);
 
-    auto cu_stream = position_ids.stream();
+  auto cu_stream = position_ids.stream();
 
-    auto rotary_embedding =
-        paddle::empty({2, batch_size, 1, max_seq_length, head_dim},
-                      paddle::DataType::FLOAT32,
-                      position_ids.place());
+  auto rotary_embedding =
+      paddle::empty({2, batch_size, 1, max_seq_length, head_dim},
+                    paddle::DataType::FLOAT32,
+                    position_ids.place());
 
-    assert(head_dim % 2 == 0);
-    const int32_t elem_cnt = batch_size * max_seq_length * head_dim / 2;
-    int32_t grid_size = 1;
-    GetNumBlocks(elem_cnt, &grid_size);
+  assert(head_dim % 2 == 0);
+  const int32_t elem_cnt = batch_size * max_seq_length * head_dim / 2;
+  int32_t grid_size = 1;
+  GetNumBlocks(elem_cnt, &grid_size);
 
-    fused_get_rotary_embedding<<<grid_size, kBlockSize, 0, cu_stream>>>(
-        position_ids.data<float>(),
-        batch_size,
-        max_seq_length,
-        max_position_seq_length,
-        head_dim,
-        prompt_num,
-        inv_head_dim,
-        elem_cnt,
-        reinterpret_cast<float *>(rotary_embedding.data<float>()));
-    return {rotary_embedding};
+  fused_get_rotary_embedding<<<grid_size, kBlockSize, 0, cu_stream>>>(
+      position_ids.data<float>(),
+      batch_size,
+      max_seq_length,
+      max_position_seq_length,
+      head_dim,
+      prompt_num,
+      inv_head_dim,
+      elem_cnt,
+      reinterpret_cast<float *>(rotary_embedding.data<float>()));
+  return {rotary_embedding};
 }
 
 std::vector<std::vector<int64_t>> GetRoPEInferShape(
     const std::vector<int64_t> &input_ids_shape,
     const std::vector<int64_t> &position_ids_shape,
     const std::vector<int64_t> &head_dim_shape_tensor_shape) {
-    const int64_t batch_size = position_ids_shape[0];
-    const int64_t max_seq_length = input_ids_shape[1];
-    const int64_t head_dim = head_dim_shape_tensor_shape[0];
-    std::vector<int64_t> out_shape = {
-        2, batch_size, 1, max_seq_length, head_dim};
-    return {out_shape};
+  const int64_t batch_size = position_ids_shape[0];
+  const int64_t max_seq_length = input_ids_shape[1];
+  const int64_t head_dim = head_dim_shape_tensor_shape[0];
+  std::vector<int64_t> out_shape = {2, batch_size, 1, max_seq_length, head_dim};
+  return {out_shape};
 }
 
 std::vector<paddle::DataType> GetRoPEInferDtype(
     const paddle::DataType &input_ids_dtype,
     const paddle::DataType &position_ids_dtype,
     const paddle::DataType &head_dim_shape_tensor_dtype) {
-    // RoPE output dtype is Float.
-    return {paddle::DataType::FLOAT32};
+  // RoPE output dtype is Float.
+  return {paddle::DataType::FLOAT32};
 }
 
 PD_BUILD_STATIC_OP(fused_get_rotary_embedding)
diff --git a/custom_ops/gpu_ops/fused_hadamard_quant_fp8.cu b/custom_ops/gpu_ops/fused_hadamard_quant_fp8.cu
index 6ad1901029..0e78a04061 100644
--- a/custom_ops/gpu_ops/fused_hadamard_quant_fp8.cu
+++ b/custom_ops/gpu_ops/fused_hadamard_quant_fp8.cu
@@ -24,14 +24,14 @@
 #include "helper.h"
 
 __device__ __forceinline__ void hadamard32_warp(__nv_bfloat16& x) {
-    int lane_id = threadIdx.x % 32;
+  int lane_id = threadIdx.x % 32;
 #pragma unroll
-    for (int step = 0; step < 5; ++step) {
-        const int lane_mask = 1 << step;
-        const __nv_bfloat16 sign = (lane_id & lane_mask) ? -1.f : 1.f;
-        __nv_bfloat16 x_val_other = __shfl_xor_sync(0xffffffff, x, lane_mask);
-        x = sign * x + x_val_other;
-    }
+  for (int step = 0; step < 5; ++step) {
+    const int lane_mask = 1 << step;
+    const __nv_bfloat16 sign = (lane_id & lane_mask) ? -1.f : 1.f;
+    __nv_bfloat16 x_val_other = __shfl_xor_sync(0xffffffff, x, lane_mask);
+    x = sign * x + x_val_other;
+  }
 }
 
 __global__ void MoeFusedHadamardQuantFp8Kernel(
@@ -41,27 +41,26 @@ __global__ void MoeFusedHadamardQuantFp8Kernel(
     __nv_fp8_e4m3* out,
     const int top_k,
     const int intermediate_size,
-    const int64_t numel
-) {
-    int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x;
-    if (out_idx >= numel) return;
+    const int64_t numel) {
+  int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (out_idx >= numel) return;
 
-    int64_t token_idx = out_idx / (top_k * intermediate_size);
-    int64_t topk_idx  = (out_idx / intermediate_size) % top_k;
-    int64_t inter_idx = out_idx % intermediate_size;
+  int64_t token_idx = out_idx / (top_k * intermediate_size);
+  int64_t topk_idx = (out_idx / intermediate_size) % top_k;
+  int64_t inter_idx = out_idx % intermediate_size;
 
-    int64_t input_idx = token_idx * intermediate_size + inter_idx;
-    if (input_idx >= numel / top_k) return;
+  int64_t input_idx = token_idx * intermediate_size + inter_idx;
+  if (input_idx >= numel / top_k) return;
 
-    int64_t expert_id = topk_ids[token_idx * top_k + topk_idx];
-    float scale_value = scale[expert_id];
+  int64_t expert_id = topk_ids[token_idx * top_k + topk_idx];
+  float scale_value = scale[expert_id];
 
-    __nv_bfloat16 x = input[input_idx];
-    hadamard32_warp(x);
+  __nv_bfloat16 x = input[input_idx];
+  hadamard32_warp(x);
 
-    float x_fp32 = __bfloat162float(x);
-    float quantized = x_fp32 / scale_value;
-    out[out_idx] = static_cast<__nv_fp8_e4m3>(quantized);
+  float x_fp32 = __bfloat162float(x);
+  float quantized = x_fp32 / scale_value;
+  out[out_idx] = static_cast<__nv_fp8_e4m3>(quantized);
 }
 
 __global__ void MoeFusedHadamardQuantFp8TiledKernel(
@@ -71,117 +70,116 @@ __global__ void MoeFusedHadamardQuantFp8TiledKernel(
     __nv_fp8_e4m3* out,
     const int top_k,
     const int intermediate_size,
-    const int64_t numel
-) {
-    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-    if (idx >= numel) return;
+    const int64_t numel) {
+  int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= numel) return;
 
-    int64_t token_idx = idx / intermediate_size;
-    int64_t expert_id = topk_ids[token_idx];
-    float scale_value = scale[expert_id];
+  int64_t token_idx = idx / intermediate_size;
+  int64_t expert_id = topk_ids[token_idx];
+  float scale_value = scale[expert_id];
 
-    __nv_bfloat16 x = input[idx];
-    hadamard32_warp(x);
+  __nv_bfloat16 x = input[idx];
+  hadamard32_warp(x);
 
-    float x_fp32 = __bfloat162float(x);
-    float quantized = x_fp32 / scale_value;
-    out[idx] = static_cast<__nv_fp8_e4m3>(quantized);
+  float x_fp32 = __bfloat162float(x);
+  float quantized = x_fp32 / scale_value;
+  out[idx] = static_cast<__nv_fp8_e4m3>(quantized);
 }
 
 std::vector<paddle::Tensor> MoeFusedHadamardQuantFp8(
-                                const paddle::Tensor &input,
-                                const paddle::Tensor &scale,
-                                const paddle::Tensor &topk_ids,
-                                const int top_k,
-                                const int intermediate_size,
-                                const bool tiled) {
-    int64_t numel = input.numel();
-    if (!tiled) numel *= top_k;
-    paddle::Tensor out = GetEmptyTensor(
-            {numel / intermediate_size, intermediate_size},
-            paddle::DataType::FLOAT8_E4M3FN,
-            input.place());
-    constexpr int64_t thread_per_block = 256;
-    int64_t block_per_grid = (numel + thread_per_block - 1) / thread_per_block;
-    auto stream = input.stream();
-    if (tiled) {
-        MoeFusedHadamardQuantFp8TiledKernel<<<block_per_grid, thread_per_block, 0, stream>>>(
-            reinterpret_cast<const __nv_bfloat16*>(input.data<paddle::bfloat16>()),
-            scale.data<float>(),
-            topk_ids.data<int64_t>(),
-            reinterpret_cast<__nv_fp8_e4m3*>(out.mutable_data<phi::dtype::float8_e4m3fn>()),
-            top_k,
-            intermediate_size,
-            numel
-        );
-    } else {
-        MoeFusedHadamardQuantFp8Kernel<<<block_per_grid, thread_per_block, 0, stream>>>(
-            reinterpret_cast<const __nv_bfloat16*>(input.data<phi::dtype::bfloat16>()),
-            scale.data<float>(),
-            topk_ids.data<int64_t>(),
-            reinterpret_cast<__nv_fp8_e4m3*>(out.mutable_data<phi::dtype::float8_e4m3fn>()),
-            top_k,
-            intermediate_size,
-            numel
-        );
-    }
-    return {out};
+    const paddle::Tensor& input,
+    const paddle::Tensor& scale,
+    const paddle::Tensor& topk_ids,
+    const int top_k,
+    const int intermediate_size,
+    const bool tiled) {
+  int64_t numel = input.numel();
+  if (!tiled) numel *= top_k;
+  paddle::Tensor out =
+      GetEmptyTensor({numel / intermediate_size, intermediate_size},
+                     paddle::DataType::FLOAT8_E4M3FN,
+                     input.place());
+  constexpr int64_t thread_per_block = 256;
+  int64_t block_per_grid = (numel + thread_per_block - 1) / thread_per_block;
+  auto stream = input.stream();
+  if (tiled) {
+    MoeFusedHadamardQuantFp8TiledKernel<<<block_per_grid,
+                                          thread_per_block,
+                                          0,
+                                          stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(input.data<paddle::bfloat16>()),
+        scale.data<float>(),
+        topk_ids.data<int64_t>(),
+        reinterpret_cast<__nv_fp8_e4m3*>(
+            out.mutable_data<phi::dtype::float8_e4m3fn>()),
+        top_k,
+        intermediate_size,
+        numel);
+  } else {
+    MoeFusedHadamardQuantFp8Kernel<<<block_per_grid,
+                                     thread_per_block,
+                                     0,
+                                     stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(
+            input.data<phi::dtype::bfloat16>()),
+        scale.data<float>(),
+        topk_ids.data<int64_t>(),
+        reinterpret_cast<__nv_fp8_e4m3*>(
+            out.mutable_data<phi::dtype::float8_e4m3fn>()),
+        top_k,
+        intermediate_size,
+        numel);
+  }
+  return {out};
 }
 
 PD_BUILD_STATIC_OP(moe_fused_hadamard_quant_fp8)
     .Inputs({"input", "scale", "topk_ids"})
     .Outputs({"output"})
-    .Attrs({"top_k: int",
-            "intermediate_size: int",
-            "tiled: bool"})
+    .Attrs({"top_k: int", "intermediate_size: int", "tiled: bool"})
     .SetKernelFn(PD_KERNEL(MoeFusedHadamardQuantFp8));
 
-
-paddle::Tensor MoeFusedHadamardQuantFp8Func(
-                const paddle::Tensor &input,
-                const paddle::Tensor &scale,
-                const paddle::Tensor &topk_ids,
-                const int top_k,
-                const int intermediate_size,
-                const bool tiled) {
-    return MoeFusedHadamardQuantFp8(input, scale, topk_ids, top_k, intermediate_size, tiled)[0];
+paddle::Tensor MoeFusedHadamardQuantFp8Func(const paddle::Tensor& input,
+                                            const paddle::Tensor& scale,
+                                            const paddle::Tensor& topk_ids,
+                                            const int top_k,
+                                            const int intermediate_size,
+                                            const bool tiled) {
+  return MoeFusedHadamardQuantFp8(
+      input, scale, topk_ids, top_k, intermediate_size, tiled)[0];
 }
 
-
 __global__ void FusedHadamardQuantFp8Kernel(
-                              const __nv_bfloat16* __restrict__ input,
-                              __nv_fp8_e4m3* out,
-                              const float scale,
-                              const int64_t numel) {
-    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-    if (idx >= numel) return;
+    const __nv_bfloat16* __restrict__ input,
+    __nv_fp8_e4m3* out,
+    const float scale,
+    const int64_t numel) {
+  int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= numel) return;
 
-    __nv_bfloat16 x = input[idx];
-    hadamard32_warp(x);
+  __nv_bfloat16 x = input[idx];
+  hadamard32_warp(x);
 
-    float x_fp32 = __bfloat162float(x);
-    float quantized = x_fp32 / scale;
-    out[idx] = static_cast<__nv_fp8_e4m3>(quantized);
+  float x_fp32 = __bfloat162float(x);
+  float quantized = x_fp32 / scale;
+  out[idx] = static_cast<__nv_fp8_e4m3>(quantized);
 }
 
-std::vector<paddle::Tensor> FusedHadamardQuantFp8(
-                                const paddle::Tensor &input,
-                                const float scale) {
-    int64_t numel = input.numel();
-    paddle::Tensor out = GetEmptyTensor(
-            input.dims(),
-            paddle::DataType::FLOAT8_E4M3FN,
-            input.place());
-    constexpr int64_t thread_per_block = 256;
-    int64_t block_per_grid = (numel + thread_per_block - 1) / thread_per_block;
-    auto stream = input.stream();
-    FusedHadamardQuantFp8Kernel<<<block_per_grid, thread_per_block, 0, stream>>>(
-        reinterpret_cast<const __nv_bfloat16*>(input.data<paddle::bfloat16>()),
-        reinterpret_cast<__nv_fp8_e4m3*>(out.mutable_data<phi::dtype::float8_e4m3fn>()),
-        scale,
-        numel
-    );
-    return {out};
+std::vector<paddle::Tensor> FusedHadamardQuantFp8(const paddle::Tensor& input,
+                                                  const float scale) {
+  int64_t numel = input.numel();
+  paddle::Tensor out = GetEmptyTensor(
+      input.dims(), paddle::DataType::FLOAT8_E4M3FN, input.place());
+  constexpr int64_t thread_per_block = 256;
+  int64_t block_per_grid = (numel + thread_per_block - 1) / thread_per_block;
+  auto stream = input.stream();
+  FusedHadamardQuantFp8Kernel<<<block_per_grid, thread_per_block, 0, stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(input.data<paddle::bfloat16>()),
+      reinterpret_cast<__nv_fp8_e4m3*>(
+          out.mutable_data<phi::dtype::float8_e4m3fn>()),
+      scale,
+      numel);
+  return {out};
 }
 
 PD_BUILD_STATIC_OP(fused_hadamard_quant_fp8)
@@ -190,9 +188,7 @@ PD_BUILD_STATIC_OP(fused_hadamard_quant_fp8)
     .Attrs({"scale: float"})
     .SetKernelFn(PD_KERNEL(FusedHadamardQuantFp8));
 
-
-paddle::Tensor FusedHadamardQuantFp8Func(
-                const paddle::Tensor &input,
-                const float scale) {
-    return FusedHadamardQuantFp8(input, scale)[0];
+paddle::Tensor FusedHadamardQuantFp8Func(const paddle::Tensor& input,
+                                         const float scale) {
+  return FusedHadamardQuantFp8(input, scale)[0];
 }
diff --git a/custom_ops/gpu_ops/fused_rotary_position_encoding.cu b/custom_ops/gpu_ops/fused_rotary_position_encoding.cu
index a0462ba346..05c94b60b7 100644
--- a/custom_ops/gpu_ops/fused_rotary_position_encoding.cu
+++ b/custom_ops/gpu_ops/fused_rotary_position_encoding.cu
@@ -15,7 +15,6 @@
 #include "helper.h"
 #include "paddle/extension.h"
 
-
 template <typename T, bool IS_NEOX>
 inline __device__ void apply_token_rotary_embedding_kernel(
     T* __restrict__ arr,
@@ -43,7 +42,6 @@ inline __device__ void apply_token_rotary_embedding_kernel(
   arr[y_index] = y * cos + x * sin;
 }
 
-
 template <typename T, bool IS_NEOX>
 __global__ void apply_rotary_embedding_kernel(
     T* __restrict__ query,  // [num_tokens, num_heads, head_size]
@@ -84,7 +82,6 @@ __global__ void apply_rotary_embedding_kernel(
   }
 }
 
-
 void FusedRotaryPositionEncoding(
     paddle::Tensor& query,  // [num_tokens, num_heads, head_size] or
                             // [num_tokens, num_heads * head_size]
diff --git a/custom_ops/gpu_ops/gather_idx.cu b/custom_ops/gpu_ops/gather_idx.cu
index 248514a560..877a8a21ce 100644
--- a/custom_ops/gpu_ops/gather_idx.cu
+++ b/custom_ops/gpu_ops/gather_idx.cu
@@ -33,27 +33,27 @@ class PDTraits;
 
 template <>
 class PDTraits<paddle::DataType::FLOAT32> {
-    public:
-    typedef float DataType;
-    typedef float data_t;
+ public:
+  typedef float DataType;
+  typedef float data_t;
 };
 
 template <>
 class PDTraits<paddle::DataType::FLOAT16> {
-    public:
-    typedef half DataType;
-    typedef paddle::float16 data_t;
+ public:
+  typedef half DataType;
+  typedef paddle::float16 data_t;
 };
 
 template <>
 class PDTraits<paddle::DataType::BFLOAT16> {
-    public:
+ public:
 #ifdef PADDLE_WITH_HIP
-    typedef hip_bfloat16 DataType;
+  typedef hip_bfloat16 DataType;
 #else
-    typedef __nv_bfloat16 DataType;
+  typedef __nv_bfloat16 DataType;
 #endif
-    typedef paddle::bfloat16 data_t;
+  typedef paddle::bfloat16 data_t;
 };
 
 template <typename T>
@@ -63,86 +63,86 @@ __global__ void get_value_by_id(const T *logits,
                                 int bs,
                                 int seq_len,
                                 int length) {
-    int bid = blockIdx.x;
-    int tid = threadIdx.x;
-    int idx = bid * blockDim.x + tid;
-    for (int i = idx; i < bs * length; i += gridDim.x * blockDim.x) {
-        int bi = i / length;
-        int lane = i % length;
-        int si = ids[bi];
-        if (si == -1) {
-            si = 0;
-        }
-        const T *logits_now = logits + bi * seq_len * length + si * length;
-        T *logits_out_now = logits_out + bi * length;
-        logits_out_now[lane] = logits_now[lane];
+  int bid = blockIdx.x;
+  int tid = threadIdx.x;
+  int idx = bid * blockDim.x + tid;
+  for (int i = idx; i < bs * length; i += gridDim.x * blockDim.x) {
+    int bi = i / length;
+    int lane = i % length;
+    int si = ids[bi];
+    if (si == -1) {
+      si = 0;
     }
+    const T *logits_now = logits + bi * seq_len * length + si * length;
+    T *logits_out_now = logits_out + bi * length;
+    logits_out_now[lane] = logits_now[lane];
+  }
 }
 
 template <paddle::DataType D>
 std::vector<paddle::Tensor> gather_idx(const paddle::Tensor &logits,
                                        const paddle::Tensor &gather_id) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    PD_CHECK(gather_id.dtype() == paddle::DataType::INT32);
-    auto cu_stream = logits.stream();
-    std::vector<int64_t> logits_shape = logits.shape();
-    std::vector<int64_t> id_shape = gather_id.shape();
-    int logits_bs = logits_shape[0];
-    int seq_len = logits_shape[1];
-    int logits_len = logits_shape[2];
-    auto logits_out =
-        paddle::empty({logits_bs, logits_len}, logits.type(), logits.place());
-    int id_bs = id_shape[0];
-    int64_t numels = logits_bs * logits_len;
-    int block_size = 128;
-    int grid_size = (numels + block_size - 1) / block_size;
-    get_value_by_id<<<grid_size, block_size, 0, cu_stream>>>(
-        reinterpret_cast<DataType_ *>(
-            const_cast<data_t *>(logits.data<data_t>())),
-        gather_id.data<int>(),
-        reinterpret_cast<DataType_ *>(
-            const_cast<data_t *>(logits_out.data<data_t>())),
-        logits_bs,
-        seq_len,
-        logits_len);
-    return {logits_out};
+  PD_CHECK(gather_id.dtype() == paddle::DataType::INT32);
+  auto cu_stream = logits.stream();
+  std::vector<int64_t> logits_shape = logits.shape();
+  std::vector<int64_t> id_shape = gather_id.shape();
+  int logits_bs = logits_shape[0];
+  int seq_len = logits_shape[1];
+  int logits_len = logits_shape[2];
+  auto logits_out =
+      paddle::empty({logits_bs, logits_len}, logits.type(), logits.place());
+  int id_bs = id_shape[0];
+  int64_t numels = logits_bs * logits_len;
+  int block_size = 128;
+  int grid_size = (numels + block_size - 1) / block_size;
+  get_value_by_id<<<grid_size, block_size, 0, cu_stream>>>(
+      reinterpret_cast<DataType_ *>(
+          const_cast<data_t *>(logits.data<data_t>())),
+      gather_id.data<int>(),
+      reinterpret_cast<DataType_ *>(
+          const_cast<data_t *>(logits_out.data<data_t>())),
+      logits_bs,
+      seq_len,
+      logits_len);
+  return {logits_out};
 }
 
 std::vector<paddle::Tensor> GatherIdx(const paddle::Tensor &logits,
                                       const paddle::Tensor &gather_id) {
-    switch (logits.type()) {
-        case paddle::DataType::BFLOAT16: {
-            return gather_idx<paddle::DataType::BFLOAT16>(logits, gather_id);
-        }
-        case paddle::DataType::FLOAT16: {
-            return gather_idx<paddle::DataType::FLOAT16>(logits, gather_id);
-        }
-        case paddle::DataType::FLOAT32: {
-            return gather_idx<paddle::DataType::FLOAT32>(logits, gather_id);
-        }
-        default: {
-            PD_THROW(
-                "NOT supported data type. "
-                "Only bfloat16, float16 and float32 are supported. ");
-            break;
-        }
+  switch (logits.type()) {
+    case paddle::DataType::BFLOAT16: {
+      return gather_idx<paddle::DataType::BFLOAT16>(logits, gather_id);
     }
+    case paddle::DataType::FLOAT16: {
+      return gather_idx<paddle::DataType::FLOAT16>(logits, gather_id);
+    }
+    case paddle::DataType::FLOAT32: {
+      return gather_idx<paddle::DataType::FLOAT32>(logits, gather_id);
+    }
+    default: {
+      PD_THROW(
+          "NOT supported data type. "
+          "Only bfloat16, float16 and float32 are supported. ");
+      break;
+    }
+  }
 }
 
 std::vector<std::vector<int64_t>> GatherIdxInferShape(
     const std::vector<int64_t> &logits_shape,
     const std::vector<int64_t> &gather_id_shape) {
-    std::vector<int64_t> out_shape = {logits_shape[0], logits_shape[2]};
-    return {out_shape};
+  std::vector<int64_t> out_shape = {logits_shape[0], logits_shape[2]};
+  return {out_shape};
 }
 
 std::vector<paddle::DataType> GatherIdxInferDtype(
     const paddle::DataType &logits_dtype,
     const paddle::DataType &gather_id_dtype) {
-    return {logits_dtype};
+  return {logits_dtype};
 }
 
 PD_BUILD_STATIC_OP(gather_idx)
diff --git a/custom_ops/gpu_ops/get_data_ptr_ipc.cu b/custom_ops/gpu_ops/get_data_ptr_ipc.cu
index 9004d76a2a..1b1eddc5ae 100644
--- a/custom_ops/gpu_ops/get_data_ptr_ipc.cu
+++ b/custom_ops/gpu_ops/get_data_ptr_ipc.cu
@@ -17,42 +17,41 @@
 
 namespace {
 int sharedMemoryOpen2(const char *name, size_t sz, sharedMemoryInfo *info) {
-    info->size = sz;
-    info->shmFd = shm_open(name, O_RDWR, 0777);
-    if (info->shmFd < 0) {
-        return errno;
-    }
+  info->size = sz;
+  info->shmFd = shm_open(name, O_RDWR, 0777);
+  if (info->shmFd < 0) {
+    return errno;
+  }
 
-    info->addr =
-        mmap(0, sz, PROT_READ | PROT_WRITE, MAP_SHARED, info->shmFd, 0);
-    if (info->addr == NULL) {
-        return errno;
-    }
+  info->addr = mmap(0, sz, PROT_READ | PROT_WRITE, MAP_SHARED, info->shmFd, 0);
+  if (info->addr == NULL) {
+    return errno;
+  }
 
-    return 0;
+  return 0;
 }
 }  // namespace
 
 std::vector<paddle::Tensor> GetDataPtrIpc(const paddle::Tensor &tmp_input,
                                           const std::string &shm_name) {
-    auto out_data_ptr_tensor =
-        paddle::full({1}, 0, paddle::DataType::INT64, paddle::CPUPlace());
-    auto out_data_ptr_tensor_ptr = out_data_ptr_tensor.data<int64_t>();
-    volatile shmStruct *shm = NULL;
-    sharedMemoryInfo info;
-    if (sharedMemoryOpen2(shm_name.c_str(), sizeof(shmStruct), &info) != 0) {
-        printf("Failed to create shared memory slab\n");
-        printf("Func GetDataPtrIpc. Shm_name: %s\n", shm_name.c_str());
-        exit(EXIT_FAILURE);
-    }
-    shm = (volatile shmStruct *)info.addr;
-    void *ptr = nullptr;
-    checkCudaErrors(cudaIpcOpenMemHandle(&ptr,
-                                         *(cudaIpcMemHandle_t *)&shm->memHandle,
-                                         cudaIpcMemLazyEnablePeerAccess));
+  auto out_data_ptr_tensor =
+      paddle::full({1}, 0, paddle::DataType::INT64, paddle::CPUPlace());
+  auto out_data_ptr_tensor_ptr = out_data_ptr_tensor.data<int64_t>();
+  volatile shmStruct *shm = NULL;
+  sharedMemoryInfo info;
+  if (sharedMemoryOpen2(shm_name.c_str(), sizeof(shmStruct), &info) != 0) {
+    printf("Failed to create shared memory slab\n");
+    printf("Func GetDataPtrIpc. Shm_name: %s\n", shm_name.c_str());
+    exit(EXIT_FAILURE);
+  }
+  shm = (volatile shmStruct *)info.addr;
+  void *ptr = nullptr;
+  checkCudaErrors(cudaIpcOpenMemHandle(&ptr,
+                                       *(cudaIpcMemHandle_t *)&shm->memHandle,
+                                       cudaIpcMemLazyEnablePeerAccess));
 
-    out_data_ptr_tensor_ptr[0] = reinterpret_cast<int64_t>(ptr);
-    return {out_data_ptr_tensor};
+  out_data_ptr_tensor_ptr[0] = reinterpret_cast<int64_t>(ptr);
+  return {out_data_ptr_tensor};
 }
 
 PD_BUILD_STATIC_OP(get_data_ptr_ipc)
diff --git a/custom_ops/gpu_ops/get_mm_split_fuse.cc b/custom_ops/gpu_ops/get_mm_split_fuse.cc
index 81531ba95e..7cad249141 100644
--- a/custom_ops/gpu_ops/get_mm_split_fuse.cc
+++ b/custom_ops/gpu_ops/get_mm_split_fuse.cc
@@ -19,126 +19,158 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-std::vector<paddle::Tensor> GetMmSplitFuse(const paddle::Tensor& task_input_ids,
-                            const paddle::Tensor& task_image_type_ids,
-                            const paddle::Tensor& task_input_ids_image_token_count,
-                            const paddle::Tensor& grid_thw,
-                            int64_t image_token_id,
-                            int64_t img_total,
-                            int batch_idx,
-                            int seq_lens_origin,
-                            int split_fuse_img_size,
-                            int split_fuse_text_size,
-                            int max_chunk_token_size) {
-    // All tensor in cpu
-    auto input_ids_cpu = task_input_ids.data<int64_t>();
-    auto image_type_ids_cpu = task_image_type_ids.data<int>();
-    auto task_input_ids_image_token_count_cpu = task_input_ids_image_token_count.data<int>();
-    auto grid_thw_cpu = grid_thw.data<int64_t>();
-    int chunk_token_count = 0;
-    int chunk_image_cout = 0;
-    int idx = 0;
-    std::vector<int> image_chunk_selections_vector;
-    std::vector<int> split_fuse_cur_seq_lens_vector;
-    std::vector<int> split_fuse_cur_idx_vector;
-    split_fuse_cur_idx_vector.emplace_back(0);
-    int image_idx = 0;
-    int last_ib = 0;
-    // 打表参数, mp记录可划分chunk的位置
-    std::map<int, int> mp;
-    int st_idx = 0;
-    int last_st_ib = 0;
-    while (st_idx < seq_lens_origin) {
-        // 1. 当前st_idx为文本，找到文本末尾
-        if (input_ids_cpu[st_idx] != image_token_id) {
-            do {
-                st_idx ++;
-            } while (st_idx < seq_lens_origin && input_ids_cpu[st_idx] != image_token_id);
-            mp[st_idx] = 1; // 记录划分chunk的末尾位置，此处为文本的末位+1
-        } else { // 2. 当前st_idx为多模，根据多模token的长度找到末尾
-            int ib = last_st_ib;
-            int cur_st_len = 0;
-            int token_times = 4;
-            cur_st_len = (grid_thw_cpu[ib * 3 + 1] * grid_thw_cpu[ib * 3 + 2]) / token_times;
-            mp[st_idx + cur_st_len] = 1;
-            last_st_ib = ++ib;
-            st_idx += cur_st_len;
-        }
+std::vector<paddle::Tensor> GetMmSplitFuse(
+    const paddle::Tensor& task_input_ids,
+    const paddle::Tensor& task_image_type_ids,
+    const paddle::Tensor& task_input_ids_image_token_count,
+    const paddle::Tensor& grid_thw,
+    int64_t image_token_id,
+    int64_t img_total,
+    int batch_idx,
+    int seq_lens_origin,
+    int split_fuse_img_size,
+    int split_fuse_text_size,
+    int max_chunk_token_size) {
+  // All tensor in cpu
+  auto input_ids_cpu = task_input_ids.data<int64_t>();
+  auto image_type_ids_cpu = task_image_type_ids.data<int>();
+  auto task_input_ids_image_token_count_cpu =
+      task_input_ids_image_token_count.data<int>();
+  auto grid_thw_cpu = grid_thw.data<int64_t>();
+  int chunk_token_count = 0;
+  int chunk_image_cout = 0;
+  int idx = 0;
+  std::vector<int> image_chunk_selections_vector;
+  std::vector<int> split_fuse_cur_seq_lens_vector;
+  std::vector<int> split_fuse_cur_idx_vector;
+  split_fuse_cur_idx_vector.emplace_back(0);
+  int image_idx = 0;
+  int last_ib = 0;
+  // 打表参数, mp记录可划分chunk的位置
+  std::map<int, int> mp;
+  int st_idx = 0;
+  int last_st_ib = 0;
+  while (st_idx < seq_lens_origin) {
+    // 1. 当前st_idx为文本，找到文本末尾
+    if (input_ids_cpu[st_idx] != image_token_id) {
+      do {
+        st_idx++;
+      } while (st_idx < seq_lens_origin &&
+               input_ids_cpu[st_idx] != image_token_id);
+      mp[st_idx] = 1;  // 记录划分chunk的末尾位置，此处为文本的末位+1
+    } else {  // 2. 当前st_idx为多模，根据多模token的长度找到末尾
+      int ib = last_st_ib;
+      int cur_st_len = 0;
+      int token_times = 4;
+      cur_st_len =
+          (grid_thw_cpu[ib * 3 + 1] * grid_thw_cpu[ib * 3 + 2]) / token_times;
+      mp[st_idx + cur_st_len] = 1;
+      last_st_ib = ++ib;
+      st_idx += cur_st_len;
     }
+  }
 
-    while (idx < seq_lens_origin) {
-        idx = idx + split_fuse_text_size;
-        if (idx >= seq_lens_origin) {
-            // idx 超过最大seq_len，应该包含n个图片和文本
-            idx = seq_lens_origin;
-            int last_idx = split_fuse_cur_idx_vector.back();
-            int chunk_image_token_number = task_input_ids_image_token_count_cpu[idx] - task_input_ids_image_token_count_cpu[last_idx];
-            int chunk_image_number = 0;
-            int cur_img_len = 0;
-            int ib = last_ib;
-            while (ib < img_total && cur_img_len < chunk_image_token_number){
-                int token_times = 4;
-                cur_img_len += (grid_thw_cpu[ib * 3 + 1] * grid_thw_cpu[ib * 3 + 2]) / token_times;
-                ib ++;
-                chunk_image_number ++;
-            }
-            image_chunk_selections_vector.emplace_back(chunk_image_number);
-            split_fuse_cur_seq_lens_vector.emplace_back(idx - last_idx);
-            split_fuse_cur_idx_vector.emplace_back(idx);
-            continue;
-        }
-        // text
-        if (input_ids_cpu[idx-1] != image_token_id) {
-            // case1. 如果切到text, 直接分chunk
-            int last_idx = split_fuse_cur_idx_vector.back();
-            int chunk_image_token_number = task_input_ids_image_token_count_cpu[idx] - task_input_ids_image_token_count_cpu[last_idx];
-            int chunk_image_number = 0;
-            int cur_img_len = 0;
-            int ib = last_ib;
-            while (ib < img_total && cur_img_len < chunk_image_token_number) {
-                int token_times = 4;
-                cur_img_len += (grid_thw_cpu[ib * 3 + 1] * grid_thw_cpu[ib * 3 + 2]) / token_times;
-                ib ++;
-                chunk_image_number ++;
-            }
-            image_chunk_selections_vector.emplace_back(chunk_image_number);
-            split_fuse_cur_seq_lens_vector.emplace_back(idx - last_idx); // split_fuse_text_size
-            last_ib = ib; // last_ib记录遍历到第几张图
-            split_fuse_cur_idx_vector.emplace_back(idx);
-            continue;
-        } else {
-            // case2. 如果切到图片，从当前图片往后找。往前找会出现边界问题
-            // case2.1 如果split_size = img_token_num, mp[idx]==1, 直接按当前idx切分chunk
-            while (idx < seq_lens_origin && mp[idx] != 1) {
-                idx++;
-            } // idx指向切分chunk的位置
-            int last_idx = split_fuse_cur_idx_vector.back();
-            int chunk_image_token_number = task_input_ids_image_token_count_cpu[idx] - task_input_ids_image_token_count_cpu[last_idx];
-            int chunk_image_number = 0;
-            int cur_img_len = 0;
-            int ib = last_ib;
-            while (ib < img_total && cur_img_len < chunk_image_token_number) {
-                int token_times = 4;
-                cur_img_len += (grid_thw_cpu[ib * 3 + 1] * grid_thw_cpu[ib * 3 + 2]) / token_times;
-                ib ++;
-                chunk_image_number ++;
-            }
-            image_chunk_selections_vector.emplace_back(chunk_image_number);
-            split_fuse_cur_seq_lens_vector.emplace_back(idx - last_idx);
-            split_fuse_cur_idx_vector.emplace_back(idx);
-            last_ib = ib;
-            continue;
-        }
+  while (idx < seq_lens_origin) {
+    idx = idx + split_fuse_text_size;
+    if (idx >= seq_lens_origin) {
+      // idx 超过最大seq_len，应该包含n个图片和文本
+      idx = seq_lens_origin;
+      int last_idx = split_fuse_cur_idx_vector.back();
+      int chunk_image_token_number =
+          task_input_ids_image_token_count_cpu[idx] -
+          task_input_ids_image_token_count_cpu[last_idx];
+      int chunk_image_number = 0;
+      int cur_img_len = 0;
+      int ib = last_ib;
+      while (ib < img_total && cur_img_len < chunk_image_token_number) {
+        int token_times = 4;
+        cur_img_len +=
+            (grid_thw_cpu[ib * 3 + 1] * grid_thw_cpu[ib * 3 + 2]) / token_times;
+        ib++;
+        chunk_image_number++;
+      }
+      image_chunk_selections_vector.emplace_back(chunk_image_number);
+      split_fuse_cur_seq_lens_vector.emplace_back(idx - last_idx);
+      split_fuse_cur_idx_vector.emplace_back(idx);
+      continue;
     }
-    auto image_chunk_selections_out_cpu = paddle::from_blob(image_chunk_selections_vector.data(), {image_chunk_selections_vector.size()}, task_image_type_ids.dtype());
-    auto split_fuse_cur_seq_lens_out_cpu = paddle::from_blob(split_fuse_cur_seq_lens_vector.data(), {split_fuse_cur_seq_lens_vector.size()}, task_image_type_ids.dtype());
-    auto image_chunk_selections_out = paddle::experimental::copy_to(image_chunk_selections_out_cpu, task_image_type_ids.place(), false);
-    auto split_fuse_cur_seq_lens_out = paddle::experimental::copy_to(split_fuse_cur_seq_lens_out_cpu, task_image_type_ids.place(), false);
-    return {image_chunk_selections_out, split_fuse_cur_seq_lens_out};
+    // text
+    if (input_ids_cpu[idx - 1] != image_token_id) {
+      // case1. 如果切到text, 直接分chunk
+      int last_idx = split_fuse_cur_idx_vector.back();
+      int chunk_image_token_number =
+          task_input_ids_image_token_count_cpu[idx] -
+          task_input_ids_image_token_count_cpu[last_idx];
+      int chunk_image_number = 0;
+      int cur_img_len = 0;
+      int ib = last_ib;
+      while (ib < img_total && cur_img_len < chunk_image_token_number) {
+        int token_times = 4;
+        cur_img_len +=
+            (grid_thw_cpu[ib * 3 + 1] * grid_thw_cpu[ib * 3 + 2]) / token_times;
+        ib++;
+        chunk_image_number++;
+      }
+      image_chunk_selections_vector.emplace_back(chunk_image_number);
+      split_fuse_cur_seq_lens_vector.emplace_back(
+          idx - last_idx);  // split_fuse_text_size
+      last_ib = ib;         // last_ib记录遍历到第几张图
+      split_fuse_cur_idx_vector.emplace_back(idx);
+      continue;
+    } else {
+      // case2. 如果切到图片，从当前图片往后找。往前找会出现边界问题
+      // case2.1 如果split_size = img_token_num, mp[idx]==1,
+      // 直接按当前idx切分chunk
+      while (idx < seq_lens_origin && mp[idx] != 1) {
+        idx++;
+      }  // idx指向切分chunk的位置
+      int last_idx = split_fuse_cur_idx_vector.back();
+      int chunk_image_token_number =
+          task_input_ids_image_token_count_cpu[idx] -
+          task_input_ids_image_token_count_cpu[last_idx];
+      int chunk_image_number = 0;
+      int cur_img_len = 0;
+      int ib = last_ib;
+      while (ib < img_total && cur_img_len < chunk_image_token_number) {
+        int token_times = 4;
+        cur_img_len +=
+            (grid_thw_cpu[ib * 3 + 1] * grid_thw_cpu[ib * 3 + 2]) / token_times;
+        ib++;
+        chunk_image_number++;
+      }
+      image_chunk_selections_vector.emplace_back(chunk_image_number);
+      split_fuse_cur_seq_lens_vector.emplace_back(idx - last_idx);
+      split_fuse_cur_idx_vector.emplace_back(idx);
+      last_ib = ib;
+      continue;
+    }
+  }
+  auto image_chunk_selections_out_cpu =
+      paddle::from_blob(image_chunk_selections_vector.data(),
+                        {image_chunk_selections_vector.size()},
+                        task_image_type_ids.dtype());
+  auto split_fuse_cur_seq_lens_out_cpu =
+      paddle::from_blob(split_fuse_cur_seq_lens_vector.data(),
+                        {split_fuse_cur_seq_lens_vector.size()},
+                        task_image_type_ids.dtype());
+  auto image_chunk_selections_out = paddle::experimental::copy_to(
+      image_chunk_selections_out_cpu, task_image_type_ids.place(), false);
+  auto split_fuse_cur_seq_lens_out = paddle::experimental::copy_to(
+      split_fuse_cur_seq_lens_out_cpu, task_image_type_ids.place(), false);
+  return {image_chunk_selections_out, split_fuse_cur_seq_lens_out};
 }
 
 PD_BUILD_STATIC_OP(get_mm_split_fuse)
-    .Inputs({"task_input_ids", "task_image_type_ids", "task_input_ids_image_token_count", "grid_thw"})
-    .Attrs({"image_token_id: int64_t", "img_total: int64_t", "batch_idx: int", "seq_lens_origin: int", "split_fuse_img_size: int", "split_fuse_text_size: int", "max_chunk_token_size: int"})
+    .Inputs({"task_input_ids",
+             "task_image_type_ids",
+             "task_input_ids_image_token_count",
+             "grid_thw"})
+    .Attrs({"image_token_id: int64_t",
+            "img_total: int64_t",
+            "batch_idx: int",
+            "seq_lens_origin: int",
+            "split_fuse_img_size: int",
+            "split_fuse_text_size: int",
+            "max_chunk_token_size: int"})
     .Outputs({"image_chunk_selections", "split_fuse_cur_seq_lens"})
     .SetKernelFn(PD_KERNEL(GetMmSplitFuse));
diff --git a/custom_ops/gpu_ops/get_output_msg_with_topk.cc b/custom_ops/gpu_ops/get_output_msg_with_topk.cc
index 4d6b5f56bb..8972815d2c 100644
--- a/custom_ops/gpu_ops/get_output_msg_with_topk.cc
+++ b/custom_ops/gpu_ops/get_output_msg_with_topk.cc
@@ -27,10 +27,10 @@
 #define K 20
 
 struct msgdata {
-    long mtype;
-    int mtext[MAX_BSZ * (K + 1) + 2];  // stop_flag, bsz, tokens
-    float mtext_f[MAX_BSZ * (K + 1)];  // score
-    int mtext_ranks[MAX_BSZ];  // ranks
+  long mtype;
+  int mtext[MAX_BSZ * (K + 1) + 2];  // stop_flag, bsz, tokens
+  float mtext_f[MAX_BSZ * (K + 1)];  // score
+  int mtext_ranks[MAX_BSZ];          // ranks
 };
 
 void GetOutputTopK(const paddle::Tensor& x,
@@ -39,72 +39,74 @@ void GetOutputTopK(const paddle::Tensor& x,
                    int k,
                    int64_t rank_id,
                    bool wait_flag) {
+  static struct msgdata msg_rcv;
+  int msg_queue_id = 1;
 
-    static struct msgdata msg_rcv;
-    int msg_queue_id = 1;
-
-    if (const char* inference_msg_queue_id_env_p =
-            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
-        std::string inference_msg_queue_id_env_str(
-            inference_msg_queue_id_env_p);
-        int inference_msg_queue_id_from_env =
-            std::stoi(inference_msg_queue_id_env_str);
+  if (const char* inference_msg_queue_id_env_p =
+          std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+    std::string inference_msg_queue_id_env_str(inference_msg_queue_id_env_p);
+    int inference_msg_queue_id_from_env =
+        std::stoi(inference_msg_queue_id_env_str);
 #ifdef GET_OUTPUT_DEBUG
-        std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
-                  << inference_msg_queue_id_from_env << std::endl;
+    std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+              << inference_msg_queue_id_from_env << std::endl;
 #endif
-        msg_queue_id = inference_msg_queue_id_from_env;
-    }
-    static key_t key = ftok("/dev/shm", msg_queue_id);
+    msg_queue_id = inference_msg_queue_id_from_env;
+  }
+  static key_t key = ftok("/dev/shm", msg_queue_id);
 
-    static int msgid = msgget(key, IPC_CREAT | 0666);
+  static int msgid = msgget(key, IPC_CREAT | 0666);
 #ifdef GET_OUTPUT_DEBUG
-    std::cout << "get_output_key: " << key << std::endl;
-    std::cout << "get_output msgid: " << msgid << std::endl;
+  std::cout << "get_output_key: " << key << std::endl;
+  std::cout << "get_output msgid: " << msgid << std::endl;
 #endif
 
-    int64_t* out_data = const_cast<int64_t*>(x.data<int64_t>());
-    float* scores_data = const_cast<float*>(scores.data<float>());
-    int64_t* ranks_data = const_cast<int64_t*>(ranks.data<int64_t>());
-    int ret = -1;
-    if (!wait_flag) {
-        ret = msgrcv(msgid,
-                     &msg_rcv,
-                     (MAX_BSZ * (K + 1) + 2) * 4 + MAX_BSZ * (K + 1) * 4 + MAX_BSZ * 4,
-                     0,
-                     IPC_NOWAIT);
-    } else {
-        ret = msgrcv(msgid,
-                     &msg_rcv,
-                     (MAX_BSZ * (K + 1) + 2) * 4 + MAX_BSZ * (K + 1) * 4 +  MAX_BSZ * 4,
-                     0,
-                     0);
-    }
-    if (ret == -1) {
-        // read none
-        out_data[0] = -2;
-        out_data[1] = 0;
-        return;
-    }
-
-    int bsz = msg_rcv.mtext[1];
-    out_data[0] = (int64_t)msg_rcv.mtext[0];
-    out_data[1] = (int64_t)msg_rcv.mtext[1];
-
-    for (int i = 0; i < bsz; i++) {
-        for (int j = 0; j < k + 1; j++) {
-            const int64_t offset = i * (K + 1) + j;
-            out_data[offset + 2] = (int64_t)msg_rcv.mtext[offset + 2];
-            scores_data[offset] = msg_rcv.mtext_f[offset];
-        }
-        ranks_data[i] = (int64_t)msg_rcv.mtext_ranks[i];
-    }
+  int64_t* out_data = const_cast<int64_t*>(x.data<int64_t>());
+  float* scores_data = const_cast<float*>(scores.data<float>());
+  int64_t* ranks_data = const_cast<int64_t*>(ranks.data<int64_t>());
+  int ret = -1;
+  if (!wait_flag) {
+    ret = msgrcv(
+        msgid,
+        &msg_rcv,
+        (MAX_BSZ * (K + 1) + 2) * 4 + MAX_BSZ * (K + 1) * 4 + MAX_BSZ * 4,
+        0,
+        IPC_NOWAIT);
+  } else {
+    ret = msgrcv(
+        msgid,
+        &msg_rcv,
+        (MAX_BSZ * (K + 1) + 2) * 4 + MAX_BSZ * (K + 1) * 4 + MAX_BSZ * 4,
+        0,
+        0);
+  }
+  if (ret == -1) {
+    // read none
+    out_data[0] = -2;
+    out_data[1] = 0;
     return;
+  }
+
+  int bsz = msg_rcv.mtext[1];
+  out_data[0] = (int64_t)msg_rcv.mtext[0];
+  out_data[1] = (int64_t)msg_rcv.mtext[1];
+
+  for (int i = 0; i < bsz; i++) {
+    for (int j = 0; j < k + 1; j++) {
+      const int64_t offset = i * (K + 1) + j;
+      out_data[offset + 2] = (int64_t)msg_rcv.mtext[offset + 2];
+      scores_data[offset] = msg_rcv.mtext_f[offset];
+    }
+    ranks_data[i] = (int64_t)msg_rcv.mtext_ranks[i];
+  }
+  return;
 }
 
 PD_BUILD_STATIC_OP(get_output_topk)
     .Inputs({"x", "scores", "ranks"})
     .Attrs({"k: int", "rank_id: int64_t", "wait_flag: bool"})
     .Outputs({"x_out", "scores_out", "ranks_out"})
-    .SetInplaceMap({{"x", "x_out"}, {"scores", "scores_out"}, {"ranks", "ranks_out"}})
+    .SetInplaceMap({{"x", "x_out"},
+                    {"scores", "scores_out"},
+                    {"ranks", "ranks_out"}})
     .SetKernelFn(PD_KERNEL(GetOutputTopK));
diff --git a/custom_ops/gpu_ops/get_padding_offset_system.cu b/custom_ops/gpu_ops/get_padding_offset_system.cu
index 841228c032..8157868dd2 100644
--- a/custom_ops/gpu_ops/get_padding_offset_system.cu
+++ b/custom_ops/gpu_ops/get_padding_offset_system.cu
@@ -30,30 +30,30 @@ __global__ void GetPaddingOffsetSystemKernel(int64_t* output_data,
                                              const int* seq_lens,
                                              const int* seq_lens_merged,
                                              const int max_seq_len) {
-    // get padding offset of each batch
-    const int bi = blockIdx.x;
-    const int real_bi = seq_mapping[bi];
-    const int ti = threadIdx.x;
-    const int cum_offset = bi == 0 ? 0 : cum_offsets[bi - 1];
+  // get padding offset of each batch
+  const int bi = blockIdx.x;
+  const int real_bi = seq_mapping[bi];
+  const int ti = threadIdx.x;
+  const int cum_offset = bi == 0 ? 0 : cum_offsets[bi - 1];
 
-    const int base_offset = bi * max_seq_len - cum_offset;
-    const int src_base_offset = real_bi * max_seq_len;
-    for (int i = ti; i < seq_lens[bi]; i += blockDim.x) {
-        const int offset_now = base_offset + i;
-        padding_offset[offset_now] = cum_offset;
-        output_data[offset_now] = input_data[src_base_offset + i];
-    }
+  const int base_offset = bi * max_seq_len - cum_offset;
+  const int src_base_offset = real_bi * max_seq_len;
+  for (int i = ti; i < seq_lens[bi]; i += blockDim.x) {
+    const int offset_now = base_offset + i;
+    padding_offset[offset_now] = cum_offset;
+    output_data[offset_now] = input_data[src_base_offset + i];
+  }
 
-    const int seq_len_merged = seq_lens_merged[bi];
-    const int cum_offset_merged = bi == 0 ? 0 : cum_offsets_merged[bi - 1];
-    const int merged_base_offset = bi * max_seq_len - cum_offset_merged;
-    for (int i = ti; i < seq_len_merged; i += blockDim.x) {
-        padding_offset_merged[merged_base_offset + i] = cum_offset_merged;
-    }
-    if (ti == 0) {
-        cum_offsets_out[bi] = cum_offset;
-        cum_offsets_out_merged[bi] = cum_offset_merged;
-    }
+  const int seq_len_merged = seq_lens_merged[bi];
+  const int cum_offset_merged = bi == 0 ? 0 : cum_offsets_merged[bi - 1];
+  const int merged_base_offset = bi * max_seq_len - cum_offset_merged;
+  for (int i = ti; i < seq_len_merged; i += blockDim.x) {
+    padding_offset_merged[merged_base_offset + i] = cum_offset_merged;
+  }
+  if (ti == 0) {
+    cum_offsets_out[bi] = cum_offset;
+    cum_offsets_out_merged[bi] = cum_offset_merged;
+  }
 }
 
 std::vector<paddle::Tensor> GetPaddingOffsetSystem(
@@ -64,41 +64,41 @@ std::vector<paddle::Tensor> GetPaddingOffsetSystem(
     const paddle::Tensor& seq_len,
     const paddle::Tensor& seq_len_merged,
     const paddle::Tensor& seq_mapping) {
-    auto cu_stream = input_ids.stream();
-    std::vector<int64_t> input_ids_shape = input_ids.shape();
-    const int bsz = seq_len.shape()[0];
-    const int seq_length = input_ids_shape[1];
-    auto cum_offsets_out = cum_offsets.copy_to(cum_offsets.place(), false);
-    auto cum_offsets_merged_out =
-        cum_offsets_merged.copy_to(cum_offsets_merged.place(), false);
-    auto cpu_token_num = token_num.copy_to(paddle::CPUPlace(), false);
+  auto cu_stream = input_ids.stream();
+  std::vector<int64_t> input_ids_shape = input_ids.shape();
+  const int bsz = seq_len.shape()[0];
+  const int seq_length = input_ids_shape[1];
+  auto cum_offsets_out = cum_offsets.copy_to(cum_offsets.place(), false);
+  auto cum_offsets_merged_out =
+      cum_offsets_merged.copy_to(cum_offsets_merged.place(), false);
+  auto cpu_token_num = token_num.copy_to(paddle::CPUPlace(), false);
 
-    const int token_num_data = cpu_token_num.data<int64_t>()[0];
-    auto x_remove_padding = paddle::full(
-        {token_num_data}, 0, paddle::DataType::INT64, input_ids.place());
-    auto padding_offset = paddle::full(
-        {token_num_data}, 0, paddle::DataType::INT32, input_ids.place());
-    auto padding_offset_merged = paddle::full(
-        {token_num_data}, 0, paddle::DataType::INT32, input_ids.place());
-    int blockSize = min((token_num_data + 32 - 1) / 32 * 32, 128);
-    GetPaddingOffsetSystemKernel<<<bsz, 1024, 0, cu_stream>>>(
-        x_remove_padding.data<int64_t>(),
-        padding_offset.data<int>(),
-        padding_offset_merged.data<int>(),
-        cum_offsets_out.data<int>(),
-        cum_offsets_merged_out.data<int>(),
-        input_ids.data<int64_t>(),
-        seq_mapping.data<int>(),
-        cum_offsets.data<int>(),
-        cum_offsets_merged.data<int>(),
-        seq_len.data<int>(),
-        seq_len_merged.data<int>(),
-        seq_length);
-    return {x_remove_padding,
-            cum_offsets_out,
-            cum_offsets_merged_out,
-            padding_offset,
-            padding_offset_merged};  // , enc_token_num, dec_token_num};
+  const int token_num_data = cpu_token_num.data<int64_t>()[0];
+  auto x_remove_padding = paddle::full(
+      {token_num_data}, 0, paddle::DataType::INT64, input_ids.place());
+  auto padding_offset = paddle::full(
+      {token_num_data}, 0, paddle::DataType::INT32, input_ids.place());
+  auto padding_offset_merged = paddle::full(
+      {token_num_data}, 0, paddle::DataType::INT32, input_ids.place());
+  int blockSize = min((token_num_data + 32 - 1) / 32 * 32, 128);
+  GetPaddingOffsetSystemKernel<<<bsz, 1024, 0, cu_stream>>>(
+      x_remove_padding.data<int64_t>(),
+      padding_offset.data<int>(),
+      padding_offset_merged.data<int>(),
+      cum_offsets_out.data<int>(),
+      cum_offsets_merged_out.data<int>(),
+      input_ids.data<int64_t>(),
+      seq_mapping.data<int>(),
+      cum_offsets.data<int>(),
+      cum_offsets_merged.data<int>(),
+      seq_len.data<int>(),
+      seq_len_merged.data<int>(),
+      seq_length);
+  return {x_remove_padding,
+          cum_offsets_out,
+          cum_offsets_merged_out,
+          padding_offset,
+          padding_offset_merged};  // , enc_token_num, dec_token_num};
 }
 
 std::vector<std::vector<int64_t>> GetPaddingOffsetSystemInferShape(
@@ -109,8 +109,8 @@ std::vector<std::vector<int64_t>> GetPaddingOffsetSystemInferShape(
     const std::vector<int64_t>& seq_len_shape,
     const std::vector<int64_t>& seq_len_merged_shape,
     const std::vector<int64_t>& seq_mapping_shape) {
-    int64_t bsz = seq_len_shape[0];
-    return {{-1}, {bsz}, {bsz}, {-1}, {-1}};
+  int64_t bsz = seq_len_shape[0];
+  return {{-1}, {bsz}, {bsz}, {-1}, {-1}};
 }
 
 std::vector<paddle::DataType> GetPaddingOffsetSystemInferDtype(
@@ -121,11 +121,11 @@ std::vector<paddle::DataType> GetPaddingOffsetSystemInferDtype(
     const paddle::DataType& seq_len_dtype,
     const paddle::DataType& seq_len_merged_dtype,
     const paddle::DataType& seq_mapping_dtype) {
-    return {input_ids_dtype,
-            seq_len_dtype,
-            seq_len_dtype,
-            seq_len_dtype,
-            seq_len_dtype};
+  return {input_ids_dtype,
+          seq_len_dtype,
+          seq_len_dtype,
+          seq_len_dtype,
+          seq_len_dtype};
 }
 
 PD_BUILD_STATIC_OP(get_padding_offset_system)
diff --git a/custom_ops/gpu_ops/get_position_ids_and_mask_encoder_batch.cu b/custom_ops/gpu_ops/get_position_ids_and_mask_encoder_batch.cu
index 9ddc1732e2..946c975407 100644
--- a/custom_ops/gpu_ops/get_position_ids_and_mask_encoder_batch.cu
+++ b/custom_ops/gpu_ops/get_position_ids_and_mask_encoder_batch.cu
@@ -15,7 +15,6 @@
 #include "helper.h"
 #include "paddle/extension.h"
 
-
 __global__ void GetPositionIdsAndMaskEncoderBatchKernel(
     const int* seq_lens_encoder,  // [bsz] 每个批次的 encoder 长度
     const int* seq_lens_decoder,  // [bsz] 每个批次的 decoder 长度
@@ -57,7 +56,6 @@ __global__ void GetPositionIdsAndMaskEncoderBatchKernel(
   }
 }
 
-
 void GetPositionIdsAndMaskEncoderBatch(
     const paddle::Tensor& seq_lens_encoder,
     const paddle::Tensor& seq_lens_decoder,
diff --git a/custom_ops/gpu_ops/glog/logging.h b/custom_ops/gpu_ops/glog/logging.h
index 7e1a45cc9a..c9fd82b0e7 100644
--- a/custom_ops/gpu_ops/glog/logging.h
+++ b/custom_ops/gpu_ops/glog/logging.h
@@ -16,13 +16,17 @@
 #include <iostream>
 
 /// @brief Defines a macro to disable log output.
-/// @details This macro returns a `NullStream` object, and all content passed via `<<` is ignored.
-///          It is designed to keep the code compilable without modifying log call statement.
+/// @details This macro returns a `NullStream` object, and all content passed
+/// via `<<` is ignored.
+///          It is designed to keep the code compilable without modifying log
+///          call statement.
 /// @param level The log level (ignored).
 class NullStream {
-public:
-    template <typename T>
-    NullStream& operator<<(const T&) { return *this; }
+ public:
+  template <typename T>
+  NullStream& operator<<(const T&) {
+    return *this;
+  }
 };
 
 #define VLOG(level) NullStream()
diff --git a/custom_ops/gpu_ops/init_signal_layerwise.cc b/custom_ops/gpu_ops/init_signal_layerwise.cc
index b55d2237d1..a47ea9cef7 100644
--- a/custom_ops/gpu_ops/init_signal_layerwise.cc
+++ b/custom_ops/gpu_ops/init_signal_layerwise.cc
@@ -21,28 +21,30 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-using cache_write_complete_signal_type = RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data;
+using cache_write_complete_signal_type =
+    RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data;
 
-paddle::Tensor InitSignalLayerwiseFunc(const paddle::Tensor& kv_signal_metadata, const int layer_id) {
-    auto kv_signal_metadata_out = kv_signal_metadata.copy_to(paddle::CPUPlace(), false);
-    kv_signal_metadata_out.data<int64_t>()[0] = static_cast<int64_t>(layer_id);
-    return kv_signal_metadata_out;
+paddle::Tensor InitSignalLayerwiseFunc(const paddle::Tensor& kv_signal_metadata,
+                                       const int layer_id) {
+  auto kv_signal_metadata_out =
+      kv_signal_metadata.copy_to(paddle::CPUPlace(), false);
+  kv_signal_metadata_out.data<int64_t>()[0] = static_cast<int64_t>(layer_id);
+  return kv_signal_metadata_out;
 }
 
-std::vector<paddle::Tensor> InitSignalLayerwise(const paddle::Tensor& kv_signal_metadata, const int layer_id) {
-    return {InitSignalLayerwiseFunc(kv_signal_metadata, layer_id)};
+std::vector<paddle::Tensor> InitSignalLayerwise(
+    const paddle::Tensor& kv_signal_metadata, const int layer_id) {
+  return {InitSignalLayerwiseFunc(kv_signal_metadata, layer_id)};
 }
 
 std::vector<std::vector<int64_t>> InitSignalLayerwiseShape(
-    const std::vector<int64_t>& kv_signal_metadata_shape,
-    const int layer_id) {
-    return {kv_signal_metadata_shape};
+    const std::vector<int64_t>& kv_signal_metadata_shape, const int layer_id) {
+  return {kv_signal_metadata_shape};
 }
 
 std::vector<paddle::DataType> InitSignalLayerwiseDtype(
-    const paddle::DataType& kv_signal_metadata_dtype,
-    const int layer_id) {
-    return {paddle::DataType::INT64};
+    const paddle::DataType& kv_signal_metadata_dtype, const int layer_id) {
+  return {paddle::DataType::INT64};
 }
 
 PD_BUILD_STATIC_OP(init_signal_layerwise)
diff --git a/custom_ops/gpu_ops/int8_gemm_with_cutlass/epilogue_tensor_op_int32.h b/custom_ops/gpu_ops/int8_gemm_with_cutlass/epilogue_tensor_op_int32.h
index 776ccce14d..37eca5eb70 100644
--- a/custom_ops/gpu_ops/int8_gemm_with_cutlass/epilogue_tensor_op_int32.h
+++ b/custom_ops/gpu_ops/int8_gemm_with_cutlass/epilogue_tensor_op_int32.h
@@ -98,17 +98,17 @@ template <typename ElementOutput,
           typename InstructionShape,
           typename ThreadMap>
 struct DefaultIteratorsTensorOp {
-    using WarpTileIterator =
-        cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
-                                                      InstructionShape,
-                                                      ElementAccumulator,
-                                                      layout::RowMajor>;
+  using WarpTileIterator =
+      cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
+                                                    InstructionShape,
+                                                    ElementAccumulator,
+                                                    layout::RowMajor>;
 
-    using SharedLoadIterator =
-        cutlass::epilogue::threadblock::SharedLoadIterator<ThreadMap,
-                                                           ElementAccumulator>;
+  using SharedLoadIterator =
+      cutlass::epilogue::threadblock::SharedLoadIterator<ThreadMap,
+                                                         ElementAccumulator>;
 
-    static int const kFragmentsPerIteration = 1;
+  static int const kFragmentsPerIteration = 1;
 };
 
 /// Partial specialization for half <= int32_t x 8 epilogues avoids shared
@@ -124,16 +124,16 @@ struct DefaultIteratorsTensorOp<cutlass::half_t,
                                 WarpShape,
                                 InstructionShape,
                                 ThreadMap> {
-    using WarpTileIterator =
-        cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
-                                                      InstructionShape,
-                                                      int32_t,
-                                                      layout::RowMajor>;
+  using WarpTileIterator =
+      cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
+                                                    InstructionShape,
+                                                    int32_t,
+                                                    layout::RowMajor>;
 
-    using SharedLoadIterator =
-        cutlass::epilogue::threadblock::SharedLoadIterator<ThreadMap, int32_t>;
+  using SharedLoadIterator =
+      cutlass::epilogue::threadblock::SharedLoadIterator<ThreadMap, int32_t>;
 
-    static int const kFragmentsPerIteration = 1;
+  static int const kFragmentsPerIteration = 1;
 };
 
 /// Partial specialization for bfloat16_t <= int32_t x 8 epilogues avoids shared
@@ -149,16 +149,16 @@ struct DefaultIteratorsTensorOp<cutlass::bfloat16_t,
                                 WarpShape,
                                 InstructionShape,
                                 ThreadMap> {
-    using WarpTileIterator =
-        cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
-                                                      InstructionShape,
-                                                      int32_t,
-                                                      layout::RowMajor>;
+  using WarpTileIterator =
+      cutlass::epilogue::warp::TileIteratorTensorOp<WarpShape,
+                                                    InstructionShape,
+                                                    int32_t,
+                                                    layout::RowMajor>;
 
-    using SharedLoadIterator =
-        cutlass::epilogue::threadblock::SharedLoadIterator<ThreadMap, int32_t>;
+  using SharedLoadIterator =
+      cutlass::epilogue::threadblock::SharedLoadIterator<ThreadMap, int32_t>;
 
-    static int const kFragmentsPerIteration = 1;
+  static int const kFragmentsPerIteration = 1;
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
@@ -174,162 +174,153 @@ struct DefaultIteratorsTensorOp<cutlass::bfloat16_t,
 template <typename ThreadMap_  ///< Thread map (concept: OutputTileThreadMap)
           >
 class SharedLoadIteratorMixed<ThreadMap_, int32_t, 32, 16, 8, 8> {
-    public:
-    using ThreadMap = ThreadMap_;
-    using Shape = typename ThreadMap::Shape;
+ public:
+  using ThreadMap = ThreadMap_;
+  using Shape = typename ThreadMap::Shape;
 
-    using Element = int32_t;
+  using Element = int32_t;
 
-    using Layout = layout::RowMajor;
-    using TensorRef = TensorRef<Element, Layout>;
-    using ConstTensorRef = typename TensorRef::ConstTensorRef;
+  using Layout = layout::RowMajor;
+  using TensorRef = TensorRef<Element, Layout>;
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
 
-    using Index = typename Layout::Index;
-    using LongIndex = typename Layout::LongIndex;
-    using TensorCoord = MatrixCoord;
+  using Index = typename Layout::Index;
+  using LongIndex = typename Layout::LongIndex;
+  using TensorCoord = MatrixCoord;
 
-    static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
+  static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
 
-    static int const kAlignment =
-        ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value / 8;
+  static int const kAlignment =
+      ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value / 8;
 
-    static int const kThreads = ThreadMap::kThreads;
+  static int const kThreads = ThreadMap::kThreads;
 
-    /// Fragment object
-    using Fragment =
-        Array<Element,
-              ThreadMap::Iterations::kColumn * ThreadMap::Iterations::kRow *
-                  ThreadMap::Iterations::kGroup *
-                  ThreadMap::Iterations::kCluster *
-                  ThreadMap::kElementsPerAccess>;
+  /// Fragment object
+  using Fragment =
+      Array<Element,
+            ThreadMap::Iterations::kColumn * ThreadMap::Iterations::kRow *
+                ThreadMap::Iterations::kGroup *
+                ThreadMap::Iterations::kCluster *
+                ThreadMap::kElementsPerAccess>;
 
-    /// Memory access size
-    using AccessType =
-        AlignedArray<Element, ThreadMap::kElementsPerAccess, kAlignment>;
+  /// Memory access size
+  using AccessType =
+      AlignedArray<Element, ThreadMap::kElementsPerAccess, kAlignment>;
 
-    /// Vector type used for SMEM loads
-    using LoadType = AlignedArray<Element,
-                                  const_min(128 / sizeof_bits<Element>::value,
-                                            ThreadMap::kElementsPerAccess),
-                                  const_min(16, kAlignment)>;
+  /// Vector type used for SMEM loads
+  using LoadType = AlignedArray<Element,
+                                const_min(128 / sizeof_bits<Element>::value,
+                                          ThreadMap::kElementsPerAccess),
+                                const_min(16, kAlignment)>;
 
-    static int const kLoadsPerAccess =
-        AccessType::kElements / LoadType::kElements;
+  static int const kLoadsPerAccess =
+      AccessType::kElements / LoadType::kElements;
 
-    private:
-    //
-    // Data members
-    //
+ private:
+  //
+  // Data members
+  //
 
-    /// Byte-level pointer
-    LoadType const *pointers_[kLoadsPerAccess];
+  /// Byte-level pointer
+  LoadType const *pointers_[kLoadsPerAccess];
 
-    /// Stride along adjacent rows in units of LoadType
-    int stride_;
+  /// Stride along adjacent rows in units of LoadType
+  int stride_;
 
-    public:
-    //
-    // Methods
-    //
+ public:
+  //
+  // Methods
+  //
 
-    /// Constructor
-    CUTLASS_DEVICE
-    SharedLoadIteratorMixed(TensorRef ref, int thread_idx)
-        : stride_((ref.stride(0) / LoadType::kElements)) {
-        TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
+  /// Constructor
+  CUTLASS_DEVICE
+  SharedLoadIteratorMixed(TensorRef ref, int thread_idx)
+      : stride_((ref.stride(0) / LoadType::kElements)) {
+    TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
 
-        // Initialize pointers
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < kLoadsPerAccess; ++i) {
-            pointers_[i] = reinterpret_cast<LoadType const *>(ref.data());
+    // Initialize pointers
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] = reinterpret_cast<LoadType const *>(ref.data());
 
-            int col_idx =
-                (thread_offset.column() / kElementsPerAccess) * kLoadsPerAccess;
-            int bank_offset =
-                (col_idx * static_cast<int>(sizeof(LoadType)) / 128) %
-                kLoadsPerAccess;
+      int col_idx =
+          (thread_offset.column() / kElementsPerAccess) * kLoadsPerAccess;
+      int bank_offset = (col_idx * static_cast<int>(sizeof(LoadType)) / 128) %
+                        kLoadsPerAccess;
 
-            col_idx += (bank_offset + i) % kLoadsPerAccess;
+      col_idx += (bank_offset + i) % kLoadsPerAccess;
 
-            pointers_[i] += thread_offset.row() * stride_ + col_idx;
-        }
+      pointers_[i] += thread_offset.row() * stride_ + col_idx;
     }
+  }
 
-    /// Adds a pointer offset in units of Element
-    CUTLASS_HOST_DEVICE
-    void add_pointer_offset(LongIndex pointer_offset) {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < kLoadsPerAccess; ++i) {
-            pointers_[i] += pointer_offset / LoadType::kElements;
-        }
+  /// Adds a pointer offset in units of Element
+  CUTLASS_HOST_DEVICE
+  void add_pointer_offset(LongIndex pointer_offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += pointer_offset / LoadType::kElements;
     }
+  }
 
-    CUTLASS_DEVICE
-    void add_tile_offset(TensorCoord const &offset) {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < kLoadsPerAccess; ++i) {
-            pointers_[i] +=
-                offset.row() * Shape::kRow * stride_ +
-                offset.column() * Shape::kColumn / LoadType::kElements;
-        }
+  CUTLASS_DEVICE
+  void add_tile_offset(TensorCoord const &offset) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kLoadsPerAccess; ++i) {
+      pointers_[i] += offset.row() * Shape::kRow * stride_ +
+                      offset.column() * Shape::kColumn / LoadType::kElements;
     }
+  }
 
-    /// Loads a fragment from memory
-    CUTLASS_DEVICE
-    void load_with_pointer_offset(Fragment &frag,  // NOLINT
-                                  Index pointer_offset) const {
+  /// Loads a fragment from memory
+  CUTLASS_DEVICE
+  void load_with_pointer_offset(Fragment &frag,  // NOLINT
+                                Index pointer_offset) const {
+    CUTLASS_PRAGMA_UNROLL
+    for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
+         ++cluster) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
         CUTLASS_PRAGMA_UNROLL
-        for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
-             ++cluster) {
+        for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+          int row_ptr_offset = row * ThreadMap::Delta::kRow * stride_ +
+                               group * ThreadMap::Delta::kGroup * stride_ +
+                               cluster * ThreadMap::Delta::kCluster * stride_ +
+                               pointer_offset / LoadType::kElements;
+
+          int frag_row_idx =
+              (row + ThreadMap::Iterations::kRow *
+                         (group + ThreadMap::Iterations::kGroup * cluster));
+
+          LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
+
+          CUTLASS_PRAGMA_UNROLL
+          for (int column = 0; column < ThreadMap::Iterations::kColumn;
+               ++column) {
+            int frag_idx =
+                frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
             CUTLASS_PRAGMA_UNROLL
-            for (int group = 0; group < ThreadMap::Iterations::kGroup;
-                 ++group) {
-                CUTLASS_PRAGMA_UNROLL
-                for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
-                    int row_ptr_offset =
-                        row * ThreadMap::Delta::kRow * stride_ +
-                        group * ThreadMap::Delta::kGroup * stride_ +
-                        cluster * ThreadMap::Delta::kCluster * stride_ +
-                        pointer_offset / LoadType::kElements;
+            for (int v = 0; v < kLoadsPerAccess; ++v) {
+              int vector_idx = (column * ThreadMap::Delta::kColumn /
+                                kElementsPerAccess * kLoadsPerAccess);
 
-                    int frag_row_idx =
-                        (row +
-                         ThreadMap::Iterations::kRow *
-                             (group + ThreadMap::Iterations::kGroup * cluster));
+              LoadType const *memory_pointer = pointers_[v] + row_ptr_offset;
 
-                    LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
-
-                    CUTLASS_PRAGMA_UNROLL
-                    for (int column = 0;
-                         column < ThreadMap::Iterations::kColumn;
-                         ++column) {
-                        int frag_idx =
-                            frag_row_idx * ThreadMap::Iterations::kColumn +
-                            column;
-
-                        CUTLASS_PRAGMA_UNROLL
-                        for (int v = 0; v < kLoadsPerAccess; ++v) {
-                            int vector_idx =
-                                (column * ThreadMap::Delta::kColumn /
-                                 kElementsPerAccess * kLoadsPerAccess);
-
-                            LoadType const *memory_pointer =
-                                pointers_[v] + row_ptr_offset;
-
-                            frag_ptr[frag_idx * kLoadsPerAccess + v] =
-                                memory_pointer[vector_idx];
-                        }
-                    }
-                }
+              frag_ptr[frag_idx * kLoadsPerAccess + v] =
+                  memory_pointer[vector_idx];
             }
+          }
         }
+      }
     }
+  }
 
-    /// Loads a fragment
-    CUTLASS_DEVICE
-    void load(Fragment &frag) const {  // NOLINT
-        load_with_pointer_offset(frag, 0);
-    }
+  /// Loads a fragment
+  CUTLASS_DEVICE
+  void load(Fragment &frag) const {  // NOLINT
+    load_with_pointer_offset(frag, 0);
+  }
 };
 
 template <typename Shape_,
@@ -340,89 +331,88 @@ template <typename Shape_,
           bool ScatterD = false,
           typename PermuteDLayout = layout::NoPermute>
 struct DequantEpilogueTensorOp {
-    using Shape = Shape_;
-    using WarpMmaTensorOp = WarpMmaTensorOp_;
-    static int const kPartitionsK = PartitionsK;
-    using OutputOp = OutputOp_;
-    static int const kElementsPerAccess = ElementsPerAccess;
+  using Shape = Shape_;
+  using WarpMmaTensorOp = WarpMmaTensorOp_;
+  static int const kPartitionsK = PartitionsK;
+  using OutputOp = OutputOp_;
+  static int const kElementsPerAccess = ElementsPerAccess;
 
-    using ElementOutput = typename OutputOp::ElementOutput;
-    using LayoutC = typename WarpMmaTensorOp::LayoutC;
-    using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
+  using ElementOutput = typename OutputOp::ElementOutput;
+  using LayoutC = typename WarpMmaTensorOp::LayoutC;
+  using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
 
-    //
-    // Thread map
-    //
+  //
+  // Thread map
+  //
 
-    using OutputTileThreadMap =
-        typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp<
-            Shape,
-            typename WarpMmaTensorOp::Shape,
-            kPartitionsK,
-            ElementOutput,
-            kElementsPerAccess>::Type;
+  using OutputTileThreadMap =
+      typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp<
+          Shape,
+          typename WarpMmaTensorOp::Shape,
+          kPartitionsK,
+          ElementOutput,
+          kElementsPerAccess>::Type;
 
-    static bool const UseCUDAStore =
-        platform::is_same<ElementOutput, double>::value;
+  static bool const UseCUDAStore =
+      platform::is_same<ElementOutput, double>::value;
 
-    using OutputTileIterator =
-        cutlass::epilogue::threadblock::PredicatedTileIterator<
-            OutputTileThreadMap,
-            ElementOutput,
-            ScatterD,
-            PermuteDLayout,
-            UseCUDAStore>;
+  using OutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIterator<
+          OutputTileThreadMap,
+          ElementOutput,
+          ScatterD,
+          PermuteDLayout,
+          UseCUDAStore>;
 
-    using AccumulatorFragmentIterator = typename platform::conditional<
-        is_complex<ElementOutput>::value,
-        cutlass::epilogue::warp::FragmentIteratorComplexTensorOp<
-            typename WarpMmaTensorOp::Shape,
-            typename WarpMmaTensorOp::Policy::Operator::Shape,
-            typename WarpMmaTensorOp::Policy::Operator::ElementC,
-            typename WarpMmaTensorOp::Policy::Operator::FragmentC,
-            LayoutC>,
-        cutlass::epilogue::warp::FragmentIteratorTensorOp<
-            typename WarpMmaTensorOp::Shape,
-            typename WarpMmaTensorOp::Policy::Operator::Shape,
-            typename WarpMmaTensorOp::Policy::Operator::ElementC,
-            typename WarpMmaTensorOp::Policy::Operator::FragmentC,
-            LayoutC>>::type;
+  using AccumulatorFragmentIterator = typename platform::conditional<
+      is_complex<ElementOutput>::value,
+      cutlass::epilogue::warp::FragmentIteratorComplexTensorOp<
+          typename WarpMmaTensorOp::Shape,
+          typename WarpMmaTensorOp::Policy::Operator::Shape,
+          typename WarpMmaTensorOp::Policy::Operator::ElementC,
+          typename WarpMmaTensorOp::Policy::Operator::FragmentC,
+          LayoutC>,
+      cutlass::epilogue::warp::FragmentIteratorTensorOp<
+          typename WarpMmaTensorOp::Shape,
+          typename WarpMmaTensorOp::Policy::Operator::Shape,
+          typename WarpMmaTensorOp::Policy::Operator::ElementC,
+          typename WarpMmaTensorOp::Policy::Operator::FragmentC,
+          LayoutC>>::type;
 
-    /// Support several implementations depending on structure of epilogue
-    using DefaultIterators = detail::DefaultIteratorsTensorOp<
-        ElementOutput,
-        ElementAccumulator,
-        kElementsPerAccess,
-        Shape,
-        typename WarpMmaTensorOp::Shape,
-        typename WarpMmaTensorOp::Policy::Operator::Shape,
-        typename OutputTileThreadMap::CompactedThreadMap>;
+  /// Support several implementations depending on structure of epilogue
+  using DefaultIterators = detail::DefaultIteratorsTensorOp<
+      ElementOutput,
+      ElementAccumulator,
+      kElementsPerAccess,
+      Shape,
+      typename WarpMmaTensorOp::Shape,
+      typename WarpMmaTensorOp::Policy::Operator::Shape,
+      typename OutputTileThreadMap::CompactedThreadMap>;
 
-    using WarpTileIterator = typename DefaultIterators::WarpTileIterator;
-    using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator;
+  using WarpTileIterator = typename DefaultIterators::WarpTileIterator;
+  using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator;
 
-    /// Hard-coded padding elements added
-    using Padding =
-        cutlass::MatrixShape<0,
-                             64 / sizeof_bits<ElementAccumulator>::value * 4>;
+  /// Hard-coded padding elements added
+  using Padding =
+      cutlass::MatrixShape<0, 64 / sizeof_bits<ElementAccumulator>::value * 4>;
 
-    static int const kFragmentsPerIteration =
-        (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1);
+  static int const kFragmentsPerIteration =
+      (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1);
 
-    //
-    // Define the epilogue
-    //
-    using Epilogue =
-        cutlass::epilogue::threadblock::Epilogue<Shape,
-                                                 WarpMmaTensorOp,
-                                                 kPartitionsK,
-                                                 OutputTileIterator,
-                                                 AccumulatorFragmentIterator,
-                                                 WarpTileIterator,
-                                                 SharedLoadIterator,
-                                                 OutputOp,
-                                                 Padding,
-                                                 kFragmentsPerIteration>;
+  //
+  // Define the epilogue
+  //
+  using Epilogue =
+      cutlass::epilogue::threadblock::Epilogue<Shape,
+                                               WarpMmaTensorOp,
+                                               kPartitionsK,
+                                               OutputTileIterator,
+                                               AccumulatorFragmentIterator,
+                                               WarpTileIterator,
+                                               SharedLoadIterator,
+                                               OutputOp,
+                                               Padding,
+                                               kFragmentsPerIteration>;
 };
 
 }  // namespace threadblock
diff --git a/custom_ops/gpu_ops/int8_gemm_with_cutlass/gemm_dequant.cu b/custom_ops/gpu_ops/int8_gemm_with_cutlass/gemm_dequant.cu
index e40cfb181c..31472d6e85 100644
--- a/custom_ops/gpu_ops/int8_gemm_with_cutlass/gemm_dequant.cu
+++ b/custom_ops/gpu_ops/int8_gemm_with_cutlass/gemm_dequant.cu
@@ -16,13 +16,13 @@
 #include "cutlass_helper.h"
 
 template <typename Type, int CtaM, int CtaN, int Threads>
-__global__ void int8_sq(int8_t const* act,
-                          int8_t const* weight,
-                          float const* scale,
-                          Type* output,
-                          int m,
-                          int n,
-                          int k) {
+__global__ void int8_sq(int8_t const *act,
+                        int8_t const *weight,
+                        float const *scale,
+                        Type *output,
+                        int m,
+                        int n,
+                        int k) {
   using VecType = int4;
   static constexpr int kStepK = 128 / (8 * sizeof(int8_t));
   static constexpr int CtaK = kStepK * Threads;
@@ -42,20 +42,20 @@ __global__ void int8_sq(int8_t const* act,
   for (int idx_k = tid * kStepK; idx_k < k; idx_k += CtaK) {
 #pragma unroll
     for (int i = 0; i < CtaN; ++i) {
-      reinterpret_cast<VecType*>(tile_w)[i] =
-          reinterpret_cast<VecType const*>(weight + i * k + idx_k)[0];
+      reinterpret_cast<VecType *>(tile_w)[i] =
+          reinterpret_cast<VecType const *>(weight + i * k + idx_k)[0];
     }
 #pragma unroll
     for (int i = 0; i < CtaM; ++i) {
-      reinterpret_cast<VecType*>(tile_a)[0] =
-          reinterpret_cast<VecType const*>(act + i * k + idx_k)[0];
+      reinterpret_cast<VecType *>(tile_a)[0] =
+          reinterpret_cast<VecType const *>(act + i * k + idx_k)[0];
 #pragma unroll
       for (int j = 0; j < CtaN; ++j) {
 #pragma unroll
         for (int l = 0; l < kStepK; l += 4) {
           acc[i * CtaN + j] =
-              __dp4a(reinterpret_cast<int*>(tile_a + l)[0],
-                     reinterpret_cast<int*>(tile_w + j * kStepK + l)[0],
+              __dp4a(reinterpret_cast<int *>(tile_a + l)[0],
+                     reinterpret_cast<int *>(tile_w + j * kStepK + l)[0],
                      acc[i * CtaN + j]);
         }
       }
@@ -90,7 +90,8 @@ __global__ void int8_sq(int8_t const* act,
     for (int jj = 0; jj < kWarpNum; ++jj) {
       val += shmem[jj * CtaM * CtaN + ii];
     }
-    output[mid * n + nid] = static_cast<Type>(static_cast<float>(val)*(float)*(scale+nid));
+    output[mid * n + nid] =
+        static_cast<Type>(static_cast<float>(val) * (float)*(scale + nid));
   }
 }
 
@@ -99,15 +100,15 @@ template <typename InputType,
           int32_t TILE_M,
           int32_t TILE_N,
           int32_t BLOCK_SIZE>
-void int8_sq_kernel(GemmDequantParams const& params) {
+void int8_sq_kernel(GemmDequantParams const &params) {
   dim3 block(BLOCK_SIZE);
   dim3 grid(params.m / TILE_M, params.n / TILE_N);
   int8_sq<OutputType, TILE_M, TILE_N, BLOCK_SIZE>
       <<<grid, block, 0, params.stream>>>(
-          reinterpret_cast<InputType const*>(params.act),
-          reinterpret_cast<InputType const*>(params.weight),
-          reinterpret_cast<float const*>(params.scale),
-          reinterpret_cast<OutputType*>(params.output),
+          reinterpret_cast<InputType const *>(params.act),
+          reinterpret_cast<InputType const *>(params.weight),
+          reinterpret_cast<float const *>(params.scale),
+          reinterpret_cast<OutputType *>(params.output),
           params.m,
           params.n,
           params.k);
@@ -118,25 +119,24 @@ template <typename InputType,
           int TILE_M,
           int TILE_N,
           int BLOCK_SIZE>
-bool int8_sq_kernel_caller(GemmDequantParams const& params) {
+bool int8_sq_kernel_caller(GemmDequantParams const &params) {
   constexpr int cudaCoreGemmTemplateMaxM = 16;
   if (params.m == TILE_M) {
-    int8_sq_kernel<InputType, OutputType, TILE_M, TILE_N, BLOCK_SIZE>(
-        params);
+    int8_sq_kernel<InputType, OutputType, TILE_M, TILE_N, BLOCK_SIZE>(params);
     return true;
   }
   if constexpr (TILE_M < cudaCoreGemmTemplateMaxM) {
     return int8_sq_kernel_caller<InputType,
-                                      OutputType,
-                                      TILE_M + 1,
-                                      TILE_N,
-                                      BLOCK_SIZE>(params);
+                                 OutputType,
+                                 TILE_M + 1,
+                                 TILE_N,
+                                 BLOCK_SIZE>(params);
   }
   return false;
 }
 
 template <typename InputType, typename OutputType>
-bool int8_sq_kernel_launcher(GemmDequantParams const& params) {
+bool int8_sq_kernel_launcher(GemmDequantParams const &params) {
   return int8_sq_kernel_caller<InputType, OutputType, 1, 2, 256>(params);
 }
 
@@ -149,176 +149,171 @@ void RunGemmDequant(const int8_t *a,
                     int k,
                     int n,
                     cudaStream_t stream) {
-    using ElementA = int8_t;
-    using ElementB = int8_t;
-    using ElementC = typename CutlassDtypeTraits<D>::DataType;
-    using ElementCompute = int32_t;
-    using ElementD = ElementC;
+  using ElementA = int8_t;
+  using ElementB = int8_t;
+  using ElementC = typename CutlassDtypeTraits<D>::DataType;
+  using ElementCompute = int32_t;
+  using ElementD = ElementC;
 
-    using LayoutA = cutlass::layout::RowMajor;
-    using LayoutB = cutlass::layout::ColumnMajor;
+  using LayoutA = cutlass::layout::RowMajor;
+  using LayoutB = cutlass::layout::ColumnMajor;
 
-    using ThreadblockShape = cutlass::gemm::GemmShape<64, 64, 64>;
-    using WarpShape = cutlass::gemm::GemmShape<32, 32, 64>;
-    using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
+  using ThreadblockShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using WarpShape = cutlass::gemm::GemmShape<32, 32, 64>;
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
 
-    using OperatorClass = cutlass::arch::OpClassTensorOp;
-    using ArchTag = cutlass::arch::Sm80;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+  using ArchTag = cutlass::arch::Sm80;
 
-    static int const kStages = 5;
+  static int const kStages = 5;
 
-    /// Linear scaling operator
-    using EpilogueFunctorOp = cutlass::epilogue::thread::LinearCombination<
-        ElementC,
-        128 / cutlass::sizeof_bits<ElementC>::value,
-        ElementCompute,
-        ElementCompute>;
+  /// Linear scaling operator
+  using EpilogueFunctorOp = cutlass::epilogue::thread::LinearCombination<
+      ElementC,
+      128 / cutlass::sizeof_bits<ElementC>::value,
+      ElementCompute,
+      ElementCompute>;
 
-    using GemmDequantT = cutlass::GemmDequant<ElementA,
-                                              LayoutA,
-                                              ElementB,
-                                              LayoutB,
-                                              ElementC,
-                                              ElementCompute,
-                                              OperatorClass,
-                                              ArchTag,
-                                              ThreadblockShape,
-                                              WarpShape,
-                                              InstructionShape,
-                                              EpilogueFunctorOp,
-                                              kStages>;
+  using GemmDequantT = cutlass::GemmDequant<ElementA,
+                                            LayoutA,
+                                            ElementB,
+                                            LayoutB,
+                                            ElementC,
+                                            ElementCompute,
+                                            OperatorClass,
+                                            ArchTag,
+                                            ThreadblockShape,
+                                            WarpShape,
+                                            InstructionShape,
+                                            EpilogueFunctorOp,
+                                            kStages>;
 
-    using LayoutC = typename GemmDequantT::LayoutC;
+  using LayoutC = typename GemmDequantT::LayoutC;
 
-    int64_t lda = LayoutA::packed({m, k}).stride(0);
-    int64_t ldb = LayoutB::packed({k, n}).stride(0);
-    int64_t ldc = LayoutC::packed({m, n}).stride(0);
+  int64_t lda = LayoutA::packed({m, k}).stride(0);
+  int64_t ldb = LayoutB::packed({k, n}).stride(0);
+  int64_t ldc = LayoutC::packed({m, n}).stride(0);
 
-    cutlass::gemm::GemmCoord problem_size(m, n, k);
+  cutlass::gemm::GemmCoord problem_size(m, n, k);
 
-    typename CutlassDtypeTraits<D>::DataType *c_tmp = nullptr;
-    typename CutlassDtypeTraits<D>::DataType *d =
-        reinterpret_cast<typename CutlassDtypeTraits<D>::DataType *>(c);
+  typename CutlassDtypeTraits<D>::DataType *c_tmp = nullptr;
+  typename CutlassDtypeTraits<D>::DataType *d =
+      reinterpret_cast<typename CutlassDtypeTraits<D>::DataType *>(c);
 
-    typename GemmDequantT::TensorRefA ref_a(const_cast<int8_t *>(a), lda);
-    typename GemmDequantT::TensorRefB ref_b(const_cast<int8_t *>(b), ldb);
-    typename GemmDequantT::TensorRefC ref_c(c_tmp, ldc);
-    typename GemmDequantT::TensorRefC ref_d(d, ldc);
-    typename GemmDequantT::TensorRefScale ref_scale(
-        const_cast<float *>(dequant_scale), 0);
+  typename GemmDequantT::TensorRefA ref_a(const_cast<int8_t *>(a), lda);
+  typename GemmDequantT::TensorRefB ref_b(const_cast<int8_t *>(b), ldb);
+  typename GemmDequantT::TensorRefC ref_c(c_tmp, ldc);
+  typename GemmDequantT::TensorRefC ref_d(d, ldc);
+  typename GemmDequantT::TensorRefScale ref_scale(
+      const_cast<float *>(dequant_scale), 0);
 
-    typename GemmDequantT::Arguments args(
-        problem_size,
-        ref_a,
-        ref_b,
-        ref_c,
-        ref_d,
-        ref_scale,
-        {ElementCompute(1.0f), ElementCompute(0.0f)});
+  typename GemmDequantT::Arguments args(
+      problem_size,
+      ref_a,
+      ref_b,
+      ref_c,
+      ref_d,
+      ref_scale,
+      {ElementCompute(1.0f), ElementCompute(0.0f)});
 
-    GemmDequantT gemm;
-    // Initialize
-    auto status = gemm.initialize(args);
-    PADDLE_ENFORCE_EQ(
-        status,
-        cutlass::Status::kSuccess,
-        phi::errors::Fatal("cutlass GemmDequant initialize error"));
+  GemmDequantT gemm;
+  // Initialize
+  auto status = gemm.initialize(args);
+  PADDLE_ENFORCE_EQ(status,
+                    cutlass::Status::kSuccess,
+                    phi::errors::Fatal("cutlass GemmDequant initialize error"));
 
-    // Run
-    status = gemm(stream);
-    PADDLE_ENFORCE_EQ(status,
-                      cutlass::Status::kSuccess,
-                      phi::errors::Fatal("cutlass GemmDequant runtime error"));
+  // Run
+  status = gemm(stream);
+  PADDLE_ENFORCE_EQ(status,
+                    cutlass::Status::kSuccess,
+                    phi::errors::Fatal("cutlass GemmDequant runtime error"));
 }
 
 std::vector<paddle::Tensor> GemmDequant(const paddle::Tensor &x,
                                         const paddle::Tensor &y,
                                         const paddle::Tensor &scale,
                                         const std::string &out_dtype) {
-    std::vector<int64_t> x_dims = x.shape(), y_dims = y.shape();
-    PADDLE_ENFORCE_EQ(
-        x_dims[x_dims.size() - 1],
-        y_dims[y_dims.size() - 1],
-        phi::errors::InvalidArgument(
-            "The last dimension of x and y should be equal. But "
-            "received x[%d] != y[%d].",
-            "Ensure that x is not transposed and y is transposed.",
-            x_dims[x_dims.size() - 1],
-            y_dims[y_dims.size() - 1]));
-    int64_t m = x_dims[x_dims.size() - 2];
-    int64_t k = x_dims[x_dims.size() - 1];
-    int64_t n = y_dims[y_dims.size() - 2];
-    if(m <= 4)
-    {
-        if (out_dtype == "bfloat16") {
-            paddle::Tensor out =
-                    paddle::empty({m, n}, paddle::DataType::BFLOAT16, x.place());
-            GemmDequantParams params = {
-                reinterpret_cast<const void*>(x.data<int8_t>()),
-                reinterpret_cast<const void*>(y.data<int8_t>()),
-                reinterpret_cast<const void*>(scale.data<float>()),
-                reinterpret_cast<void*>(out.data<paddle::bfloat16>()),
-                m,
-                n,
-                k,
-                x.stream()
-            };
-            if (!int8_sq_kernel_launcher<int8_t, __nv_bfloat16>(params)) {
-                PADDLE_THROW(common::errors::Fatal("gemm dequamt kernel run error"));
-            }
-            return {out};
-        } else if (out_dtype == "float16") {
-            paddle::Tensor out =
-                    paddle::empty({m, n}, paddle::DataType::FLOAT16, x.place());
-            GemmDequantParams params = {
-                reinterpret_cast<const void*>(x.data<int8_t>()),
-                reinterpret_cast<const void*>(y.data<int8_t>()),
-                reinterpret_cast<const void*>(scale.data<float>()),
-                reinterpret_cast<void*>(out.data<paddle::float16>()),
-                m,
-                n,
-                k,
-                x.stream()
-            };
-            if (!int8_sq_kernel_launcher<int8_t, half>(params)) {
-                PADDLE_THROW(common::errors::Fatal("gemm dequamt kernel run error"));
-            }
-            return {out};
-        } else {
-            PADDLE_THROW(phi::errors::InvalidArgument(
-                "only support bfloat16 and float16, but got %s", out_dtype));
-        }
-    }
+  std::vector<int64_t> x_dims = x.shape(), y_dims = y.shape();
+  PADDLE_ENFORCE_EQ(x_dims[x_dims.size() - 1],
+                    y_dims[y_dims.size() - 1],
+                    phi::errors::InvalidArgument(
+                        "The last dimension of x and y should be equal. But "
+                        "received x[%d] != y[%d].",
+                        "Ensure that x is not transposed and y is transposed.",
+                        x_dims[x_dims.size() - 1],
+                        y_dims[y_dims.size() - 1]));
+  int64_t m = x_dims[x_dims.size() - 2];
+  int64_t k = x_dims[x_dims.size() - 1];
+  int64_t n = y_dims[y_dims.size() - 2];
+  if (m <= 4) {
     if (out_dtype == "bfloat16") {
-        paddle::Tensor out =
-            paddle::empty({m, n}, paddle::DataType::BFLOAT16, x.place());
-        RunGemmDequant<paddle::DataType::BFLOAT16, paddle::bfloat16>(
-            x.data<int8_t>(),
-            y.data<int8_t>(),
-            scale.data<float>(),
-            out.data<paddle::bfloat16>(),
-            m,
-            k,
-            n,
-            x.stream());
-        return {out};
+      paddle::Tensor out =
+          paddle::empty({m, n}, paddle::DataType::BFLOAT16, x.place());
+      GemmDequantParams params = {
+          reinterpret_cast<const void *>(x.data<int8_t>()),
+          reinterpret_cast<const void *>(y.data<int8_t>()),
+          reinterpret_cast<const void *>(scale.data<float>()),
+          reinterpret_cast<void *>(out.data<paddle::bfloat16>()),
+          m,
+          n,
+          k,
+          x.stream()};
+      if (!int8_sq_kernel_launcher<int8_t, __nv_bfloat16>(params)) {
+        PADDLE_THROW(common::errors::Fatal("gemm dequamt kernel run error"));
+      }
+      return {out};
     } else if (out_dtype == "float16") {
-        paddle::Tensor out =
-            paddle::empty({m, n}, paddle::DataType::FLOAT16, x.place());
-        RunGemmDequant<paddle::DataType::FLOAT16, paddle::float16>(
-            x.data<int8_t>(),
-            y.data<int8_t>(),
-            scale.data<float>(),
-            out.data<paddle::float16>(),
-            m,
-            k,
-            n,
-            x.stream());
-        return {out};
+      paddle::Tensor out =
+          paddle::empty({m, n}, paddle::DataType::FLOAT16, x.place());
+      GemmDequantParams params = {
+          reinterpret_cast<const void *>(x.data<int8_t>()),
+          reinterpret_cast<const void *>(y.data<int8_t>()),
+          reinterpret_cast<const void *>(scale.data<float>()),
+          reinterpret_cast<void *>(out.data<paddle::float16>()),
+          m,
+          n,
+          k,
+          x.stream()};
+      if (!int8_sq_kernel_launcher<int8_t, half>(params)) {
+        PADDLE_THROW(common::errors::Fatal("gemm dequamt kernel run error"));
+      }
+      return {out};
     } else {
-        PADDLE_THROW(phi::errors::InvalidArgument(
-            "only support bfloat16 and float16, but got %s", out_dtype));
+      PADDLE_THROW(phi::errors::InvalidArgument(
+          "only support bfloat16 and float16, but got %s", out_dtype));
     }
+  }
+  if (out_dtype == "bfloat16") {
+    paddle::Tensor out =
+        paddle::empty({m, n}, paddle::DataType::BFLOAT16, x.place());
+    RunGemmDequant<paddle::DataType::BFLOAT16, paddle::bfloat16>(
+        x.data<int8_t>(),
+        y.data<int8_t>(),
+        scale.data<float>(),
+        out.data<paddle::bfloat16>(),
+        m,
+        k,
+        n,
+        x.stream());
+    return {out};
+  } else if (out_dtype == "float16") {
+    paddle::Tensor out =
+        paddle::empty({m, n}, paddle::DataType::FLOAT16, x.place());
+    RunGemmDequant<paddle::DataType::FLOAT16, paddle::float16>(
+        x.data<int8_t>(),
+        y.data<int8_t>(),
+        scale.data<float>(),
+        out.data<paddle::float16>(),
+        m,
+        k,
+        n,
+        x.stream());
+    return {out};
+  } else {
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "only support bfloat16 and float16, but got %s", out_dtype));
+  }
 }
 
 std::vector<std::vector<int64_t>> GemmDequantShape(
@@ -326,21 +321,21 @@ std::vector<std::vector<int64_t>> GemmDequantShape(
     const std::vector<int64_t> &y,
     const std::vector<int64_t> &scale,
     const std::string &out_dtype) {
-    return {{x[x.size() - 2], y[y.size() - 2]}};
+  return {{x[x.size() - 2], y[y.size() - 2]}};
 }
 
 std::vector<paddle::DataType> GemmDequantDtype(const paddle::DataType &x,
                                                const paddle::DataType &y,
                                                const paddle::DataType &scale,
                                                const std::string &out_dtype) {
-    if (out_dtype == "bfloat16") {
-        return {paddle::DataType::BFLOAT16};
-    } else if (out_dtype == "float16") {
-        return {paddle::DataType::FLOAT16};
-    } else {
-        PADDLE_THROW(phi::errors::InvalidArgument(
-            "only support bfloat16 and float16, but got %s", out_dtype));
-    }
+  if (out_dtype == "bfloat16") {
+    return {paddle::DataType::BFLOAT16};
+  } else if (out_dtype == "float16") {
+    return {paddle::DataType::FLOAT16};
+  } else {
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "only support bfloat16 and float16, but got %s", out_dtype));
+  }
 }
 
 PD_BUILD_STATIC_OP(gemm_dequant)
diff --git a/custom_ops/gpu_ops/int8_gemm_with_cutlass/gemm_dequant.h b/custom_ops/gpu_ops/int8_gemm_with_cutlass/gemm_dequant.h
index e14b8de74e..b339de6b26 100644
--- a/custom_ops/gpu_ops/int8_gemm_with_cutlass/gemm_dequant.h
+++ b/custom_ops/gpu_ops/int8_gemm_with_cutlass/gemm_dequant.h
@@ -159,50 +159,50 @@ struct DefaultDequantGemm<ElementA,
                           GatherB,
                           ScatterD,
                           PermuteDLayout> {
-    static_assert(platform::is_same<LayoutC, layout::RowMajor>::value ||
-                      platform::is_same<LayoutC, layout::AffineRankN<2>>::value,
-                  "Epilogue in the kernel level must be row major");
+  static_assert(platform::is_same<LayoutC, layout::RowMajor>::value ||
+                    platform::is_same<LayoutC, layout::AffineRankN<2>>::value,
+                "Epilogue in the kernel level must be row major");
 
-    /// Define the threadblock-scoped matrix multiply-accumulate
-    using Mma = typename cutlass::gemm::threadblock::DefaultMma<
-        ElementA,
-        LayoutA,
-        kAlignmentA,
-        ElementB,
-        LayoutB,
-        kAlignmentB,
-        ElementAccumulator,
-        LayoutC,
-        arch::OpClassTensorOp,
-        arch::Sm80,
-        ThreadblockShape,
-        WarpShape,
-        InstructionShape,
-        Stages,
-        Operator,
-        false,
-        SharedMemoryClear,
-        GatherA,
-        GatherB>::ThreadblockMma;
+  /// Define the threadblock-scoped matrix multiply-accumulate
+  using Mma =
+      typename cutlass::gemm::threadblock::DefaultMma<ElementA,
+                                                      LayoutA,
+                                                      kAlignmentA,
+                                                      ElementB,
+                                                      LayoutB,
+                                                      kAlignmentB,
+                                                      ElementAccumulator,
+                                                      LayoutC,
+                                                      arch::OpClassTensorOp,
+                                                      arch::Sm80,
+                                                      ThreadblockShape,
+                                                      WarpShape,
+                                                      InstructionShape,
+                                                      Stages,
+                                                      Operator,
+                                                      false,
+                                                      SharedMemoryClear,
+                                                      GatherA,
+                                                      GatherB>::ThreadblockMma;
 
-    static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
+  static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
 
-    /// Define the epilogue
-    using RegularEpilogue =
-        typename epilogue::threadblock::DequantEpilogueTensorOp<
-            ThreadblockShape,
-            typename Mma::Operator,
-            kPartitionsK,
-            EpilogueOutputOp,
-            EpilogueOutputOp::kCount,
-            ScatterD,
-            PermuteDLayout>::Epilogue;
+  /// Define the epilogue
+  using RegularEpilogue =
+      typename epilogue::threadblock::DequantEpilogueTensorOp<
+          ThreadblockShape,
+          typename Mma::Operator,
+          kPartitionsK,
+          EpilogueOutputOp,
+          EpilogueOutputOp::kCount,
+          ScatterD,
+          PermuteDLayout>::Epilogue;
 
-    using Epilogue = RegularEpilogue;
+  using Epilogue = RegularEpilogue;
 
-    /// Define the kernel-level GEMM operator.
-    using GemmKernel =
-        kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+  /// Define the kernel-level GEMM operator.
+  using GemmKernel =
+      kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
 };
 
 template <
@@ -340,49 +340,49 @@ struct DefaultDequantGemmUniversal<
     PermuteDLayout,
     typename platform::enable_if<
         !cutlass::is_complex<ElementAccumulator>::value>::type> {
-    using DefaultGemmKernel =
-        typename kernel::DefaultDequantGemm<ElementA,
-                                            LayoutA,
-                                            kAlignmentA,
-                                            ElementB,
-                                            LayoutB,
-                                            kAlignmentB,
-                                            ElementC,
-                                            LayoutC,
-                                            ElementAccumulator,
-                                            OperatorClass,
-                                            ArchTag,
-                                            ThreadblockShape,
-                                            WarpShape,
-                                            InstructionShape,
-                                            EpilogueOutputOp,
-                                            ThreadblockSwizzle,
-                                            Stages,
-                                            true,
-                                            Operator,
-                                            SharedMemoryClear,
-                                            GatherA,
-                                            GatherB,
-                                            ScatterD,
-                                            PermuteDLayout>::GemmKernel;
+  using DefaultGemmKernel =
+      typename kernel::DefaultDequantGemm<ElementA,
+                                          LayoutA,
+                                          kAlignmentA,
+                                          ElementB,
+                                          LayoutB,
+                                          kAlignmentB,
+                                          ElementC,
+                                          LayoutC,
+                                          ElementAccumulator,
+                                          OperatorClass,
+                                          ArchTag,
+                                          ThreadblockShape,
+                                          WarpShape,
+                                          InstructionShape,
+                                          EpilogueOutputOp,
+                                          ThreadblockSwizzle,
+                                          Stages,
+                                          true,
+                                          Operator,
+                                          SharedMemoryClear,
+                                          GatherA,
+                                          GatherB,
+                                          ScatterD,
+                                          PermuteDLayout>::GemmKernel;
 
-    /// Universal kernel without StreamkFeature member type
-    template <class SwizzleT, class Enable = void>
-    class SelectBase
-        : public kernel::GemmUniversal<typename DefaultGemmKernel::Mma,
-                                       typename DefaultGemmKernel::Epilogue,
-                                       SwizzleT> {};
+  /// Universal kernel without StreamkFeature member type
+  template <class SwizzleT, class Enable = void>
+  class SelectBase
+      : public kernel::GemmUniversal<typename DefaultGemmKernel::Mma,
+                                     typename DefaultGemmKernel::Epilogue,
+                                     SwizzleT> {};
 
-    /// Universal kernel with StreamkFeature member type
-    template <class SwizzleT>
-    class SelectBase<SwizzleT, typename SwizzleT::StreamkFeature>
-        : public kernel::GemmUniversalStreamk<
-              typename DefaultGemmKernel::Mma,
-              typename DefaultGemmKernel::Epilogue,
-              SwizzleT> {};
+  /// Universal kernel with StreamkFeature member type
+  template <class SwizzleT>
+  class SelectBase<SwizzleT, typename SwizzleT::StreamkFeature>
+      : public kernel::GemmUniversalStreamk<
+            typename DefaultGemmKernel::Mma,
+            typename DefaultGemmKernel::Epilogue,
+            SwizzleT> {};
 
-    /// Select kernel by ThreadblockSwizzle's support for StreamkFeature
-    using GemmKernel = SelectBase<ThreadblockSwizzle>;
+  /// Select kernel by ThreadblockSwizzle's support for StreamkFeature
+  using GemmKernel = SelectBase<ThreadblockSwizzle>;
 };
 
 template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
@@ -390,431 +390,413 @@ template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
           typename ThreadblockSwizzle_  ///! Threadblock swizzling function
           >
 struct GemmWithEpilogueVisitorDequant {
-    public:
-    using Mma = Mma_;
-    using Epilogue = Epilogue_;
-    using EpilogueVisitor = typename Epilogue::Visitor;
-    using ThreadblockSwizzle = ThreadblockSwizzle_;
+ public:
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using EpilogueVisitor = typename Epilogue::Visitor;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
 
-    using ElementA = typename Mma::IteratorA::Element;
-    using LayoutA = typename Mma::IteratorA::Layout;
-    using TensorRefA = TensorRef<ElementA, LayoutA>;
+  using ElementA = typename Mma::IteratorA::Element;
+  using LayoutA = typename Mma::IteratorA::Layout;
+  using TensorRefA = TensorRef<ElementA, LayoutA>;
 
-    using ElementB = typename Mma::IteratorB::Element;
-    using LayoutB = typename Mma::IteratorB::Layout;
-    using TensorRefB = TensorRef<ElementB, LayoutB>;
+  using ElementB = typename Mma::IteratorB::Element;
+  using LayoutB = typename Mma::IteratorB::Layout;
+  using TensorRefB = TensorRef<ElementB, LayoutB>;
 
-    using ElementC = typename EpilogueVisitor::ElementOutput;
-    using LayoutC = typename Epilogue::Layout;
-    using TensorRefC = TensorRef<ElementC, LayoutC>;
+  using ElementC = typename EpilogueVisitor::ElementOutput;
+  using LayoutC = typename Epilogue::Layout;
+  using TensorRefC = TensorRef<ElementC, LayoutC>;
 
-    using ElementScale = typename EpilogueVisitor::AlphaScaleElementType;
-    using LayoutScale = cutlass::layout::RowMajor;
-    using TensorRefScale = TensorRef<ElementScale, LayoutScale>;
+  using ElementScale = typename EpilogueVisitor::AlphaScaleElementType;
+  using LayoutScale = cutlass::layout::RowMajor;
+  using TensorRefScale = TensorRef<ElementScale, LayoutScale>;
 
-    static ComplexTransform const kTransformA = Mma::kTransformA;
-    static ComplexTransform const kTransformB = Mma::kTransformB;
-    using Operator = typename Mma::Operator;
+  static ComplexTransform const kTransformA = Mma::kTransformA;
+  static ComplexTransform const kTransformB = Mma::kTransformB;
+  using Operator = typename Mma::Operator;
 
-    using OperatorClass = typename Mma::Operator::OperatorClass;
-    using ThreadblockShape = typename Mma::Shape;
-    using WarpShape = typename Mma::Operator::Shape;
-    using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
-    using ArchTag = typename Mma::ArchTag;
+  using OperatorClass = typename Mma::Operator::OperatorClass;
+  using ThreadblockShape = typename Mma::Shape;
+  using WarpShape = typename Mma::Operator::Shape;
+  using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
+  using ArchTag = typename Mma::ArchTag;
 
-    static int const kStages = Mma::kStages;
-    static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
-    static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
-    static int const kAlignmentC = EpilogueVisitor::kElementsPerAccess;
+  static int const kStages = Mma::kStages;
+  static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+  static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+  static int const kAlignmentC = EpilogueVisitor::kElementsPerAccess;
 
-    /// Warp count (concept: GemmShape)
-    using WarpCount = typename Mma::WarpCount;
-    static int const kThreadCount = 32 * WarpCount::kCount;
+  /// Warp count (concept: GemmShape)
+  using WarpCount = typename Mma::WarpCount;
+  static int const kThreadCount = 32 * WarpCount::kCount;
 
-    /// Split-K preserves splits that are 128b aligned
-    static int const kSplitKAlignment = const_max(
-        128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value);
+  /// Split-K preserves splits that are 128b aligned
+  static int const kSplitKAlignment = const_max(
+      128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value);
 
+  //
+  // Structures
+  //
+
+  /// Argument structure
+  struct Arguments {
     //
-    // Structures
+    // Data members
     //
 
-    /// Argument structure
-    struct Arguments {
-        //
-        // Data members
-        //
+    GemmUniversalMode mode;
+    GemmCoord problem_size;
 
-        GemmUniversalMode mode;
-        GemmCoord problem_size;
+    TensorRefA ref_A;
+    TensorRefB ref_B;
+    TensorRefC ref_C;
+    TensorRefC ref_D;
+    TensorRefScale ref_scale;
 
-        TensorRefA ref_A;
-        TensorRefB ref_B;
-        TensorRefC ref_C;
-        TensorRefC ref_D;
-        TensorRefScale ref_scale;
+    typename EpilogueVisitor::Arguments epilogue_visitor;
 
-        typename EpilogueVisitor::Arguments epilogue_visitor;
-
-        //
-        // Methods
-        //
-
-        Arguments() : mode(GemmUniversalMode::kGemm) {}
-
-        /// constructs an arguments structure
-        Arguments(GemmUniversalMode mode_,
-                  GemmCoord problem_size_,
-                  TensorRefA ref_A_,
-                  TensorRefB ref_B_,
-                  TensorRefC ref_C_,
-                  TensorRefC ref_D_,
-                  TensorRefScale ref_scale_,
-                  typename EpilogueVisitor::Arguments epilogue_visitor_)
-            : mode(mode_),
-              problem_size(problem_size_),
-              ref_A(ref_A_),
-              ref_B(ref_B_),
-              ref_C(ref_C_),
-              ref_D(ref_D_),
-              ref_scale(ref_scale_),
-              epilogue_visitor(epilogue_visitor_) {}
-    };
-
-    //
-    // Structure for precomputing values in host memory and passing to kernels
-    //
-
-    /// Parameters structure
-    struct Params {
-        cutlass::gemm::GemmCoord problem_size;
-        cutlass::gemm::GemmCoord grid_tiled_shape;
-        int swizzle_log_tile;
-
-        typename Mma::IteratorA::Params params_A;
-        typename Mma::IteratorB::Params params_B;
-        typename EpilogueVisitor::OutputTileIterator::Params params_C;
-        typename EpilogueVisitor::OutputTileIterator::Params params_D;
-        typename EpilogueVisitor::ScaleTileIterator::Params params_scale;
-
-        GemmUniversalMode mode;
-        int gemm_k_size;
-
-        void *ptr_A;
-        void *ptr_B;
-        ElementC *ptr_C;
-        ElementC *ptr_D;
-        typename EpilogueVisitor::AlphaScaleElementType *ptr_dequant_scale;
-
-        typename EpilogueVisitor::Params epilogue_visitor;
-
-        //
-        // Methods
-        //
-
-        CUTLASS_HOST_DEVICE
-        Params()
-            : swizzle_log_tile(0),
-              params_A(0),
-              params_B(0),
-              params_C(0),
-              params_D(0),
-              params_scale(0),
-              gemm_k_size(0),
-              mode(cutlass::gemm::GemmUniversalMode::kGemm),
-              ptr_A(nullptr),
-              ptr_B(nullptr),
-              ptr_C(nullptr),
-              ptr_D(nullptr),
-              ptr_dequant_scale(nullptr) {}
-
-        explicit Params(Arguments const &args)
-            : problem_size(args.problem_size),
-              swizzle_log_tile(0),
-              params_A(args.ref_A.layout()),
-              params_B(args.ref_B.layout()),
-              params_C(args.ref_C.layout()),
-              params_D(args.ref_D.layout()),
-              mode(args.mode),
-              gemm_k_size(args.problem_size.k()),
-              ptr_A(args.ref_A.data()),
-              ptr_B(args.ref_B.data()),
-              ptr_C(args.ref_C.data()),
-              ptr_D(args.ref_D.data()),
-              ptr_dequant_scale(args.ref_scale.data()),
-              epilogue_visitor(args.epilogue_visitor) {
-            ThreadblockSwizzle threadblock_swizzle;
-
-            grid_tiled_shape =
-                threadblock_swizzle.get_tiled_shape(args.problem_size,
-                                                    {ThreadblockShape::kM,
-                                                     ThreadblockShape::kN,
-                                                     ThreadblockShape::kK},
-                                                    1);
-
-            if (args.mode == GemmUniversalMode::kGemm ||
-                args.mode == GemmUniversalMode::kGemmSplitKParallel) {
-                int const kAlignK =
-                    const_max(const_max(128 / sizeof_bits<ElementA>::value,
-                                        128 / sizeof_bits<ElementB>::value),
-                              1);
-
-                gemm_k_size =
-                    round_up(ceil_div(args.problem_size.k(), 1), kAlignK);
-
-                if (gemm_k_size) {
-                    grid_tiled_shape.k() =
-                        ceil_div(args.problem_size.k(), gemm_k_size);
-                }
-            }
-
-            swizzle_log_tile =
-                threadblock_swizzle.get_log_tile(grid_tiled_shape);
-        }
-    };
-
-    /// Shared memory storage structure
-    union SharedStorage {
-        typename Mma::SharedStorage main_loop;
-
-        struct {
-            typename Epilogue::SharedStorage epilogue;
-            typename EpilogueVisitor::SharedStorage visitor;
-        } epilogue;
-    };
-
-    public:
     //
     // Methods
     //
 
-    CUTLASS_DEVICE
-    GemmWithEpilogueVisitorDequant() {}
+    Arguments() : mode(GemmUniversalMode::kGemm) {}
 
-    /// Determines whether kernel satisfies alignment
-    static Status can_implement(cutlass::gemm::GemmCoord const &problem_size) {
-        CUTLASS_TRACE_HOST("GemmWithEpilogueVisitorDequant::can_implement()");
+    /// constructs an arguments structure
+    Arguments(GemmUniversalMode mode_,
+              GemmCoord problem_size_,
+              TensorRefA ref_A_,
+              TensorRefB ref_B_,
+              TensorRefC ref_C_,
+              TensorRefC ref_D_,
+              TensorRefScale ref_scale_,
+              typename EpilogueVisitor::Arguments epilogue_visitor_)
+        : mode(mode_),
+          problem_size(problem_size_),
+          ref_A(ref_A_),
+          ref_B(ref_B_),
+          ref_C(ref_C_),
+          ref_D(ref_D_),
+          ref_scale(ref_scale_),
+          epilogue_visitor(epilogue_visitor_) {}
+  };
 
-        static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
-        static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
-        static int const kAlignmentC =
-            Epilogue::OutputTileIterator::kElementsPerAccess;
+  //
+  // Structure for precomputing values in host memory and passing to kernels
+  //
 
-        bool isAMisaligned = false;
-        bool isBMisaligned = false;
-        bool isCMisaligned = false;
+  /// Parameters structure
+  struct Params {
+    cutlass::gemm::GemmCoord problem_size;
+    cutlass::gemm::GemmCoord grid_tiled_shape;
+    int swizzle_log_tile;
 
-        if (platform::is_same<LayoutA, layout::RowMajor>::value) {
-            isAMisaligned = problem_size.k() % kAlignmentA;
-        } else if (platform::is_same<LayoutA, layout::ColumnMajor>::value) {
-            isAMisaligned = problem_size.m() % kAlignmentA;
-        } else if (platform::is_same<
-                       LayoutA,
-                       layout::ColumnMajorInterleaved<32>>::value ||
-                   platform::is_same<
-                       LayoutA,
-                       layout::ColumnMajorInterleaved<64>>::value) {
-            isAMisaligned = problem_size.k() % kAlignmentA;
+    typename Mma::IteratorA::Params params_A;
+    typename Mma::IteratorB::Params params_B;
+    typename EpilogueVisitor::OutputTileIterator::Params params_C;
+    typename EpilogueVisitor::OutputTileIterator::Params params_D;
+    typename EpilogueVisitor::ScaleTileIterator::Params params_scale;
+
+    GemmUniversalMode mode;
+    int gemm_k_size;
+
+    void *ptr_A;
+    void *ptr_B;
+    ElementC *ptr_C;
+    ElementC *ptr_D;
+    typename EpilogueVisitor::AlphaScaleElementType *ptr_dequant_scale;
+
+    typename EpilogueVisitor::Params epilogue_visitor;
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : swizzle_log_tile(0),
+          params_A(0),
+          params_B(0),
+          params_C(0),
+          params_D(0),
+          params_scale(0),
+          gemm_k_size(0),
+          mode(cutlass::gemm::GemmUniversalMode::kGemm),
+          ptr_A(nullptr),
+          ptr_B(nullptr),
+          ptr_C(nullptr),
+          ptr_D(nullptr),
+          ptr_dequant_scale(nullptr) {}
+
+    explicit Params(Arguments const &args)
+        : problem_size(args.problem_size),
+          swizzle_log_tile(0),
+          params_A(args.ref_A.layout()),
+          params_B(args.ref_B.layout()),
+          params_C(args.ref_C.layout()),
+          params_D(args.ref_D.layout()),
+          mode(args.mode),
+          gemm_k_size(args.problem_size.k()),
+          ptr_A(args.ref_A.data()),
+          ptr_B(args.ref_B.data()),
+          ptr_C(args.ref_C.data()),
+          ptr_D(args.ref_D.data()),
+          ptr_dequant_scale(args.ref_scale.data()),
+          epilogue_visitor(args.epilogue_visitor) {
+      ThreadblockSwizzle threadblock_swizzle;
+
+      grid_tiled_shape = threadblock_swizzle.get_tiled_shape(
+          args.problem_size,
+          {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+          1);
+
+      if (args.mode == GemmUniversalMode::kGemm ||
+          args.mode == GemmUniversalMode::kGemmSplitKParallel) {
+        int const kAlignK =
+            const_max(const_max(128 / sizeof_bits<ElementA>::value,
+                                128 / sizeof_bits<ElementB>::value),
+                      1);
+
+        gemm_k_size = round_up(ceil_div(args.problem_size.k(), 1), kAlignK);
+
+        if (gemm_k_size) {
+          grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size);
         }
+      }
 
-        if (platform::is_same<LayoutB, layout::RowMajor>::value) {
-            isBMisaligned = problem_size.n() % kAlignmentB;
-        } else if (platform::is_same<LayoutB, layout::ColumnMajor>::value) {
-            isBMisaligned = problem_size.k() % kAlignmentB;
-        } else if (platform::is_same<LayoutB,
-                                     layout::RowMajorInterleaved<32>>::value ||
-                   platform::is_same<LayoutB,
-                                     layout::RowMajorInterleaved<64>>::value) {
-            isBMisaligned = problem_size.k() % kAlignmentB;
-        }
+      swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape);
+    }
+  };
 
-        if (platform::is_same<LayoutC, layout::RowMajor>::value) {
-            isCMisaligned = problem_size.n() % kAlignmentC;
-        } else if (platform::is_same<LayoutC, layout::ColumnMajor>::value) {
-            isCMisaligned = problem_size.m() % kAlignmentC;
-        } else if (platform::is_same<
-                       LayoutC,
-                       layout::ColumnMajorInterleaved<32>>::value ||
-                   platform::is_same<
-                       LayoutC,
-                       layout::ColumnMajorInterleaved<64>>::value) {
-            isCMisaligned = problem_size.n() % kAlignmentC;
-        }
+  /// Shared memory storage structure
+  union SharedStorage {
+    typename Mma::SharedStorage main_loop;
 
-        if (isAMisaligned) {
-            CUTLASS_TRACE_HOST(
-                "  returning kErrorMisalignedOperand for A operand");
-            return Status::kErrorMisalignedOperand;
-        }
+    struct {
+      typename Epilogue::SharedStorage epilogue;
+      typename EpilogueVisitor::SharedStorage visitor;
+    } epilogue;
+  };
 
-        if (isBMisaligned) {
-            CUTLASS_TRACE_HOST(
-                "  returning kErrorMisalignedOperand for B operand");
-            return Status::kErrorMisalignedOperand;
-        }
+ public:
+  //
+  // Methods
+  //
 
-        if (isCMisaligned) {
-            CUTLASS_TRACE_HOST(
-                "  returning kErrorMisalignedOperand for C operand");
-            return Status::kErrorMisalignedOperand;
-        }
+  CUTLASS_DEVICE
+  GemmWithEpilogueVisitorDequant() {}
 
-        CUTLASS_TRACE_HOST("  returning kSuccess");
+  /// Determines whether kernel satisfies alignment
+  static Status can_implement(cutlass::gemm::GemmCoord const &problem_size) {
+    CUTLASS_TRACE_HOST("GemmWithEpilogueVisitorDequant::can_implement()");
 
-        return Status::kSuccess;
+    static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+    static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+    static int const kAlignmentC =
+        Epilogue::OutputTileIterator::kElementsPerAccess;
+
+    bool isAMisaligned = false;
+    bool isBMisaligned = false;
+    bool isCMisaligned = false;
+
+    if (platform::is_same<LayoutA, layout::RowMajor>::value) {
+      isAMisaligned = problem_size.k() % kAlignmentA;
+    } else if (platform::is_same<LayoutA, layout::ColumnMajor>::value) {
+      isAMisaligned = problem_size.m() % kAlignmentA;
+    } else if (platform::is_same<LayoutA,
+                                 layout::ColumnMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutA,
+                                 layout::ColumnMajorInterleaved<64>>::value) {
+      isAMisaligned = problem_size.k() % kAlignmentA;
     }
 
-    static Status can_implement(Arguments const &args) {
-        return can_implement(args.problem_size);
+    if (platform::is_same<LayoutB, layout::RowMajor>::value) {
+      isBMisaligned = problem_size.n() % kAlignmentB;
+    } else if (platform::is_same<LayoutB, layout::ColumnMajor>::value) {
+      isBMisaligned = problem_size.k() % kAlignmentB;
+    } else if (platform::is_same<LayoutB,
+                                 layout::RowMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutB,
+                                 layout::RowMajorInterleaved<64>>::value) {
+      isBMisaligned = problem_size.k() % kAlignmentB;
     }
 
+    if (platform::is_same<LayoutC, layout::RowMajor>::value) {
+      isCMisaligned = problem_size.n() % kAlignmentC;
+    } else if (platform::is_same<LayoutC, layout::ColumnMajor>::value) {
+      isCMisaligned = problem_size.m() % kAlignmentC;
+    } else if (platform::is_same<LayoutC,
+                                 layout::ColumnMajorInterleaved<32>>::value ||
+               platform::is_same<LayoutC,
+                                 layout::ColumnMajorInterleaved<64>>::value) {
+      isCMisaligned = problem_size.n() % kAlignmentC;
+    }
+
+    if (isAMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for A operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (isBMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for B operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (isCMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for C operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    CUTLASS_TRACE_HOST("  returning kSuccess");
+
+    return Status::kSuccess;
+  }
+
+  static Status can_implement(Arguments const &args) {
+    return can_implement(args.problem_size);
+  }
+
 #define SPLIT_K_ENABLED 1
 
-    /// Executes one GEMM
-    CUTLASS_DEVICE
-    void operator()(Params const &params,
-                    SharedStorage &shared_storage) {  // NOLINT
-        // Compute threadblock location
-        ThreadblockSwizzle threadblock_swizzle;
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const &params,
+                  SharedStorage &shared_storage) {  // NOLINT
+    // Compute threadblock location
+    ThreadblockSwizzle threadblock_swizzle;
 
-        cutlass::gemm::GemmCoord threadblock_tile_offset =
-            threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+    cutlass::gemm::GemmCoord threadblock_tile_offset =
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-        // Early exit if CTA is out of range
+    // Early exit if CTA is out of range
 
-        if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
-            params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
-            return;
-        }
+    if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
+        params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
+      return;
+    }
 
-        int offset_k = 0;
-        int problem_size_k = params.problem_size.k();
+    int offset_k = 0;
+    int problem_size_k = params.problem_size.k();
 
-        ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A);
-        ElementB *ptr_B = static_cast<ElementB *>(params.ptr_B);
+    ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A);
+    ElementB *ptr_B = static_cast<ElementB *>(params.ptr_B);
 
 #if SPLIT_K_ENABLED
-        //
-        // Fetch pointers based on mode.
-        //
-        if (params.mode == GemmUniversalMode::kGemm ||
-            params.mode == GemmUniversalMode::kGemmSplitKParallel) {
-            if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
-                problem_size_k =
-                    (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
-            }
+    //
+    // Fetch pointers based on mode.
+    //
+    if (params.mode == GemmUniversalMode::kGemm ||
+        params.mode == GemmUniversalMode::kGemmSplitKParallel) {
+      if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
+        problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
+      }
 
-            offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
-        }
+      offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
+    }
 #endif
 
-        // Compute initial location in logical coordinates
-        cutlass::MatrixCoord tb_offset_A{
-            threadblock_tile_offset.m() * Mma::Shape::kM,
-            offset_k,
-        };
+    // Compute initial location in logical coordinates
+    cutlass::MatrixCoord tb_offset_A{
+        threadblock_tile_offset.m() * Mma::Shape::kM,
+        offset_k,
+    };
 
-        cutlass::MatrixCoord tb_offset_B{
-            offset_k, threadblock_tile_offset.n() * Mma::Shape::kN};
+    cutlass::MatrixCoord tb_offset_B{
+        offset_k, threadblock_tile_offset.n() * Mma::Shape::kN};
 
-        // Compute position within threadblock
-        int thread_idx = threadIdx.x;
+    // Compute position within threadblock
+    int thread_idx = threadIdx.x;
 
-        // Construct iterators to A and B operands
-        typename Mma::IteratorA iterator_A(
-            params.params_A,
-            ptr_A,
-            {params.problem_size.m(), problem_size_k},
-            thread_idx,
-            tb_offset_A);
+    // Construct iterators to A and B operands
+    typename Mma::IteratorA iterator_A(
+        params.params_A,
+        ptr_A,
+        {params.problem_size.m(), problem_size_k},
+        thread_idx,
+        tb_offset_A);
 
-        typename Mma::IteratorB iterator_B(
-            params.params_B,
-            ptr_B,
-            {problem_size_k, params.problem_size.n()},
-            thread_idx,
-            tb_offset_B);
+    typename Mma::IteratorB iterator_B(
+        params.params_B,
+        ptr_B,
+        {problem_size_k, params.problem_size.n()},
+        thread_idx,
+        tb_offset_B);
 
-        // Broadcast the warp_id computed by lane 0 to ensure dependent code
-        // is compiled as warp-uniform.
-        int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+    // Broadcast the warp_id computed by lane 0 to ensure dependent code
+    // is compiled as warp-uniform.
+    int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
 
-        int lane_idx = threadIdx.x % 32;
+    int lane_idx = threadIdx.x % 32;
 
-        //
-        // Main loop
-        //
+    //
+    // Main loop
+    //
 
-        // Construct thread-scoped matrix multiply
-        Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
+    // Construct thread-scoped matrix multiply
+    Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
 
-        typename Mma::FragmentC accumulators;
+    typename Mma::FragmentC accumulators;
 
-        accumulators.clear();
+    accumulators.clear();
 
-        // Compute threadblock-scoped matrix multiply-add
-        int gemm_k_iterations =
-            (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
+    // Compute threadblock-scoped matrix multiply-add
+    int gemm_k_iterations =
+        (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
 
-        // Compute threadblock-scoped matrix multiply-add
-        mma(gemm_k_iterations,
-            accumulators,
-            iterator_A,
-            iterator_B,
-            accumulators);
+    // Compute threadblock-scoped matrix multiply-add
+    mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators);
 
-        //
-        // Masked tile iterators constructed from members
-        //
+    //
+    // Masked tile iterators constructed from members
+    //
 
-        threadblock_tile_offset =
-            threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+    threadblock_tile_offset =
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
-        // assume identity swizzle
-        MatrixCoord threadblock_offset(
-            threadblock_tile_offset.m() * Mma::Shape::kM,
-            threadblock_tile_offset.n() * Mma::Shape::kN);
+    // assume identity swizzle
+    MatrixCoord threadblock_offset(
+        threadblock_tile_offset.m() * Mma::Shape::kM,
+        threadblock_tile_offset.n() * Mma::Shape::kN);
 
-        int block_idx =
-            threadblock_tile_offset.m() +
-            threadblock_tile_offset.n() * params.grid_tiled_shape.m();
+    int block_idx = threadblock_tile_offset.m() +
+                    threadblock_tile_offset.n() * params.grid_tiled_shape.m();
 
-        //
-        // Construct the epilogue visitor
-        //
+    //
+    // Construct the epilogue visitor
+    //
 
-        EpilogueVisitor epilogue_visitor(params.epilogue_visitor,
-                                         shared_storage.epilogue.visitor,
-                                         params.problem_size.mn(),
-                                         thread_idx,
-                                         warp_idx,
-                                         lane_idx,
-                                         params.params_C,
-                                         params.params_D,
-                                         params.params_scale,
-                                         params.ptr_C,
-                                         params.ptr_D,
-                                         params.ptr_dequant_scale,
-                                         threadblock_offset,
-                                         blockIdx.y * params.problem_size.m());
+    EpilogueVisitor epilogue_visitor(params.epilogue_visitor,
+                                     shared_storage.epilogue.visitor,
+                                     params.problem_size.mn(),
+                                     thread_idx,
+                                     warp_idx,
+                                     lane_idx,
+                                     params.params_C,
+                                     params.params_D,
+                                     params.params_scale,
+                                     params.ptr_C,
+                                     params.ptr_D,
+                                     params.ptr_dequant_scale,
+                                     threadblock_offset,
+                                     blockIdx.y * params.problem_size.m());
 
-        if (params.mode == GemmUniversalMode::kGemm) {
-            // Indicate which position in a serial reduction the output operator
-            // is currently updating
-            epilogue_visitor.set_k_partition(threadblock_tile_offset.k(),
-                                             params.grid_tiled_shape.k());
-        } else if (params.mode == GemmUniversalMode::kBatched ||
-                   params.mode == GemmUniversalMode::kArray) {
-            epilogue_visitor.set_batch_index(threadblock_tile_offset.k());
-        }
-
-        // Construct the epilogue
-        Epilogue epilogue(
-            shared_storage.epilogue.epilogue, thread_idx, warp_idx, lane_idx);
-
-        // Execute the epilogue operator to update the destination tensor.
-        epilogue(epilogue_visitor, accumulators);
+    if (params.mode == GemmUniversalMode::kGemm) {
+      // Indicate which position in a serial reduction the output operator
+      // is currently updating
+      epilogue_visitor.set_k_partition(threadblock_tile_offset.k(),
+                                       params.grid_tiled_shape.k());
+    } else if (params.mode == GemmUniversalMode::kBatched ||
+               params.mode == GemmUniversalMode::kArray) {
+      epilogue_visitor.set_batch_index(threadblock_tile_offset.k());
     }
+
+    // Construct the epilogue
+    Epilogue epilogue(
+        shared_storage.epilogue.epilogue, thread_idx, warp_idx, lane_idx);
+
+    // Execute the epilogue operator to update the destination tensor.
+    epilogue(epilogue_visitor, accumulators);
+  }
 };
 }  // namespace kernel
 }  // namespace gemm
@@ -830,286 +812,284 @@ template <typename ThreadblockShape_,
           typename ElementwiseFunctor_,
           bool UseMasking_ = false>
 class DequantEpilogueVisitor {
-    public:
-    using ThreadblockShape = ThreadblockShape_;
-    static int const kThreadCount = ThreadCount;
+ public:
+  using ThreadblockShape = ThreadblockShape_;
+  static int const kThreadCount = ThreadCount;
 
-    using ScaleTileIterator = ScaleTileIterator_;
-    using OutputTileIterator = OutputTileIterator_;
-    using ElementwiseFunctor = ElementwiseFunctor_;
+  using ScaleTileIterator = ScaleTileIterator_;
+  using OutputTileIterator = OutputTileIterator_;
+  using ElementwiseFunctor = ElementwiseFunctor_;
 
-    static int const kIterations = OutputTileIterator::kIterations;
-    static int const kElementsPerAccess =
-        OutputTileIterator::kElementsPerAccess;
+  static int const kIterations = OutputTileIterator::kIterations;
+  static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
 
-    using ElementOutput = typename OutputTileIterator::Element;
-    using LayoutOutput = cutlass::layout::RowMajor;
-    using ElementAccumulator = ElementAccumulator_;
+  using ElementOutput = typename OutputTileIterator::Element;
+  using LayoutOutput = cutlass::layout::RowMajor;
+  using ElementAccumulator = ElementAccumulator_;
 
-    using AlphaScaleElementType = typename ScaleTileIterator::Element;
+  using AlphaScaleElementType = typename ScaleTileIterator::Element;
 
-    using ElementCompute = ElementCompute_;
-    using AccumulatorFragment = Array<ElementAccumulator, kElementsPerAccess>;
-    using ComputeFragment = Array<ElementCompute_, kElementsPerAccess>;
-    using OutputVector = Array<ElementOutput, kElementsPerAccess>;
+  using ElementCompute = ElementCompute_;
+  using AccumulatorFragment = Array<ElementAccumulator, kElementsPerAccess>;
+  using ComputeFragment = Array<ElementCompute_, kElementsPerAccess>;
+  using OutputVector = Array<ElementOutput, kElementsPerAccess>;
 
-    static int const kThreadsPerRow =
-        OutputTileIterator::ThreadMap::Detail::kAccessWidth;
-    static bool const kHasMultiStepsInRow =
-        (OutputTileIterator::ThreadMap::Iterations::kColumn > 1);
+  static int const kThreadsPerRow =
+      OutputTileIterator::ThreadMap::Detail::kAccessWidth;
+  static bool const kHasMultiStepsInRow =
+      (OutputTileIterator::ThreadMap::Iterations::kColumn > 1);
 
-    /// Argument structure
-    struct Arguments {
-        typename ElementwiseFunctor::Params elementwise;
+  /// Argument structure
+  struct Arguments {
+    typename ElementwiseFunctor::Params elementwise;
 
-        //
-        // Methods
-        //
+    //
+    // Methods
+    //
 
-        explicit Arguments(typename ElementwiseFunctor::Params elementwise_)
-            : elementwise(elementwise_) {}
-    };
+    explicit Arguments(typename ElementwiseFunctor::Params elementwise_)
+        : elementwise(elementwise_) {}
+  };
 
-    struct Params {
-        typename ElementwiseFunctor::Params elementwise;
-        //
-        // Methods
-        //
-        CUTLASS_HOST_DEVICE
-        Params() {}
+  struct Params {
+    typename ElementwiseFunctor::Params elementwise;
+    //
+    // Methods
+    //
+    CUTLASS_HOST_DEVICE
+    Params() {}
 
-        CUTLASS_HOST_DEVICE
-        explicit Params(Arguments const &args)
-            : elementwise(args.elementwise) {}
-    };
+    CUTLASS_HOST_DEVICE
+    explicit Params(Arguments const &args) : elementwise(args.elementwise) {}
+  };
 
-    /// Shared storage
-    struct SharedStorage {};
+  /// Shared storage
+  struct SharedStorage {};
 
-    private:
-    Params const &params_;
-    SharedStorage &shared_storage_;
-    MatrixCoord extent_;
-    MatrixCoord extent_real_;
-    ElementwiseFunctor elementwise_;
+ private:
+  Params const &params_;
+  SharedStorage &shared_storage_;
+  MatrixCoord extent_;
+  MatrixCoord extent_real_;
+  ElementwiseFunctor elementwise_;
 
-    AlphaScaleElementType *ptr_dequant_scale_;
-    ScaleTileIterator iterator_dequant_scale_;
-    OutputTileIterator iterator_C_;
-    OutputTileIterator iterator_D_;
+  AlphaScaleElementType *ptr_dequant_scale_;
+  ScaleTileIterator iterator_dequant_scale_;
+  OutputTileIterator iterator_C_;
+  OutputTileIterator iterator_D_;
 
-    AlphaScaleElementType element_alpha_row_ = 1.0f;
-    AlphaScaleElementType element_alpha_col_ = 1.0f;
-    typename ScaleTileIterator::Fragment fragment_dequant_scale_;
-    typename OutputTileIterator::Fragment fragment_C_;
-    typename OutputTileIterator::Fragment fragment_D_;
+  AlphaScaleElementType element_alpha_row_ = 1.0f;
+  AlphaScaleElementType element_alpha_col_ = 1.0f;
+  typename ScaleTileIterator::Fragment fragment_dequant_scale_;
+  typename OutputTileIterator::Fragment fragment_C_;
+  typename OutputTileIterator::Fragment fragment_D_;
 
-    ElementAccumulator beta_;
+  ElementAccumulator beta_;
 
-    int column_offset_;
+  int column_offset_;
 
-    MatrixCoord thread_offset_;
+  MatrixCoord thread_offset_;
 
-    public:
-    CUTLASS_DEVICE
-    DequantEpilogueVisitor(
-        Params const &params,
-        SharedStorage &shared_storage,  // NOLINT
-        cutlass::MatrixCoord const &problem_size,
-        int thread_idx,
-        int warp_idx,
-        int lane_idx,
-        typename OutputTileIterator::Params params_C,
-        typename OutputTileIterator::Params params_D,
-        typename ScaleTileIterator::Params params_dequant_scale,
-        typename OutputTileIterator::Element *ptr_C,
-        typename OutputTileIterator::Element *ptr_D,
-        AlphaScaleElementType *ptr_dequant_scale,
-        cutlass::MatrixCoord const &threadblock_offset =
-            cutlass::MatrixCoord(0, 0),
-        int column_offset = 0,
-        cutlass::MatrixCoord const &problem_size_real = cutlass::MatrixCoord(0,
-                                                                             0))
-        : params_(params),
-          shared_storage_(shared_storage),
-          extent_(problem_size),
-          elementwise_(params.elementwise),
-          ptr_dequant_scale_(ptr_dequant_scale),
-          iterator_dequant_scale_(params_dequant_scale,
-                                  ptr_dequant_scale,
-                                  problem_size,
-                                  thread_idx,
-                                  threadblock_offset),
-          iterator_C_(
-              params_C, ptr_C, problem_size, thread_idx, threadblock_offset),
-          iterator_D_(
-              params_D, ptr_D, problem_size, thread_idx, threadblock_offset),
-          extent_real_(problem_size_real) {
-        beta_ = (params.elementwise.beta_ptr ? *params.elementwise.beta_ptr
-                                             : params.elementwise.beta);
+ public:
+  CUTLASS_DEVICE
+  DequantEpilogueVisitor(
+      Params const &params,
+      SharedStorage &shared_storage,  // NOLINT
+      cutlass::MatrixCoord const &problem_size,
+      int thread_idx,
+      int warp_idx,
+      int lane_idx,
+      typename OutputTileIterator::Params params_C,
+      typename OutputTileIterator::Params params_D,
+      typename ScaleTileIterator::Params params_dequant_scale,
+      typename OutputTileIterator::Element *ptr_C,
+      typename OutputTileIterator::Element *ptr_D,
+      AlphaScaleElementType *ptr_dequant_scale,
+      cutlass::MatrixCoord const &threadblock_offset = cutlass::MatrixCoord(0,
+                                                                            0),
+      int column_offset = 0,
+      cutlass::MatrixCoord const &problem_size_real = cutlass::MatrixCoord(0,
+                                                                           0))
+      : params_(params),
+        shared_storage_(shared_storage),
+        extent_(problem_size),
+        elementwise_(params.elementwise),
+        ptr_dequant_scale_(ptr_dequant_scale),
+        iterator_dequant_scale_(params_dequant_scale,
+                                ptr_dequant_scale,
+                                problem_size,
+                                thread_idx,
+                                threadblock_offset),
+        iterator_C_(
+            params_C, ptr_C, problem_size, thread_idx, threadblock_offset),
+        iterator_D_(
+            params_D, ptr_D, problem_size, thread_idx, threadblock_offset),
+        extent_real_(problem_size_real) {
+    beta_ = (params.elementwise.beta_ptr ? *params.elementwise.beta_ptr
+                                         : params.elementwise.beta);
 
-        if (beta_ == ElementAccumulator()) {
-            iterator_C_.clear_mask();
-        }
+    if (beta_ == ElementAccumulator()) {
+      iterator_C_.clear_mask();
+    }
+  }
+
+  /// Helper to indicate split-K behavior
+  CUTLASS_DEVICE
+  void set_k_partition(
+      int split_k_index,     ///< Index of this threadblock within
+                             ///< split-K partitioned scheme
+      int split_k_slices) {  ///< Total number of split-K slices
+  }
+
+  /// Called to set the batch index
+  CUTLASS_DEVICE
+  void set_batch_index(int batch_idx) {}
+
+  /// Called at the start of the epilogue just before iterating over
+  /// accumulator slices
+  CUTLASS_DEVICE
+  void begin_epilogue() {
+    iterator_dequant_scale_.load(fragment_dequant_scale_);
+  }
+
+  /// Called at the start of one step before starting accumulator exchange
+  CUTLASS_DEVICE
+  void begin_step(int step_idx) {
+    fragment_D_.clear();
+    fragment_C_.clear();
+
+    iterator_C_.load(fragment_C_);
+    ++iterator_C_;
+  }
+
+  /// Called at the start of a row
+  CUTLASS_DEVICE
+  void begin_row(int row_idx) {
+    // Clear accumulators for max and sum when starting a whole row
+  }
+
+  /// Called after accumulators have been exchanged for each accumulator
+  /// vector
+  CUTLASS_DEVICE
+  void visit(int iter_idx,
+             int row_idx,
+             int column_idx,
+             int frag_idx,
+             AccumulatorFragment const &accum) {
+    NumericArrayConverter<ElementCompute,
+                          ElementAccumulator,
+                          kElementsPerAccess>
+        source_converter;
+
+    ComputeFragment result = source_converter(accum);
+
+    // printf("start_row:%d start_col:%d\niter_idx: %d, row_idx: %d,
+    // column_idx: %d, frag_idx: %d, i: %d, item: %f\n",
+    // iterator_D_.thread_start_row(), iterator_D_.thread_start_column(),
+    // iter_idx, row_idx, column_idx, frag_idx, i, result[i]);
+
+    ComputeFragment alpha_col =
+        reinterpret_cast<ComputeFragment *>(&fragment_dequant_scale_)[frag_idx];
+    result = per_token_channel_scale_accumulator_(
+        result, alpha_col, element_alpha_row_);
+
+    // printf("%d %e\n", accum[0], result[0]);
+    // scale_accumulator_(result, alpha_row_vector[0]); //TODO(mseznec)
+
+    // if (elementwise_.kScale ==
+    // cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling) {
+    //   result = source_converter(elementwise_(result));
+    // } else {
+    //   result = source_converter(elementwise_(result, source_vector));
+    // }
+
+    // Convert to the output
+    NumericArrayConverter<ElementOutput, ElementCompute, kElementsPerAccess>
+        output_converter;
+    OutputVector &output =
+        reinterpret_cast<OutputVector *>(&fragment_D_)[frag_idx];
+    output = output_converter(result);
+  }
+
+  /// Called at the end of a row
+  CUTLASS_DEVICE
+  void end_row(int row_idx) {
+    /* using ConvertSumOutput = cutlass::NumericConverter<ElementSum,
+     * ElementSoftmaxCompute>; */
+    /* using ConvertNormOutput = cutlass::NumericConverter<ElementNorm,
+     * ElementSoftmaxCompute>; */
+
+    /* ConvertSumOutput   convert_sum_output; */
+    /* ConvertNormOutput  convert_norm_output; */
+
+    /* Compute accumulate sum only in the last step */
+    /* accum_sum_ = warp_reduce_sum_(accum_sum_); */
+
+    /* bool is_first_thread_in_tile = ((threadIdx.x % kThreadsPerRow) == 0);
+     */
+    /* bool row_guard = thread_offset_.row() < extent_.row(); */
+    /* bool is_write_thread = row_guard && is_first_thread_in_tile; */
+
+    /* int block_batch = blockIdx.z; */
+
+    /* ElementNorm *curr_ptr_max = ptr_Max_ + thread_offset_.row() +
+     * column_offset_ + block_batch * params_.batch_stride_Max; */
+    /* ElementSum *curr_ptr_sum = ptr_Sum_ + thread_offset_.row() +
+     * column_offset_ + block_batch * params_.batch_stride_Sum; */
+
+    /* arch::global_store<ElementNorm, sizeof(ElementNorm)>( */
+    /*           convert_norm_output(accum_max_), */
+    /*           (void *)curr_ptr_max, */
+    /*           is_write_thread); */
+
+    /* arch::global_store<ElementSum, sizeof(ElementSum)>( */
+    /*           convert_sum_output(accum_sum_), */
+    /*           (void *)curr_ptr_sum, */
+    /*           is_write_thread); */
+
+    // Clear accumulators for max and sum when finishing a whole row
+    /* clear_accum_(); */
+  }
+
+  /// Called after all accumulator elements have been visited
+  CUTLASS_DEVICE
+  void end_step(int step_idx) {
+    iterator_D_.store(fragment_D_);
+    ++iterator_D_;
+  }
+
+  /// Called after all steps have been completed
+  CUTLASS_DEVICE
+  void end_epilogue() {}
+
+ private:
+  CUTLASS_DEVICE
+  ComputeFragment per_token_channel_scale_accumulator_(
+      ComputeFragment const &accum,
+      ComputeFragment const &scale_col,
+      AlphaScaleElementType const &scale_row) {
+    ComputeFragment result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < ComputeFragment::kElements; ++i) {
+      result[i] = accum[i] * (scale_col[i] * scale_row);
     }
 
-    /// Helper to indicate split-K behavior
-    CUTLASS_DEVICE
-    void set_k_partition(
-        int split_k_index,     ///< Index of this threadblock within
-                               ///< split-K partitioned scheme
-        int split_k_slices) {  ///< Total number of split-K slices
+    return result;
+  }
+
+  CUTLASS_DEVICE
+  ComputeFragment per_token_scale_accumulator_(
+      ComputeFragment const &accum,
+      AlphaScaleElementType const &scale_col,
+      AlphaScaleElementType const &scale_row) {
+    ComputeFragment result;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < ComputeFragment::kElements; ++i) {
+      result[i] = accum[i] * (scale_col * scale_row);
     }
 
-    /// Called to set the batch index
-    CUTLASS_DEVICE
-    void set_batch_index(int batch_idx) {}
-
-    /// Called at the start of the epilogue just before iterating over
-    /// accumulator slices
-    CUTLASS_DEVICE
-    void begin_epilogue() {
-        iterator_dequant_scale_.load(fragment_dequant_scale_);
-    }
-
-    /// Called at the start of one step before starting accumulator exchange
-    CUTLASS_DEVICE
-    void begin_step(int step_idx) {
-        fragment_D_.clear();
-        fragment_C_.clear();
-
-        iterator_C_.load(fragment_C_);
-        ++iterator_C_;
-    }
-
-    /// Called at the start of a row
-    CUTLASS_DEVICE
-    void begin_row(int row_idx) {
-        // Clear accumulators for max and sum when starting a whole row
-    }
-
-    /// Called after accumulators have been exchanged for each accumulator
-    /// vector
-    CUTLASS_DEVICE
-    void visit(int iter_idx,
-               int row_idx,
-               int column_idx,
-               int frag_idx,
-               AccumulatorFragment const &accum) {
-        NumericArrayConverter<ElementCompute,
-                              ElementAccumulator,
-                              kElementsPerAccess>
-            source_converter;
-
-        ComputeFragment result = source_converter(accum);
-
-        // printf("start_row:%d start_col:%d\niter_idx: %d, row_idx: %d,
-        // column_idx: %d, frag_idx: %d, i: %d, item: %f\n",
-        // iterator_D_.thread_start_row(), iterator_D_.thread_start_column(),
-        // iter_idx, row_idx, column_idx, frag_idx, i, result[i]);
-
-        ComputeFragment alpha_col = reinterpret_cast<ComputeFragment *>(
-            &fragment_dequant_scale_)[frag_idx];
-        result = per_token_channel_scale_accumulator_(
-            result, alpha_col, element_alpha_row_);
-
-        // printf("%d %e\n", accum[0], result[0]);
-        // scale_accumulator_(result, alpha_row_vector[0]); //TODO(mseznec)
-
-        // if (elementwise_.kScale ==
-        // cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling) {
-        //   result = source_converter(elementwise_(result));
-        // } else {
-        //   result = source_converter(elementwise_(result, source_vector));
-        // }
-
-        // Convert to the output
-        NumericArrayConverter<ElementOutput, ElementCompute, kElementsPerAccess>
-            output_converter;
-        OutputVector &output =
-            reinterpret_cast<OutputVector *>(&fragment_D_)[frag_idx];
-        output = output_converter(result);
-    }
-
-    /// Called at the end of a row
-    CUTLASS_DEVICE
-    void end_row(int row_idx) {
-        /* using ConvertSumOutput = cutlass::NumericConverter<ElementSum,
-         * ElementSoftmaxCompute>; */
-        /* using ConvertNormOutput = cutlass::NumericConverter<ElementNorm,
-         * ElementSoftmaxCompute>; */
-
-        /* ConvertSumOutput   convert_sum_output; */
-        /* ConvertNormOutput  convert_norm_output; */
-
-        /* Compute accumulate sum only in the last step */
-        /* accum_sum_ = warp_reduce_sum_(accum_sum_); */
-
-        /* bool is_first_thread_in_tile = ((threadIdx.x % kThreadsPerRow) == 0);
-         */
-        /* bool row_guard = thread_offset_.row() < extent_.row(); */
-        /* bool is_write_thread = row_guard && is_first_thread_in_tile; */
-
-        /* int block_batch = blockIdx.z; */
-
-        /* ElementNorm *curr_ptr_max = ptr_Max_ + thread_offset_.row() +
-         * column_offset_ + block_batch * params_.batch_stride_Max; */
-        /* ElementSum *curr_ptr_sum = ptr_Sum_ + thread_offset_.row() +
-         * column_offset_ + block_batch * params_.batch_stride_Sum; */
-
-        /* arch::global_store<ElementNorm, sizeof(ElementNorm)>( */
-        /*           convert_norm_output(accum_max_), */
-        /*           (void *)curr_ptr_max, */
-        /*           is_write_thread); */
-
-        /* arch::global_store<ElementSum, sizeof(ElementSum)>( */
-        /*           convert_sum_output(accum_sum_), */
-        /*           (void *)curr_ptr_sum, */
-        /*           is_write_thread); */
-
-        // Clear accumulators for max and sum when finishing a whole row
-        /* clear_accum_(); */
-    }
-
-    /// Called after all accumulator elements have been visited
-    CUTLASS_DEVICE
-    void end_step(int step_idx) {
-        iterator_D_.store(fragment_D_);
-        ++iterator_D_;
-    }
-
-    /// Called after all steps have been completed
-    CUTLASS_DEVICE
-    void end_epilogue() {}
-
-    private:
-    CUTLASS_DEVICE
-    ComputeFragment per_token_channel_scale_accumulator_(
-        ComputeFragment const &accum,
-        ComputeFragment const &scale_col,
-        AlphaScaleElementType const &scale_row) {
-        ComputeFragment result;
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < ComputeFragment::kElements; ++i) {
-            result[i] = accum[i] * (scale_col[i] * scale_row);
-        }
-
-        return result;
-    }
-
-    CUTLASS_DEVICE
-    ComputeFragment per_token_scale_accumulator_(
-        ComputeFragment const &accum,
-        AlphaScaleElementType const &scale_col,
-        AlphaScaleElementType const &scale_row) {
-        ComputeFragment result;
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < ComputeFragment::kElements; ++i) {
-            result[i] = accum[i] * (scale_col * scale_row);
-        }
-
-        return result;
-    }
+    return result;
+  }
 };
 
 /// Epilogue operator
@@ -1140,231 +1120,226 @@ class DequantEpilogueWithVisitor
                           WarpTileIterator_,
                           Padding_,
                           FragmentsPerPartition> {
-    public:
-    using Visitor = Visitor_;
+ public:
+  using Visitor = Visitor_;
 
-    using Base = EpilogueBase<Shape_,
-                              typename WarpMmaOperator_::Shape,
-                              PartitionsK,
-                              AccumulatorFragmentIterator_,
-                              WarpTileIterator_,
-                              Padding_,
-                              FragmentsPerPartition>;
+  using Base = EpilogueBase<Shape_,
+                            typename WarpMmaOperator_::Shape,
+                            PartitionsK,
+                            AccumulatorFragmentIterator_,
+                            WarpTileIterator_,
+                            Padding_,
+                            FragmentsPerPartition>;
 
-    using Shape = Shape_;
-    using WarpMmaOperator = WarpMmaOperator_;
-    static int const kPartitionsK = PartitionsK;
+  using Shape = Shape_;
+  using WarpMmaOperator = WarpMmaOperator_;
+  static int const kPartitionsK = PartitionsK;
 
-    using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
-    using WarpTileIterator = WarpTileIterator_;
-    using SharedLoadIterator = SharedLoadIterator_;
-    using Padding = Padding_;
+  using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
+  using WarpTileIterator = WarpTileIterator_;
+  using SharedLoadIterator = SharedLoadIterator_;
+  using Padding = Padding_;
 
-    using Layout = layout::RowMajor;
-    using LongIndex = typename Layout::LongIndex;
+  using Layout = layout::RowMajor;
+  using LongIndex = typename Layout::LongIndex;
 
-    /// The complete warp-level accumulator tile
-    using AccumulatorTile = typename Base::AccumulatorTile;
+  /// The complete warp-level accumulator tile
+  using AccumulatorTile = typename Base::AccumulatorTile;
 
-    /// Accumulator element
-    using ElementAccumulator = typename WarpTileIterator::Element;
+  /// Accumulator element
+  using ElementAccumulator = typename WarpTileIterator::Element;
 
-    /// Output access size
-    static int const kElementsPerAccess = Visitor::kElementsPerAccess;
+  /// Output access size
+  static int const kElementsPerAccess = Visitor::kElementsPerAccess;
 
-    /// Tensor reference to sync tensor
-    using SyncTensorRef =
-        typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
+  /// Tensor reference to sync tensor
+  using SyncTensorRef =
+      typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
 
-    /// Array type used by output functor
-    using AccumulatorAccessType =
-        Array<typename WarpTileIterator::Element, kElementsPerAccess>;
+  /// Array type used by output functor
+  using AccumulatorAccessType =
+      Array<typename WarpTileIterator::Element, kElementsPerAccess>;
 
-    /// Number of warps
-    using WarpCount = typename Base::WarpCount;
+  /// Number of warps
+  using WarpCount = typename Base::WarpCount;
 
-    static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1
-                                          ? Base::kFragmentsPerIteration
-                                          : kPartitionsK;
-    static int constexpr kSmemPointerOffset =
-        Base::SharedStorage::StorageShape::kCount / kSmemTiles;
+  static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1
+                                        ? Base::kFragmentsPerIteration
+                                        : kPartitionsK;
+  static int constexpr kSmemPointerOffset =
+      Base::SharedStorage::StorageShape::kCount / kSmemTiles;
 
-    using SharedStorage = typename Base::SharedStorage;
+  using SharedStorage = typename Base::SharedStorage;
 
-    private:
-    /// Loads fragment from shared memory aligned with output tensor
-    SharedLoadIterator shared_load_iterator_;
+ private:
+  /// Loads fragment from shared memory aligned with output tensor
+  SharedLoadIterator shared_load_iterator_;
 
-    public:
-    /// Constructor
-    CUTLASS_DEVICE
-    DequantEpilogueWithVisitor(
-        SharedStorage &shared_storage,  ///< Shared storage object.  // NOLINT
-        int thread_idx,  ///< ID of a thread within the threadblock
-        int warp_idx,    ///< ID of warp within threadblock
-        int lane_idx     ///< Id of thread within warp
-        )
-        : Base(shared_storage, thread_idx, warp_idx, lane_idx),
-          shared_load_iterator_(shared_storage.reference(), thread_idx) {}
+ public:
+  /// Constructor
+  CUTLASS_DEVICE
+  DequantEpilogueWithVisitor(
+      SharedStorage &shared_storage,  ///< Shared storage object.  // NOLINT
+      int thread_idx,                 ///< ID of a thread within the threadblock
+      int warp_idx,                   ///< ID of warp within threadblock
+      int lane_idx                    ///< Id of thread within warp
+      )
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
+        shared_load_iterator_(shared_storage.reference(), thread_idx) {}
 
-    /// Streams the result to global memory
-    CUTLASS_DEVICE
-    void operator()(
-        Visitor &visitor,  // NOLINT
-        AccumulatorTile const
-            &accumulators) {  ///< Threadblock tile coordinate in GEMM
-                              ///< (in units of threadblock tiles)
+  /// Streams the result to global memory
+  CUTLASS_DEVICE
+  void operator()(Visitor &visitor,  // NOLINT
+                  AccumulatorTile const
+                      &accumulators) {  ///< Threadblock tile coordinate in GEMM
+                                        ///< (in units of threadblock tiles)
 
-        visitor.begin_epilogue();
+    visitor.begin_epilogue();
 
-        //
-        // Iterator over warp-level accumulator fragment
-        //
+    //
+    // Iterator over warp-level accumulator fragment
+    //
 
-        AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
+    AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
 
-        //
-        // Iterate over accumulator tile
-        //
+    //
+    // Iterate over accumulator tile
+    //
 
 #pragma unroll(IterationsUnroll ? Visitor::kIterations : 1)
-        for (int iter_idx = 0; iter_idx < Visitor::kIterations; ++iter_idx) {
-            //
-            // Load the source
-            //
+    for (int iter_idx = 0; iter_idx < Visitor::kIterations; ++iter_idx) {
+      //
+      // Load the source
+      //
 
-            visitor.begin_step(iter_idx);
+      visitor.begin_step(iter_idx);
 
-            //
-            // Convert and store fragment
-            //
+      //
+      // Convert and store fragment
+      //
 
-            __syncthreads();
+      __syncthreads();
 
-            acc2smem_source_needed<cutlass::make_index_sequence<
-                Visitor::kIterations>>::push(iter_idx,
-                                             accum_fragment_iterator,
-                                             this->warp_tile_iterator_);
+      acc2smem_source_needed<cutlass::make_index_sequence<
+          Visitor::kIterations>>::push(iter_idx,
+                                       accum_fragment_iterator,
+                                       this->warp_tile_iterator_);
 
-            __syncthreads();
+      __syncthreads();
 
-            //
-            // Load fragments from shared memory
-            //
+      //
+      // Load fragments from shared memory
+      //
 
-            typename SharedLoadIterator::Fragment
-                aligned_accum_fragment[kPartitionsK];
+      typename SharedLoadIterator::Fragment
+          aligned_accum_fragment[kPartitionsK];
 
-            shared_load_iterator_.load(aligned_accum_fragment[0]);
+      shared_load_iterator_.load(aligned_accum_fragment[0]);
 
-            // If the number of k-slices is > 1 - perform a reduction amongst
-            // the k-slices if (kPartitionsK > 1) {
+      // If the number of k-slices is > 1 - perform a reduction amongst
+      // the k-slices if (kPartitionsK > 1) {
 
-            //   plus <typename SharedLoadIterator::Fragment> add_fragments;
+      //   plus <typename SharedLoadIterator::Fragment> add_fragments;
 
-            //   CUTLASS_PRAGMA_UNROLL
-            //   for ( int i = 1; i < kPartitionsK; ++i) {
-            //     shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
-            //     shared_load_iterator_.load(aligned_accum_fragment[i]);
-            //     aligned_accum_fragment[0] =
-            //     add_fragments(aligned_accum_fragment[0],
-            //     aligned_accum_fragment[i]);
-            //   }
+      //   CUTLASS_PRAGMA_UNROLL
+      //   for ( int i = 1; i < kPartitionsK; ++i) {
+      //     shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
+      //     shared_load_iterator_.load(aligned_accum_fragment[i]);
+      //     aligned_accum_fragment[0] =
+      //     add_fragments(aligned_accum_fragment[0],
+      //     aligned_accum_fragment[i]);
+      //   }
 
-            //   shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) *
-            //   kSmemPointerOffset);
-            // }
+      //   shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) *
+      //   kSmemPointerOffset);
+      // }
 
-            //
-            // Iterate over output fragments
-            //
+      //
+      // Iterate over output fragments
+      //
 
-            AccumulatorAccessType const *accum_frag_ptr =
-                reinterpret_cast<AccumulatorAccessType const *>(
-                    &aligned_accum_fragment[0]);
+      AccumulatorAccessType const *accum_frag_ptr =
+          reinterpret_cast<AccumulatorAccessType const *>(
+              &aligned_accum_fragment[0]);
 
-            int const kAccumulatorFragmentCount =
-                AccumulatorTile::kElements /
-                (Visitor::kIterations * AccumulatorAccessType::kElements);
+      int const kAccumulatorFragmentCount =
+          AccumulatorTile::kElements /
+          (Visitor::kIterations * AccumulatorAccessType::kElements);
 
-            CUTLASS_PRAGMA_UNROLL
-            for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
-                int row_idx =
-                    idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
-                int col_idx =
-                    idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
+      CUTLASS_PRAGMA_UNROLL
+      for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
+        int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
+        int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
 
-                // Start a new row of the output fragment
-                // if (!col_idx) {
-                //   visitor.begin_row(row_idx);
-                // }
+        // Start a new row of the output fragment
+        // if (!col_idx) {
+        //   visitor.begin_row(row_idx);
+        // }
 
-                visitor.visit(
-                    iter_idx, row_idx, col_idx, idx, accum_frag_ptr[idx]);
+        visitor.visit(iter_idx, row_idx, col_idx, idx, accum_frag_ptr[idx]);
 
-                // End the row of the output fragment
-                // if (col_idx + 1 ==
-                // SharedLoadIterator::ThreadMap::Iterations::kColumn) {
-                //   visitor.end_row(row_idx);
-                // }
-            }
+        // End the row of the output fragment
+        // if (col_idx + 1 ==
+        // SharedLoadIterator::ThreadMap::Iterations::kColumn) {
+        //   visitor.end_row(row_idx);
+        // }
+      }
 
-            //
-            // Conclude the step
-            //
+      //
+      // Conclude the step
+      //
 
-            visitor.end_step(iter_idx);
-        }
-
-        visitor.end_epilogue();
+      visitor.end_step(iter_idx);
     }
 
-    private:
-    template <class Seq>
-    struct acc2smem_source_needed;
+    visitor.end_epilogue();
+  }
 
-    template <size_t... Seq>
-    struct acc2smem_source_needed<cutlass::index_sequence<Seq...>> {
-        template <int Advance>
-        CUTLASS_DEVICE static void helper(
-            AccumulatorFragmentIterator accum_fragment_iterator,
-            WarpTileIterator &warp_tile_iterator) {  // NOLINT
-            CUTLASS_PRAGMA_UNROLL
-            for (int i = 0; i < Advance; i++) {
-                ++accum_fragment_iterator;
-            }
+ private:
+  template <class Seq>
+  struct acc2smem_source_needed;
 
-            typename AccumulatorFragmentIterator::Fragment accum_fragment;
-            accum_fragment_iterator.load(accum_fragment);
-            warp_tile_iterator.store(accum_fragment);
-        }
+  template <size_t... Seq>
+  struct acc2smem_source_needed<cutlass::index_sequence<Seq...>> {
+    template <int Advance>
+    CUTLASS_DEVICE static void helper(
+        AccumulatorFragmentIterator accum_fragment_iterator,
+        WarpTileIterator &warp_tile_iterator) {  // NOLINT
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Advance; i++) {
+        ++accum_fragment_iterator;
+      }
 
-        CUTLASS_DEVICE
-        static void push(size_t pos,
-                         AccumulatorFragmentIterator const &iterator_begin,
-                         WarpTileIterator &warp_tile_iterator) {  // NOLINT
-            int dummy[] = {
-                (pos == Seq) &&
-                (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
-        }
-    };
+      typename AccumulatorFragmentIterator::Fragment accum_fragment;
+      accum_fragment_iterator.load(accum_fragment);
+      warp_tile_iterator.store(accum_fragment);
+    }
+
+    CUTLASS_DEVICE
+    static void push(size_t pos,
+                     AccumulatorFragmentIterator const &iterator_begin,
+                     WarpTileIterator &warp_tile_iterator) {  // NOLINT
+      int dummy[] = {(pos == Seq) &&
+                     (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
+    }
+  };
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 /// Helper to create an EpilogueWithVisitor from an existing epilogue
 template <typename Visitor_, typename Existing_, bool IterationsUnroll = true>
 struct DequantEpilogueWithVisitorFromExistingEpilogue {
-    using Epilogue = DequantEpilogueWithVisitor<
-        Visitor_,
-        typename Existing_::Shape,
-        typename Existing_::WarpMmaOperator,
-        Existing_::kPartitionsK,
-        typename Existing_::AccumulatorFragmentIterator,
-        typename Existing_::WarpTileIterator,
-        typename Existing_::SharedLoadIterator,
-        typename Existing_::Padding,
-        Existing_::kFragmentsPerIteration,
-        IterationsUnroll>;
+  using Epilogue = DequantEpilogueWithVisitor<
+      Visitor_,
+      typename Existing_::Shape,
+      typename Existing_::WarpMmaOperator,
+      Existing_::kPartitionsK,
+      typename Existing_::AccumulatorFragmentIterator,
+      typename Existing_::WarpTileIterator,
+      typename Existing_::SharedLoadIterator,
+      typename Existing_::Padding,
+      Existing_::kFragmentsPerIteration,
+      IterationsUnroll>;
 };
 
 }  // namespace threadblock
@@ -1386,236 +1361,235 @@ template <typename ElementA_,
           int AlignmentA_ = 128 / cutlass::sizeof_bits<ElementA_>::value,
           int AlignmentB_ = 128 / cutlass::sizeof_bits<ElementB_>::value>
 class GemmDequant {
-    public:
-    ///////////////////////////////////////////////////////////////////////////////////////////////
+ public:
+  ///////////////////////////////////////////////////////////////////////////////////////////////
 
-    //
-    // Type definitions
-    //
+  //
+  // Type definitions
+  //
 
-    using ElementA = ElementA_;
-    using ElementB = ElementB_;
-    using ElementC = ElementC_;
-    using ElementCompute = ElementCompute_;
+  using ElementA = ElementA_;
+  using ElementB = ElementB_;
+  using ElementC = ElementC_;
+  using ElementCompute = ElementCompute_;
 
-    using LayoutA = LayoutA_;
-    using LayoutB = LayoutB_;
+  using LayoutA = LayoutA_;
+  using LayoutB = LayoutB_;
 
-    using EpilogueFunctorOp = EpilogueFunctorOp_;
+  using EpilogueFunctorOp = EpilogueFunctorOp_;
 
-    // These are mandatory layouts.
-    using LayoutC = cutlass::layout::RowMajor;
+  // These are mandatory layouts.
+  using LayoutC = cutlass::layout::RowMajor;
 
-    using TensorRefA = TensorRef<ElementA, LayoutA>;
-    using TensorRefB = TensorRef<ElementB, LayoutB>;
-    using TensorRefC = TensorRef<ElementC, LayoutC>;
+  using TensorRefA = TensorRef<ElementA, LayoutA>;
+  using TensorRefB = TensorRef<ElementB, LayoutB>;
+  using TensorRefC = TensorRef<ElementC, LayoutC>;
 
-    using ThreadblockShape = ThreadblockShape_;
-    using WarpShape = WarpShape_;
-    using InstructionShape = InstructionShape_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
 
-    using OperatorClass = OperatorClass_;
-    using ArchTag = ArchTag_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
 
-    static int const kStages = kStages_;
-    static int const AlignmentA = AlignmentA_;
-    static int const AlignmentB = AlignmentB_;
+  static int const kStages = kStages_;
+  static int const AlignmentA = AlignmentA_;
+  static int const AlignmentB = AlignmentB_;
 
-    using ThreadblockSwizzle =
-        cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>;
+  using ThreadblockSwizzle =
+      cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>;
 
-    ///////////////////////////////////////////////////////////////////////////////////////////////
+  ///////////////////////////////////////////////////////////////////////////////////////////////
 
-    // basic GEMM kernel
-    using DefaultGemmKernel =
-        typename cutlass::gemm::kernel::DefaultDequantGemmUniversal<
-            // using DefaultGemmKernel = typename
-            // cutlass::gemm::kernel::DefaultGemm<
-            ElementA,
-            LayoutA,
-            cutlass::ComplexTransform::kNone,
-            AlignmentA,
-            ElementB,
-            LayoutB,
-            cutlass::ComplexTransform::kNone,
-            AlignmentB,
-            ElementC,
-            LayoutC,
-            ElementCompute,
-            OperatorClass,
-            ArchTag,
-            ThreadblockShape,
-            WarpShape,
-            InstructionShape,
-            EpilogueFunctorOp,
-            ThreadblockSwizzle,
-            kStages,
-            cutlass::arch::OpMultiplyAddSaturate
-            // typename cutlass::gemm::device::DefaultGemmConfiguration<
-            //     OperatorClass, ArchTag, ElementA, ElementB, ElementC,
-            //     ElementCompute>::Operator,
-            // cutlass::gemm::SharedMemoryClearOption::kNone
-            >::GemmKernel;
+  // basic GEMM kernel
+  using DefaultGemmKernel =
+      typename cutlass::gemm::kernel::DefaultDequantGemmUniversal<
+          // using DefaultGemmKernel = typename
+          // cutlass::gemm::kernel::DefaultGemm<
+          ElementA,
+          LayoutA,
+          cutlass::ComplexTransform::kNone,
+          AlignmentA,
+          ElementB,
+          LayoutB,
+          cutlass::ComplexTransform::kNone,
+          AlignmentB,
+          ElementC,
+          LayoutC,
+          ElementCompute,
+          OperatorClass,
+          ArchTag,
+          ThreadblockShape,
+          WarpShape,
+          InstructionShape,
+          EpilogueFunctorOp,
+          ThreadblockSwizzle,
+          kStages,
+          cutlass::arch::OpMultiplyAddSaturate
+          // typename cutlass::gemm::device::DefaultGemmConfiguration<
+          //     OperatorClass, ArchTag, ElementA, ElementB, ElementC,
+          //     ElementCompute>::Operator,
+          // cutlass::gemm::SharedMemoryClearOption::kNone
+          >::GemmKernel;
 
-    ///////////////////////////////////////////////////////////////////////////////////////////////
-    using ElementEpilogueCompute = float;
-    using ElementEpilogueAcc = int32_t;
+  ///////////////////////////////////////////////////////////////////////////////////////////////
+  using ElementEpilogueCompute = float;
+  using ElementEpilogueAcc = int32_t;
 
-    using DequantScaleIterator =
-        cutlass::epilogue::threadblock::PredicatedTileIterator<
-            cutlass::epilogue::threadblock::OutputTileOptimalThreadMap<
-                typename DefaultGemmKernel::Epilogue::OutputTileIterator::
-                    ThreadMap::Shape,
-                typename DefaultGemmKernel::Epilogue::OutputTileIterator::
-                    ThreadMap::Count,
-                DefaultGemmKernel::Epilogue::OutputTileIterator::ThreadMap::
-                    kThreads,
-                DefaultGemmKernel::Epilogue::OutputTileIterator::
-                    kElementsPerAccess,
-                cutlass::sizeof_bits<ElementEpilogueCompute>::value>,
-            ElementEpilogueCompute>;
+  using DequantScaleIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIterator<
+          cutlass::epilogue::threadblock::OutputTileOptimalThreadMap<
+              typename DefaultGemmKernel::Epilogue::OutputTileIterator::
+                  ThreadMap::Shape,
+              typename DefaultGemmKernel::Epilogue::OutputTileIterator::
+                  ThreadMap::Count,
+              DefaultGemmKernel::Epilogue::OutputTileIterator::ThreadMap::
+                  kThreads,
+              DefaultGemmKernel::Epilogue::OutputTileIterator::
+                  kElementsPerAccess,
+              cutlass::sizeof_bits<ElementEpilogueCompute>::value>,
+          ElementEpilogueCompute>;
 
-    // Epilogue visitor
-    using EpilogueVisitor =
-        typename cutlass::epilogue::threadblock::DequantEpilogueVisitor<
-            ThreadblockShape,
-            DefaultGemmKernel::kThreadCount,
-            DequantScaleIterator,
-            typename DefaultGemmKernel::Epilogue::OutputTileIterator,
-            ElementEpilogueAcc,
-            ElementEpilogueCompute,
-            EpilogueFunctorOp>;
+  // Epilogue visitor
+  using EpilogueVisitor =
+      typename cutlass::epilogue::threadblock::DequantEpilogueVisitor<
+          ThreadblockShape,
+          DefaultGemmKernel::kThreadCount,
+          DequantScaleIterator,
+          typename DefaultGemmKernel::Epilogue::OutputTileIterator,
+          ElementEpilogueAcc,
+          ElementEpilogueCompute,
+          EpilogueFunctorOp>;
 
-    using ElementScale = typename EpilogueVisitor::AlphaScaleElementType;
-    using LayoutScale = cutlass::layout::RowMajor;
-    using TensorRefScale = TensorRef<ElementScale, LayoutScale>;
+  using ElementScale = typename EpilogueVisitor::AlphaScaleElementType;
+  using LayoutScale = cutlass::layout::RowMajor;
+  using TensorRefScale = TensorRef<ElementScale, LayoutScale>;
 
-    /// Epilogue
-    using Epilogue = typename cutlass::epilogue::threadblock::
-        DequantEpilogueWithVisitorFromExistingEpilogue<
-            EpilogueVisitor,
-            typename DefaultGemmKernel::Epilogue>::Epilogue;
+  /// Epilogue
+  using Epilogue = typename cutlass::epilogue::threadblock::
+      DequantEpilogueWithVisitorFromExistingEpilogue<
+          EpilogueVisitor,
+          typename DefaultGemmKernel::Epilogue>::Epilogue;
 
-    // GEMM
-    using GemmKernel = gemm::kernel::GemmWithEpilogueVisitorDequant<
-        typename DefaultGemmKernel::Mma,
-        Epilogue,
-        ThreadblockSwizzle>;
+  // GEMM
+  using GemmKernel = gemm::kernel::GemmWithEpilogueVisitorDequant<
+      typename DefaultGemmKernel::Mma,
+      Epilogue,
+      ThreadblockSwizzle>;
 
-    public:
-    /// Arguments class
-    struct Arguments {
-        typename GemmKernel::Arguments gemm;
-        cutlass::gemm::GemmCoord extend;
-
-        //
-        // Methods
-        //
-        Arguments() : gemm(), extend() {}
-
-        Arguments(cutlass::gemm::GemmCoord problem_size,
-                  TensorRefA ref_A_,
-                  TensorRefB ref_B_,
-                  TensorRefC ref_C_,
-                  TensorRefC ref_D_,
-                  TensorRefScale ref_scale_,
-                  typename EpilogueFunctorOp::Params linear_scaling)
-            : gemm(cutlass::gemm::GemmUniversalMode::kGemm,
-                   problem_size,
-                   ref_A_,
-                   ref_B_,
-                   ref_C_,
-                   ref_D_,
-                   ref_scale_,
-                   typename EpilogueVisitor::Arguments(linear_scaling)),
-              extend(problem_size) {}
-    };
-
-    struct Params {
-        typename GemmKernel::Params gemm;
-        MatrixCoord extend;
-        //
-        // Methods
-        //
-        Params() {}
-
-        explicit Params(Arguments const &args)
-            : gemm(args.gemm),
-              extend(MatrixCoord(args.extend.m(), args.extend.n())) {}
-    };
-
-    public:
-    // Gemm
+ public:
+  /// Arguments class
+  struct Arguments {
+    typename GemmKernel::Arguments gemm;
+    cutlass::gemm::GemmCoord extend;
 
     //
     // Methods
     //
+    Arguments() : gemm(), extend() {}
 
-    private:
-    Params params_;
+    Arguments(cutlass::gemm::GemmCoord problem_size,
+              TensorRefA ref_A_,
+              TensorRefB ref_B_,
+              TensorRefC ref_C_,
+              TensorRefC ref_D_,
+              TensorRefScale ref_scale_,
+              typename EpilogueFunctorOp::Params linear_scaling)
+        : gemm(cutlass::gemm::GemmUniversalMode::kGemm,
+               problem_size,
+               ref_A_,
+               ref_B_,
+               ref_C_,
+               ref_D_,
+               ref_scale_,
+               typename EpilogueVisitor::Arguments(linear_scaling)),
+          extend(problem_size) {}
+  };
 
-    public:
-    /// Ctor
-    GemmDequant() {}
+  struct Params {
+    typename GemmKernel::Params gemm;
+    MatrixCoord extend;
+    //
+    // Methods
+    //
+    Params() {}
 
-    /// Initialize
-    Status initialize(Arguments const &args) {
-        params_ = Params(args);
+    explicit Params(Arguments const &args)
+        : gemm(args.gemm),
+          extend(MatrixCoord(args.extend.m(), args.extend.n())) {}
+  };
 
-        return cutlass::Status::kSuccess;
+ public:
+  // Gemm
+
+  //
+  // Methods
+  //
+
+ private:
+  Params params_;
+
+ public:
+  /// Ctor
+  GemmDequant() {}
+
+  /// Initialize
+  Status initialize(Arguments const &args) {
+    params_ = Params(args);
+
+    return cutlass::Status::kSuccess;
+  }
+
+  /// Run
+  Status run(cudaStream_t stream) {
+    //
+    // Launch the GEMM + max kernel
+    //
+
+    dim3 gemm_grid =
+        ThreadblockSwizzle().get_grid_shape(params_.gemm.grid_tiled_shape);
+    dim3 gemm_block(GemmKernel::kThreadCount, 1, 1);
+
+    int gemm_smem_size =
+        static_cast<int>(sizeof(typename GemmKernel::SharedStorage));
+
+    cudaError_t result;
+
+    if (gemm_smem_size >= (48 << 10)) {
+      result = cudaFuncSetAttribute(cutlass::Kernel<GemmKernel>,
+                                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                    gemm_smem_size);
+
+      if (result != cudaSuccess) {
+        return Status::kErrorInternal;
+      }
     }
 
-    /// Run
-    Status run(cudaStream_t stream) {
-        //
-        // Launch the GEMM + max kernel
-        //
+    cutlass::Kernel<GemmKernel>
+        <<<gemm_grid, gemm_block, gemm_smem_size, stream>>>(params_.gemm);
 
-        dim3 gemm_grid =
-            ThreadblockSwizzle().get_grid_shape(params_.gemm.grid_tiled_shape);
-        dim3 gemm_block(GemmKernel::kThreadCount, 1, 1);
+    result = cudaGetLastError();
 
-        int gemm_smem_size =
-            static_cast<int>(sizeof(typename GemmKernel::SharedStorage));
-
-        cudaError_t result;
-
-        if (gemm_smem_size >= (48 << 10)) {
-            result = cudaFuncSetAttribute(
-                cutlass::Kernel<GemmKernel>,
-                cudaFuncAttributeMaxDynamicSharedMemorySize,
-                gemm_smem_size);
-
-            if (result != cudaSuccess) {
-                return Status::kErrorInternal;
-            }
-        }
-
-        cutlass::Kernel<GemmKernel>
-            <<<gemm_grid, gemm_block, gemm_smem_size, stream>>>(params_.gemm);
-
-        result = cudaGetLastError();
-
-        if (result != cudaSuccess) {
-            std::cerr << "gemm_with_dequant kernel error: "
-                      << cudaGetErrorString(result) << std::endl;
-            return cutlass::Status::kErrorInternal;
-        }
-
-        return cutlass::Status::kSuccess;
+    if (result != cudaSuccess) {
+      std::cerr << "gemm_with_dequant kernel error: "
+                << cudaGetErrorString(result) << std::endl;
+      return cutlass::Status::kErrorInternal;
     }
 
-    /// Function call operator
-    Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
+    return cutlass::Status::kSuccess;
+  }
+
+  /// Function call operator
+  Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
 };
 
 }  // namespace cutlass
 
 typedef struct {
-  void const* act;
-  void const* weight;
-  void const* scale;
-  void* output;
+  void const *act;
+  void const *weight;
+  void const *scale;
+  void *output;
   int32_t m, n, k;
   cudaStream_t stream;
 } GemmDequantParams;
diff --git a/custom_ops/gpu_ops/machete/machete_collective_builder.cuh b/custom_ops/gpu_ops/machete/machete_collective_builder.cuh
index c9a0cc4e9b..3b39294313 100644
--- a/custom_ops/gpu_ops/machete/machete_collective_builder.cuh
+++ b/custom_ops/gpu_ops/machete/machete_collective_builder.cuh
@@ -8,24 +8,48 @@ using namespace cute;
 
 struct MacheteKernelTag {};
 
-template <class ElementPairA_, class GmemLayoutA_, int AlignmentA,
-          class ElementPairB_, class GmemLayoutB_, int AlignmentB,
-          class ElementAccumulator, class TileShape_MNK, class ClusterShape_MNK,
-          class StageCountType, class KernelScheduleType>
+template <class ElementPairA_,
+          class GmemLayoutA_,
+          int AlignmentA,
+          class ElementPairB_,
+          class GmemLayoutB_,
+          int AlignmentB,
+          class ElementAccumulator,
+          class TileShape_MNK,
+          class ClusterShape_MNK,
+          class StageCountType,
+          class KernelScheduleType>
 struct MacheteCollectiveBuilder<
-    MacheteKernelTag, arch::Sm90, arch::OpClassTensorOp, ElementPairA_,
-    GmemLayoutA_, AlignmentA, ElementPairB_, GmemLayoutB_, AlignmentB,
-    ElementAccumulator, TileShape_MNK, ClusterShape_MNK, StageCountType,
+    MacheteKernelTag,
+    arch::Sm90,
+    arch::OpClassTensorOp,
+    ElementPairA_,
+    GmemLayoutA_,
+    AlignmentA,
+    ElementPairB_,
+    GmemLayoutB_,
+    AlignmentB,
+    ElementAccumulator,
+    TileShape_MNK,
+    ClusterShape_MNK,
+    StageCountType,
     KernelScheduleType,
     cute::enable_if_t<(
         cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecialized> ||
         cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedPingpong> ||
         cute::is_same_v<KernelScheduleType,
                         KernelTmaWarpSpecializedCooperative>)>> {
-  using CollectiveOp = machete::MacheteCollectiveMma<
-      ElementPairA_, GmemLayoutA_, AlignmentA, ElementPairB_, GmemLayoutB_,
-      AlignmentB, ElementAccumulator, TileShape_MNK, ClusterShape_MNK,
-      StageCountType, KernelScheduleType>;
+  using CollectiveOp = machete::MacheteCollectiveMma<ElementPairA_,
+                                                     GmemLayoutA_,
+                                                     AlignmentA,
+                                                     ElementPairB_,
+                                                     GmemLayoutB_,
+                                                     AlignmentB,
+                                                     ElementAccumulator,
+                                                     TileShape_MNK,
+                                                     ClusterShape_MNK,
+                                                     StageCountType,
+                                                     KernelScheduleType>;
 };
 
 };  // namespace cutlass::gemm::collective
diff --git a/custom_ops/gpu_ops/machete/machete_mainloop.cuh b/custom_ops/gpu_ops/machete/machete_mainloop.cuh
index 9df7adf905..18a338b8ba 100644
--- a/custom_ops/gpu_ops/machete/machete_mainloop.cuh
+++ b/custom_ops/gpu_ops/machete/machete_mainloop.cuh
@@ -56,10 +56,16 @@ using namespace cutlass::gemm;
 using namespace cutlass::gemm::collective;
 using namespace cutlass::gemm::collective::detail;
 
-template <class ElementATuple_, class GmemLayoutA, int AlignmentA,
-          class ElementB_, class GmemLayoutB, int AlignmentB,
-          class ElementAccumulator_, class TileShape_MNK,
-          class ClusterShape_MNK, class StageCountType,
+template <class ElementATuple_,
+          class GmemLayoutA,
+          int AlignmentA,
+          class ElementB_,
+          class GmemLayoutB,
+          int AlignmentB,
+          class ElementAccumulator_,
+          class TileShape_MNK,
+          class ClusterShape_MNK,
+          class StageCountType,
           class KernelScheduleType>
 struct MacheteCollectiveMma {
   using Schedule = KernelScheduleType;
@@ -100,7 +106,8 @@ struct MacheteCollectiveMma {
   using ElementMma = ElementB;
   using ElementATuple =
       cute::conditional_t<!cute::is_tuple<ElementATuple_>::value,
-                          cute::tuple<ElementA>, ElementATuple_>;
+                          cute::tuple<ElementA>,
+                          ElementATuple_>;
 
   static constexpr cute::GMMA::Major GmmaMajorA =
       gmma_rs_tag_to_major_A<layout::RowMajor>();
@@ -111,11 +118,16 @@ struct MacheteCollectiveMma {
   // instructions so we use 2 atoms along the M dim (one for each warpgroup)
   using AtomLayoutMNK = cute::conditional_t<
       cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedCooperative>,
-      Layout<Shape<_2, _1, _1>>, Layout<Shape<_1, _1, _1>>>;
+      Layout<Shape<_2, _1, _1>>,
+      Layout<Shape<_1, _1, _1>>>;
 
   using TiledMma = decltype(cute::make_tiled_mma(
-      cute::GMMA::rs_op_selector<ElementMma, ElementMma, ElementAccumulator,
-                                 TileShape_MNK, GMMA::Major::K, GmmaMajorB>(),
+      cute::GMMA::rs_op_selector<ElementMma,
+                                 ElementMma,
+                                 ElementAccumulator,
+                                 TileShape_MNK,
+                                 GMMA::Major::K,
+                                 GmmaMajorB>(),
       AtomLayoutMNK{}));
 
  private:
@@ -148,8 +160,12 @@ struct MacheteCollectiveMma {
  public:
   static constexpr int PipelineStages =
       compute_stage_count_or_override_single_affine_transformed_input<
-          sm90_smem_capacity_bytes, ElementA, ElementB, ElementScale,
-          ElementZero, TileShape_MNK>(StageCountType{});
+          sm90_smem_capacity_bytes,
+          ElementA,
+          ElementB,
+          ElementScale,
+          ElementZero,
+          TileShape_MNK>(StageCountType{});
 
   struct DispatchPolicy {
     constexpr static int Stages = PipelineStages;
@@ -164,15 +180,18 @@ struct MacheteCollectiveMma {
 
   // ((T, V), (BlocksM, BlocksK), pipe) -> offset
   using SmemLayoutA = decltype(GmemLayoutA::TVbNbKL_to_offset(
-      make_shape(size<0>(TileShape_MNK{}), size<2>(TileShape_MNK{}),
+      make_shape(size<0>(TileShape_MNK{}),
+                 size<2>(TileShape_MNK{}),
                  Int<DispatchPolicy::Stages>{})));
 
   using SmemLayoutACopy = decltype(GmemLayoutA::TVbNbKL_to_offset_copy(
-      make_shape(size<0>(TileShape_MNK{}), size<2>(TileShape_MNK{}),
+      make_shape(size<0>(TileShape_MNK{}),
+                 size<2>(TileShape_MNK{}),
                  Int<DispatchPolicy::Stages>{})));
 
   using SmemLayoutAtomARowMajor =
-      decltype(rs_smem_selector<GmmaMajorA, ElementA,
+      decltype(rs_smem_selector<GmmaMajorA,
+                                ElementA,
                                 decltype(cute::get<0>(TileShape_MNK{})),
                                 decltype(cute::get<2>(TileShape_MNK{}))>());
 
@@ -180,7 +199,8 @@ struct MacheteCollectiveMma {
       Shape<decltype(cute::shape<0>(SmemLayoutAtomARowMajor{})), cute::Int<1>>>;
 
   using SmemLayoutAtomB =
-      decltype(rs_smem_selector<GmmaMajorB, ElementB,
+      decltype(rs_smem_selector<GmmaMajorB,
+                                ElementB,
                                 decltype(cute::get<1>(TileShape_MNK{})),
                                 decltype(cute::get<2>(TileShape_MNK{}))>());
 
@@ -192,7 +212,10 @@ struct MacheteCollectiveMma {
   //
   static_assert(is_static<TileShape_MNK>::value);
   static_assert(is_static<ClusterShape_MNK>::value);
-  static_assert(is_aligned<ElementA, AlignmentA, ElementB, AlignmentB,
+  static_assert(is_aligned<ElementA,
+                           AlignmentA,
+                           ElementB,
+                           AlignmentB,
                            tma_alignment_bytes>(),
                 "Should meet TMA alignment requirement\n");
 #ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED
@@ -226,7 +249,8 @@ struct MacheteCollectiveMma {
   // code to compile when the scale is void.
   using NonVoidStrideScale =
       cute::conditional_t<cute::is_void_v<StrideScale>,
-                          cute::Stride<_1, int64_t, int64_t>, StrideScale>;
+                          cute::Stride<_1, int64_t, int64_t>,
+                          StrideScale>;
 
   static_assert((cutlass::gemm::detail::is_k_major<StrideA>()),
                 "The transformed matrix (A) must be K-major.");
@@ -254,10 +278,12 @@ struct MacheteCollectiveMma {
   static constexpr bool ConvertF32toTF32A = cute::is_same_v<float, ElementA>;
   static constexpr bool ConvertF32toTF32B = cute::is_same_v<float, ElementB>;
   using InternalElementA =
-      cute::conditional_t<ConvertF32toTF32A, tfloat32_t,
+      cute::conditional_t<ConvertF32toTF32A,
+                          tfloat32_t,
                           uint_bit_t<sizeof_bits_v<ElementA>>>;
   using InternalElementB =
-      cute::conditional_t<ConvertF32toTF32B, tfloat32_t,
+      cute::conditional_t<ConvertF32toTF32B,
+                          tfloat32_t,
                           uint_bit_t<sizeof_bits_v<ElementB>>>;
 
   using TransformA = cute::identity;
@@ -295,16 +321,19 @@ struct MacheteCollectiveMma {
   // Tile along modes in a way that maximizes the TMA box size
   using SmemLayoutB = decltype(tile_to_shape(
       SmemLayoutAtomB{},
-      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}),
+      make_shape(shape<1>(TileShape{}),
+                 shape<2>(TileShape{}),
                  Int<DispatchPolicy::Stages>{}),
       conditional_t<::cutlass::gemm::detail::is_major<0, StrideB>(),
-                    Step<_2, _1, _3>, Step<_1, _2, _3>>{}));
+                    Step<_2, _1, _3>,
+                    Step<_1, _2, _3>>{}));
 
   // It is assumed that the scales and zero-points share the same smem layout
-  using SmemLayoutScale = decltype(tile_to_shape(
-      SmemLayoutAtomScale{},
-      make_shape(shape<0>(ScaleTileShape{}), shape<1>(ScaleTileShape{}),
-                 Int<PipelineStages>{})));
+  using SmemLayoutScale =
+      decltype(tile_to_shape(SmemLayoutAtomScale{},
+                             make_shape(shape<0>(ScaleTileShape{}),
+                                        shape<1>(ScaleTileShape{}),
+                                        Int<PipelineStages>{})));
 
   // If A mn-layout and B mn-layout, transposing B matrix since WGMMA is k-major
   // only (e.g. tf32, fp32, fp8, int8).
@@ -331,10 +360,12 @@ struct MacheteCollectiveMma {
 
   using GmmaSmemLayoutB = decltype(tile_to_shape(
       SmemLayoutAtomB{},
-      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}),
+      make_shape(shape<1>(TileShape{}),
+                 shape<2>(TileShape{}),
                  Int<DispatchPolicy::Stages>{}),
       conditional_t<::cutlass::gemm::detail::is_major<0, StrideB>(),
-                    Step<_2, _1, _3>, Step<_1, _2, _3>>{}));
+                    Step<_2, _1, _3>,
+                    Step<_1, _2, _3>>{}));
 
   // These two restrictions are related, so we place the assertions together.
   // To relax them, we need to handle loading more than 1 row of scales for
@@ -436,25 +467,31 @@ struct MacheteCollectiveMma {
 
   using BTensor = decltype(make_tensor(
       get_logical_ptr(static_cast<InternalElementB const*>(nullptr)),
-      repeat_like(StrideB{}, int32_t(0)), StrideB{}));
+      repeat_like(StrideB{}, int32_t(0)),
+      StrideB{}));
   using ScaleTensor = decltype(make_tensor(
       get_logical_ptr(static_cast<NonVoidElementScale const*>(nullptr)),
-      repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{}));
+      repeat_like(NonVoidStrideScale{}, int32_t(0)),
+      NonVoidStrideScale{}));
 
   using ZeroTensor = decltype(make_tensor(
       get_logical_ptr(static_cast<NonVoidElementZero const*>(nullptr)),
-      repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{}));
+      repeat_like(NonVoidStrideScale{}, int32_t(0)),
+      NonVoidStrideScale{}));
 
   static constexpr auto make_tma_copy_A(ATensor tensor_a = ATensor{}) {
     return make_tma_copy<TmaElementA>(
-        GmemTiledCopyA{}, tensor_a, SmemLayoutACopy{}(_, _, cute::Int<0>{}),
+        GmemTiledCopyA{},
+        tensor_a,
+        SmemLayoutACopy{}(_, _, cute::Int<0>{}),
         shape(SmemLayoutACopy{}(_, _, cute::Int<0>{})),
         size<1>(ClusterShape{}));  // mcast along N mode for this M load, if any
   }
 
   static constexpr auto make_tma_copy_scale(
       ScaleTensor tensor_scale = ScaleTensor{}) {
-    return make_tma_copy(GmemTiledCopyScale{}, tensor_scale,
+    return make_tma_copy(GmemTiledCopyScale{},
+                         tensor_scale,
                          SmemLayoutScale{}(_, _, cute::Int<0>{}),
                          ScaleTileShape{},
                          _1{});  // mcast along N mode for this M load, if any
@@ -462,7 +499,8 @@ struct MacheteCollectiveMma {
 
   static constexpr auto make_tma_copy_zero(
       ZeroTensor tensor_zero = ZeroTensor{}) {
-    return make_tma_copy(GmemTiledCopyScale{}, tensor_zero,
+    return make_tma_copy(GmemTiledCopyScale{},
+                         tensor_zero,
                          SmemLayoutScale{}(_, _, cute::Int<0>{}),
                          ScaleTileShape{},
                          _1{});  // mcast along N mode for this M load, if any
@@ -470,7 +508,9 @@ struct MacheteCollectiveMma {
 
   static constexpr auto make_tma_copy_B(BTensor tensor_b = BTensor{}) {
     return make_tma_copy(
-        GmemTiledCopyB{}, tensor_b, SmemLayoutB{}(_, _, cute::Int<0>{}),
+        GmemTiledCopyB{},
+        tensor_b,
+        SmemLayoutB{}(_, _, cute::Int<0>{}),
         make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
         size<0>(ClusterShape{}));  // mcast along M mode for this N load, if any
   }
@@ -560,7 +600,8 @@ struct MacheteCollectiveMma {
   //  to handle the prepacked layout
   template <class ProblemShape>
   static constexpr Params to_underlying_arguments(
-      ProblemShape const& problem_shape, Arguments const& args,
+      ProblemShape const& problem_shape,
+      Arguments const& args,
       void* workspace) {
     (void)workspace;
 
@@ -607,8 +648,12 @@ struct MacheteCollectiveMma {
                   KernelConversionMode == ConversionMode::ConvertAndScale ||
                   KernelConversionMode ==
                       ConversionMode::ConvertAndScaleWithZero) {
-      return {tma_load_a,    tma_load_b, tma_load_scale,
-              tma_load_zero, scale_k,    group_size};
+      return {tma_load_a,
+              tma_load_b,
+              tma_load_scale,
+              tma_load_zero,
+              scale_k,
+              group_size};
     } else {
       static_assert(cutlass::detail::dependent_false<KernelSchedule>,
                     "Conversion mode not handled in to_underlying_arguments.");
@@ -940,8 +985,10 @@ struct MacheteCollectiveMma {
   // Consumer Perspective
   template <class FrgTensorC>
   CUTLASS_DEVICE void mma(MainloopPipeline pipeline,
-                          PipelineState smem_pipe_read, FrgTensorC& accum,
-                          int k_tile_count, int thread_idx,
+                          PipelineState smem_pipe_read,
+                          FrgTensorC& accum,
+                          int k_tile_count,
+                          int thread_idx,
                           TensorStorage& shared_tensors,
                           Params const& mainloop_params) {
     static_assert(is_rmem<FrgTensorC>::value,
@@ -971,8 +1018,10 @@ struct MacheteCollectiveMma {
     // which can be thought of as:
     //   (T, MMA, (MMA_M, MMA_K), pipe) -> offset
     auto constexpr smem_A_mma_ =
-        make_layout(get<0, 0>(smem_A), get<0, 1, 0>(smem_A),
-                    zip(get<0, 1, 1>(smem_A), get<1>(smem_A)), get<2>(smem_A));
+        make_layout(get<0, 0>(smem_A),
+                    get<0, 1, 0>(smem_A),
+                    zip(get<0, 1, 1>(smem_A), get<1>(smem_A)),
+                    get<2>(smem_A));
     // flatten to:
     //   (T, MMA, MMA_M, MMA_K, pipe) -> offset
     auto constexpr smem_A_mma = smem_A_mma_(_, _, make_coord(_, _), _);
@@ -1007,7 +1056,8 @@ struct MacheteCollectiveMma {
 
     auto load_A_to_registers = [&](int read_stage) {
       copy(create_auto_vectorizing_copy<ElementA, decltype(A_CPY_VEC)>(),
-           tCsA(_, _, _, read_stage), tCrA_load(_, _, _));
+           tCsA(_, _, _, read_stage),
+           tCrA_load(_, _, _));
     };
 
     // Partition of thread -> shared and thread -> RF
@@ -1028,10 +1078,11 @@ struct MacheteCollectiveMma {
     auto convert_A = [&, a_vec = Int<CONVERSION_WIDTH>{}](int k_block,
                                                           int read_stage) {
       load_extra_info_to_registers(partitioned_extra_info,
-                                   copy_partitions_extra_info, k_block,
+                                   copy_partitions_extra_info,
+                                   k_block,
                                    read_stage);
-      transform_A_kblock(tCrA_load, a_vec, tCrA_mma, partitioned_extra_info,
-                         k_block);
+      transform_A_kblock(
+          tCrA_load, a_vec, tCrA_mma, partitioned_extra_info, k_block);
     };
 
     // We release buffers to producer warps(dma load) with some mmas in flight
@@ -1065,8 +1116,10 @@ struct MacheteCollectiveMma {
         }
         warpgroup_arrive();
         // (V,M) x (V,N) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA_mma(_, _, k_block),
-                   tCrB(_, _, k_block, read_stage), accum);
+        cute::gemm(tiled_mma,
+                   tCrA_mma(_, _, k_block),
+                   tCrB(_, _, k_block, read_stage),
+                   accum);
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
         warpgroup_commit_batch();
       }
@@ -1103,8 +1156,10 @@ struct MacheteCollectiveMma {
       for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
         warpgroup_arrive();
         // (V,M) x (V,N) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA_mma(_, _, k_block),
-                   tCrB(_, _, k_block, read_stage), accum);
+        cute::gemm(tiled_mma,
+                   tCrA_mma(_, _, k_block),
+                   tCrB(_, _, k_block, read_stage),
+                   accum);
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
         warpgroup_commit_batch();
 
@@ -1149,8 +1204,10 @@ struct MacheteCollectiveMma {
       for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
         warpgroup_arrive();
         // (V,M) x (V,N) => (V,M,N)
-        cute::gemm(tiled_mma, tCrA_mma(_, _, k_block),
-                   tCrB(_, _, k_block, read_stage), accum);
+        cute::gemm(tiled_mma,
+                   tCrA_mma(_, _, k_block),
+                   tCrB(_, _, k_block, read_stage),
+                   accum);
         tiled_mma.accumulate_ = GMMA::ScaleOut::One;
         warpgroup_commit_batch();
         warpgroup_wait<K_BLOCK_MAX - 1>();
@@ -1324,7 +1381,8 @@ struct MacheteCollectiveMma {
   template <class... Ts, class... Us>
   CUTLASS_DEVICE void load_extra_info_to_registers(
       cute::tuple<Ts...> const& partitioned_mma_extra_info,
-      cute::tuple<Us...> const& tiled_copy_and_views, int k_block,
+      cute::tuple<Us...> const& tiled_copy_and_views,
+      int k_block,
       int read_stage) {
     if (k_block == 0) {
       // We are starting a new k-tile so copy the scale
@@ -1334,7 +1392,8 @@ struct MacheteCollectiveMma {
         auto smem_tiled_copy_S = cute::get<0>(tiled_copy_and_views);
         auto tCrS_copy_view = cute::get<1>(tiled_copy_and_views);
         auto tCsS = cute::get<0>(partitioned_mma_extra_info);
-        copy(smem_tiled_copy_S, tCsS(_, _, k_block, read_stage),
+        copy(smem_tiled_copy_S,
+             tCsS(_, _, k_block, read_stage),
              tCrS_copy_view(_, _, k_block));
         if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
           // Nothing extra to do
@@ -1342,7 +1401,8 @@ struct MacheteCollectiveMma {
                              ConversionMode::ConvertAndScaleWithZero) {
           auto tCsZ = cute::get<2>(partitioned_mma_extra_info);
           auto tCrZ_copy_view = cute::get<2>(tiled_copy_and_views);
-          copy(smem_tiled_copy_S, tCsZ(_, _, k_block, read_stage),
+          copy(smem_tiled_copy_S,
+               tCsZ(_, _, k_block, read_stage),
                tCrZ_copy_view(_, _, k_block));
         } else {
           static_assert(cutlass::detail::dependent_false<KernelSchedule>,
@@ -1362,8 +1422,10 @@ struct MacheteCollectiveMma {
   // Utilities to transform A.
   template <class TCrA_load, int VectorWidthA, class TCrA_mma, class... Ts>
   CUTLASS_DEVICE void transform_A_kblock(
-      TCrA_load const& tCrA_load, cute::Int<VectorWidthA> vec_A,
-      TCrA_mma& tCrA_mma, cute::tuple<Ts...> const& partitioned_extra_info,
+      TCrA_load const& tCrA_load,
+      cute::Int<VectorWidthA> vec_A,
+      TCrA_mma& tCrA_mma,
+      cute::tuple<Ts...> const& partitioned_extra_info,
       int const k_block) {
     auto in = tCrA_load(_, _, k_block);
     auto out = tCrA_mma(_, _, k_block);
@@ -1385,12 +1447,13 @@ struct MacheteCollectiveMma {
       // float, which nvcc will not optimize to using vectorized fma
       // instructions (i.e. hfma.bf16_v2)
       if constexpr (std::is_same_v<ElementScale, cutlass::bfloat16_t>) {
-        cute::transform(
-            recast<nv_bfloat16>(converted_inputs), recast<nv_bfloat16>(scales),
-            recast<nv_bfloat16>(converted_inputs), cute::multiplies{});
-      } else {
-        cute::transform(converted_inputs, scales, converted_inputs,
+        cute::transform(recast<nv_bfloat16>(converted_inputs),
+                        recast<nv_bfloat16>(scales),
+                        recast<nv_bfloat16>(converted_inputs),
                         cute::multiplies{});
+      } else {
+        cute::transform(
+            converted_inputs, scales, converted_inputs, cute::multiplies{});
       }
 
       // Apply zeros if required
@@ -1403,9 +1466,12 @@ struct MacheteCollectiveMma {
         if constexpr (std::is_same_v<ElementScale, cutlass::bfloat16_t>) {
           cute::transform(recast<nv_bfloat16>(converted_inputs),
                           recast<nv_bfloat16>(converted_zeros),
-                          recast<nv_bfloat16>(converted_inputs), cute::plus{});
+                          recast<nv_bfloat16>(converted_inputs),
+                          cute::plus{});
         } else {
-          cute::transform(converted_inputs, converted_zeros, converted_inputs,
+          cute::transform(converted_inputs,
+                          converted_zeros,
+                          converted_inputs,
                           cute::plus{});
         }
       }
@@ -1421,7 +1487,9 @@ struct MacheteCollectiveMma {
   // Modified from upstream, should be kept the same when possible
   //   the main differences is that this version supports interleaved converts
   // Utilities for transforming the A operand prior to issuing tensorcore math.
-  template <typename IlvdBlkLayout, class EngineIn, class EngineOut,
+  template <typename IlvdBlkLayout,
+            class EngineIn,
+            class EngineOut,
             class TensorLayout,
             int ConversionVectorWidth = cosize_v<TensorLayout>>
   CUTLASS_DEVICE void convert_tensor(
@@ -1455,8 +1523,12 @@ struct MacheteCollectiveMma {
     constexpr cutlass::FloatRoundStyle RoundStyle =
         cutlass::FloatRoundStyle::round_to_nearest;
 
-    using Converter = cutlass::InterleavedNumericArrayConverter<
-        IlvdBlkLayout, DstType, SrcType, ConversionVectorWidth, RoundStyle>;
+    using Converter =
+        cutlass::InterleavedNumericArrayConverter<IlvdBlkLayout,
+                                                  DstType,
+                                                  SrcType,
+                                                  ConversionVectorWidth,
+                                                  RoundStyle>;
 
     constexpr int NumIterations = N / ConversionVectorWidth;
 
diff --git a/custom_ops/gpu_ops/machete/machete_mm.cu b/custom_ops/gpu_ops/machete/machete_mm.cu
index 920c2e5ad5..86a4e95ee1 100644
--- a/custom_ops/gpu_ops/machete/machete_mm.cu
+++ b/custom_ops/gpu_ops/machete/machete_mm.cu
@@ -17,20 +17,22 @@
 
 template <typename T>
 std::optional<T> ConvertToStdOptional(const paddle::optional<T>& paddle_opt) {
-    return paddle_opt ? std::optional<T>(paddle_opt.get()) : std::nullopt;
+  return paddle_opt ? std::optional<T>(paddle_opt.get()) : std::nullopt;
 }
 
-paddle::Tensor mm(paddle::Tensor const& A, paddle::Tensor const& B,
-                 int64_t b_type_id,
-                 std::optional<paddle::DataType> const& maybe_out_type,
-                 std::optional<paddle::Tensor> const& maybe_group_scales,
-                 std::optional<paddle::Tensor> const& maybe_group_zeros,
-                 int64_t maybe_group_size,
-                 std::optional<paddle::Tensor> const& maybe_channel_scales,
-                 std::optional<paddle::Tensor> const& maybe_token_scales,
-                 std::string maybe_schedule) {
+paddle::Tensor mm(paddle::Tensor const& A,
+                  paddle::Tensor const& B,
+                  int64_t b_type_id,
+                  std::optional<paddle::DataType> const& maybe_out_type,
+                  std::optional<paddle::Tensor> const& maybe_group_scales,
+                  std::optional<paddle::Tensor> const& maybe_group_zeros,
+                  int64_t maybe_group_size,
+                  std::optional<paddle::Tensor> const& maybe_channel_scales,
+                  std::optional<paddle::Tensor> const& maybe_token_scales,
+                  std::string maybe_schedule) {
   machete::ScalarType const b_type = machete::ScalarType::from_id(b_type_id);
-  std::optional<int64_t> maybe_group_size_opt = std::optional<int64_t>(maybe_group_size);
+  std::optional<int64_t> maybe_group_size_opt =
+      std::optional<int64_t>(maybe_group_size);
   std::optional<std::string> maybe_schedule_opt;
   if (maybe_schedule == "") {
     maybe_schedule_opt = std::nullopt;
@@ -38,19 +40,20 @@ paddle::Tensor mm(paddle::Tensor const& A, paddle::Tensor const& B,
     maybe_schedule_opt = std::optional<std::string>(maybe_schedule);
   }
   return machete::mm_dispatch({.A = A,
-                      .B = B,
-                      .b_type = b_type,
-                      .maybe_out_type = maybe_out_type,
-                      .maybe_group_scales = maybe_group_scales,
-                      .maybe_group_zeros = maybe_group_zeros,
-                      .maybe_group_size = maybe_group_size_opt,
-                      .maybe_channel_scales = maybe_channel_scales,
-                      .maybe_token_scales = maybe_token_scales,
-                      .maybe_schedule = maybe_schedule_opt});
+                               .B = B,
+                               .b_type = b_type,
+                               .maybe_out_type = maybe_out_type,
+                               .maybe_group_scales = maybe_group_scales,
+                               .maybe_group_zeros = maybe_group_zeros,
+                               .maybe_group_size = maybe_group_size_opt,
+                               .maybe_channel_scales = maybe_channel_scales,
+                               .maybe_token_scales = maybe_token_scales,
+                               .maybe_schedule = maybe_schedule_opt});
 }
 
 std::vector<paddle::Tensor> MacheteMMKernel(
-    paddle::Tensor const& A, paddle::Tensor const& B,
+    paddle::Tensor const& A,
+    paddle::Tensor const& B,
     paddle::optional<paddle::Tensor> const& maybe_group_scales,
     paddle::optional<paddle::Tensor> const& maybe_group_zeros,
     paddle::optional<paddle::Tensor> const& maybe_channel_scales,
@@ -58,9 +61,7 @@ std::vector<paddle::Tensor> MacheteMMKernel(
     std::string const& b_type_str,
     std::string const& maybe_out_type_str,
     int64_t const& maybe_group_size,
-    std::string const& maybe_schedule
-  ) {
-
+    std::string const& maybe_schedule) {
   machete::ScalarTypeId b_type_id;
   paddle::DataType maybe_out_type;
   if (b_type_str == "uint4b8") {
@@ -77,7 +78,10 @@ std::vector<paddle::Tensor> MacheteMMKernel(
   } else {
     maybe_out_type = A.dtype();
   }
-  auto out = mm(A, B, b_type_id, maybe_out_type,
+  auto out = mm(A,
+                B,
+                b_type_id,
+                maybe_out_type,
                 ConvertToStdOptional<paddle::Tensor>(maybe_group_scales),
                 ConvertToStdOptional<paddle::Tensor>(maybe_group_zeros),
                 maybe_group_size,
@@ -112,7 +116,6 @@ std::vector<paddle::DataType> MacheteMMKernelInferDtype(
     std::string const& maybe_out_type_str,
     int64_t const& maybe_group_size,
     std::string const& maybe_schedule) {
-
   paddle::DataType maybe_out_type;
   if (maybe_out_type_str == "float16") {
     maybe_out_type = paddle::DataType::FLOAT16;
@@ -125,13 +128,17 @@ std::vector<paddle::DataType> MacheteMMKernelInferDtype(
 }
 
 PD_BUILD_STATIC_OP(machete_mm)
-    .Inputs({"A", "B",
+    .Inputs({"A",
+             "B",
              paddle::Optional("maybe_group_scales"),
              paddle::Optional("maybe_group_zeros"),
              paddle::Optional("maybe_channel_scales"),
              paddle::Optional("maybe_token_scales")})
     .Outputs({"out"})
-    .Attrs({"b_type_str:std::string", "maybe_out_type_str:std::string", "maybe_group_size:int64_t", "maybe_schedule:std::string"})
+    .Attrs({"b_type_str:std::string",
+            "maybe_out_type_str:std::string",
+            "maybe_group_size:int64_t",
+            "maybe_schedule:std::string"})
     .SetKernelFn(PD_KERNEL(MacheteMMKernel))
     .SetInferShapeFn(PD_INFER_SHAPE(MacheteMMKernelInferShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(MacheteMMKernelInferDtype));
diff --git a/custom_ops/gpu_ops/machete/machete_mm_kernel.cuh b/custom_ops/gpu_ops/machete/machete_mm_kernel.cuh
index 1f6d548d5b..a25a918cb0 100644
--- a/custom_ops/gpu_ops/machete/machete_mm_kernel.cuh
+++ b/custom_ops/gpu_ops/machete/machete_mm_kernel.cuh
@@ -34,9 +34,15 @@ using namespace cute;
 //   register. Since the primary use case we want to support is Y = XW^t where
 //   W is quantized, in this situation or right-hand operand is quantized so
 //   we compute the transpose to move it to the left-hand side.
-template <typename ElementA_, typename ElementB_, typename ElementD_,
-          typename AccumulatorT, typename GroupScaleT, typename GroupZeroT,
-          typename ChannelScaleT, typename TokenScaleT, class KernelSchedule,
+template <typename ElementA_,
+          typename ElementB_,
+          typename ElementD_,
+          typename AccumulatorT,
+          typename GroupScaleT,
+          typename GroupZeroT,
+          typename ChannelScaleT,
+          typename TokenScaleT,
+          class KernelSchedule,
           typename ScheduleConfig>
 struct MacheteKernelTemplate {
   static constexpr bool with_C = false;  // not ever used
@@ -97,9 +103,12 @@ struct MacheteKernelTemplate {
   using ArchTag = cutlass::arch::Sm90;
   using OperatorClass = cutlass::arch::OpClassTensorOp;
 
-  using PrepackedLayoutB =
-      PrepackedLayoutBTemplate<ElementA_, ElementB_, ElementConvertGroup,
-                               AccumulatorT, LayoutA_Transpose, KernelSchedule>;
+  using PrepackedLayoutB = PrepackedLayoutBTemplate<ElementA_,
+                                                    ElementB_,
+                                                    ElementConvertGroup,
+                                                    AccumulatorT,
+                                                    LayoutA_Transpose,
+                                                    KernelSchedule>;
 
   static int constexpr TileShapeK =
       128 * 8 / cutlass::sizeof_bits<MmaType>::value;
@@ -123,9 +132,8 @@ struct MacheteKernelTemplate {
       "Currently token and channel scales (if present) must be the same type");
 
   // Currently only supports float scales
-  using ChTokScalesEpilogue =
-      typename fastdeploy::c3x::ScaledEpilogue<ElementAccumulator, ElementD,
-                                         TileShape>;
+  using ChTokScalesEpilogue = typename fastdeploy::c3x::
+      ScaledEpilogue<ElementAccumulator, ElementD, TileShape>;
   static_assert((with_channel_scales || with_token_scales) ||
                     (std::is_same_v<ElementSChannel, float> &&
                      std::is_same_v<ElementSToken, float>),
@@ -143,23 +151,45 @@ struct MacheteKernelTemplate {
   // EVTCompute
   using CollectiveEpilogue =
       typename cutlass::epilogue::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, TileShape, ClusterShape, EpilogueTileType,
-          ElementAccumulator, ElementSChannel, ElementC, LayoutC_Transpose,
-          AlignmentC, ElementD, LayoutD_Transpose, AlignmentD, EpilogueSchedule,
+          ArchTag,
+          OperatorClass,
+          TileShape,
+          ClusterShape,
+          EpilogueTileType,
+          ElementAccumulator,
+          ElementSChannel,
+          ElementC,
+          LayoutC_Transpose,
+          AlignmentC,
+          ElementD,
+          LayoutD_Transpose,
+          AlignmentD,
+          EpilogueSchedule,
           EVTCompute>::CollectiveOp;
 
   using CollectiveMainloop =
       typename cutlass::gemm::collective::MacheteCollectiveBuilder<
-          cutlass::gemm::collective::MacheteKernelTag, ArchTag, OperatorClass,
-          BTypeTuple, PrepackedLayoutB, AlignmentB, ElementA, LayoutA_Transpose,
-          AlignmentA, ElementAccumulator, TileShape, ClusterShape,
+          cutlass::gemm::collective::MacheteKernelTag,
+          ArchTag,
+          OperatorClass,
+          BTypeTuple,
+          PrepackedLayoutB,
+          AlignmentB,
+          ElementA,
+          LayoutA_Transpose,
+          AlignmentA,
+          ElementAccumulator,
+          TileShape,
+          ClusterShape,
           cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
               sizeof(typename CollectiveEpilogue::SharedStorage))>,
           KernelSchedule>::CollectiveOp;
 
   using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
       Shape<int, int, int, int>,  // Indicates ProblemShape
-      CollectiveMainloop, CollectiveEpilogue, TileScheduler>;
+      CollectiveMainloop,
+      CollectiveEpilogue,
+      TileScheduler>;
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
 
   // stride_B is unused (since B is prepacked), but still required by cutlass
@@ -194,9 +224,7 @@ struct MacheteKernelTemplate {
     int64_t numel_S_channel = maybe_ch_scales ? maybe_ch_scales->numel() : 0;
     int64_t numel_S_token = maybe_tok_scales ? maybe_tok_scales->numel() : 0;
 
-    auto unwrap = [](auto const& t) {
-      return t ? t->data() : nullptr;
-    };
+    auto unwrap = [](auto const& t) { return t ? t->data() : nullptr; };
     auto A_ptr = static_cast<ElementA const*>(A.data());
     auto B_ptr = static_cast<ElementB const*>(B.data());
     auto D_ptr = static_cast<ElementD*>(D.data());
@@ -218,7 +246,7 @@ struct MacheteKernelTemplate {
     if constexpr (with_group_scales) {
       PD_CHECK(S_group_ptr && layout_S_group);
       PD_CHECK((size<0>(*layout_S_group) == scale_k &&
-                   size<1>(*layout_S_group) == N));
+                size<1>(*layout_S_group) == N));
     } else {
       PD_CHECK(!S_group_ptr, "Scales not supported");
     }
@@ -226,9 +254,9 @@ struct MacheteKernelTemplate {
     if constexpr (with_group_zeropoints) {
       PD_CHECK(Z_group_ptr && layout_Z_group);
       PD_CHECK((size<0>(*layout_Z_group) == scale_k &&
-                   size<1>(*layout_Z_group) == N));
+                size<1>(*layout_Z_group) == N));
       PD_CHECK(layout_S_group && *layout_Z_group == *layout_S_group,
-                  "Scales and zeros must have the same layout");
+               "Scales and zeros must have the same layout");
     } else {
       PD_CHECK(!Z_group_ptr, "Zeropoints not supported");
     }
@@ -263,14 +291,23 @@ struct MacheteKernelTemplate {
 
     if constexpr (with_group_scales && with_group_zeropoints) {
       auto stride_S_group = permute_layout<1, 0, 2>(*layout_S_group).stride();
-      mainloop_arguments = MainloopArguments{
-          B_ptr,       _StrideB{},     A_ptr,      stride_At,
-          S_group_ptr, stride_S_group, group_size, Z_group_ptr};
+      mainloop_arguments = MainloopArguments{B_ptr,
+                                             _StrideB{},
+                                             A_ptr,
+                                             stride_At,
+                                             S_group_ptr,
+                                             stride_S_group,
+                                             group_size,
+                                             Z_group_ptr};
     } else if constexpr (with_group_scales) {
       auto stride_S_group = permute_layout<1, 0, 2>(*layout_S_group).stride();
-      mainloop_arguments =
-          MainloopArguments{B_ptr,       _StrideB{},     A_ptr,     stride_At,
-                            S_group_ptr, stride_S_group, group_size};
+      mainloop_arguments = MainloopArguments{B_ptr,
+                                             _StrideB{},
+                                             A_ptr,
+                                             stride_At,
+                                             S_group_ptr,
+                                             stride_S_group,
+                                             group_size};
     } else {
       mainloop_arguments =
           MainloopArguments{B_ptr, _StrideB{}, A_ptr, stride_At};
@@ -295,7 +332,7 @@ struct MacheteKernelTemplate {
 
     cutlass::Status status = gemm_op.initialize(args, workspace, stream);
     PD_CHECK(status == cutlass::Status::kSuccess,
-                "Machete kernel failed to initialize workspace");
+             "Machete kernel failed to initialize workspace");
 
     status = gemm_op.run(stream);
     PD_CHECK(status == cutlass::Status::kSuccess, "Machete kernel failed");
diff --git a/custom_ops/gpu_ops/machete/machete_prepack_B.cu b/custom_ops/gpu_ops/machete/machete_prepack_B.cu
index 223a3e654d..ab16d56e29 100644
--- a/custom_ops/gpu_ops/machete/machete_prepack_B.cu
+++ b/custom_ops/gpu_ops/machete/machete_prepack_B.cu
@@ -15,24 +15,21 @@
 #include "machete_mm_launcher.cuh"
 #include "machete_prepack_launcher.cuh"
 
-paddle::Tensor prepack_B(
-    paddle::Tensor const& B, paddle::DataType const& a_type, int64_t b_type_id,
-    std::string const& maybe_group_scales_type_str) {
+paddle::Tensor prepack_B(paddle::Tensor const& B,
+                         paddle::DataType const& a_type,
+                         int64_t b_type_id,
+                         std::string const& maybe_group_scales_type_str) {
   machete::ScalarType const b_type = machete::ScalarType::from_id(b_type_id);
   std::optional<paddle::DataType> maybe_group_scales_type;
   if (maybe_group_scales_type_str == "float16") {
     maybe_group_scales_type = paddle::DataType::FLOAT16;
-  }
-  else if (maybe_group_scales_type_str == "bfloat16") {
+  } else if (maybe_group_scales_type_str == "bfloat16") {
     maybe_group_scales_type = paddle::DataType::BFLOAT16;
-  }
-  else if (maybe_group_scales_type_str == "float32") {
+  } else if (maybe_group_scales_type_str == "float32") {
     maybe_group_scales_type = paddle::DataType::FLOAT32;
-  }
-  else if (maybe_group_scales_type_str == "") {
+  } else if (maybe_group_scales_type_str == "") {
     maybe_group_scales_type = std::nullopt;
-  }
-  else {
+  } else {
     PADDLE_ENFORCE(false, "maybe_group_scales_type_str not supported!");
   }
   return machete::prepack_B_dispatch(
@@ -43,9 +40,10 @@ paddle::Tensor prepack_B(
 }
 
 std::vector<paddle::Tensor> MachetePrepackBKernel(
-    paddle::Tensor const& B, std::string const& a_type_str, std::string const& b_type_str,
+    paddle::Tensor const& B,
+    std::string const& a_type_str,
+    std::string const& b_type_str,
     std::string const& maybe_group_scales_type_str) {
-
   machete::ScalarTypeId b_type_id;
   paddle::DataType a_type, maybe_group_scales_type;
 
@@ -59,27 +57,29 @@ std::vector<paddle::Tensor> MachetePrepackBKernel(
 
   if (a_type_str == "float16") {
     a_type = paddle::DataType::FLOAT16;
-  }
-  else if (a_type_str == "bfloat16") {
+  } else if (a_type_str == "bfloat16") {
     a_type = paddle::DataType::BFLOAT16;
-  }
-  else {
+  } else {
     PADDLE_ENFORCE(false, "a_type_str not supported!");
   }
   auto Bt = paddle::experimental::transpose(B, {1, 0});
-  paddle::Tensor B_prepacked = prepack_B(Bt, a_type, b_type_id, maybe_group_scales_type_str);
+  paddle::Tensor B_prepacked =
+      prepack_B(Bt, a_type, b_type_id, maybe_group_scales_type_str);
   return {B_prepacked};
-
 }
 
 std::vector<std::vector<int64_t>> MachetePrepackBKernelInferShape(
-    std::vector<int64_t> const& B_shape, std::string const& a_type_str, std::string const& b_type_str,
+    std::vector<int64_t> const& B_shape,
+    std::string const& a_type_str,
+    std::string const& b_type_str,
     std::string const& maybe_group_scales_type_str) {
   return {{B_shape[1], B_shape[0]}};
 }
 
 std::vector<paddle::DataType> MachetePrepackBKernelInferDtype(
-    paddle::DataType const& B_dtype, std::string const& a_type_str, std::string const& b_type_str,
+    paddle::DataType const& B_dtype,
+    std::string const& a_type_str,
+    std::string const& b_type_str,
     std::string const& maybe_group_scales_type_str) {
   return {B_dtype};
 }
@@ -87,7 +87,9 @@ std::vector<paddle::DataType> MachetePrepackBKernelInferDtype(
 PD_BUILD_STATIC_OP(machete_prepack_B)
     .Inputs({"B"})
     .Outputs({"B_prepacked"})
-    .Attrs({"a_type_str:std::string", "b_type_str:std::string", "maybe_group_scales_type_str:std::string"})
+    .Attrs({"a_type_str:std::string",
+            "b_type_str:std::string",
+            "maybe_group_scales_type_str:std::string"})
     .SetKernelFn(PD_KERNEL(MachetePrepackBKernel))
     .SetInferShapeFn(PD_INFER_SHAPE(MachetePrepackBKernelInferShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(MachetePrepackBKernelInferDtype));
diff --git a/custom_ops/gpu_ops/machete/machete_prepack_kernel.cuh b/custom_ops/gpu_ops/machete/machete_prepack_kernel.cuh
index af84fdfef5..a97961580a 100644
--- a/custom_ops/gpu_ops/machete/machete_prepack_kernel.cuh
+++ b/custom_ops/gpu_ops/machete/machete_prepack_kernel.cuh
@@ -6,7 +6,9 @@
 
 namespace machete {
 
-template <int threads, typename PrepackedLayoutB, typename BInTensor,
+template <int threads,
+          typename PrepackedLayoutB,
+          typename BInTensor,
           typename ElementB>
 static __global__ void prepack_B_kernel(BInTensor B_in, ElementB* B_out_ptr) {
   auto constexpr block_size =
@@ -17,9 +19,10 @@ static __global__ void prepack_B_kernel(BInTensor B_in, ElementB* B_out_ptr) {
 
   // Which pre-packed are we responsible for
   auto blk_coord = make_coord(blockIdx.x, blockIdx.y, blockIdx.z);
-  auto tB_in = local_tile(
-      B_in, append(typename PrepackedLayoutB::PPBlockShape_NK{}, _1{}),
-      blk_coord);
+  auto tB_in =
+      local_tile(B_in,
+                 append(typename PrepackedLayoutB::PPBlockShape_NK{}, _1{}),
+                 blk_coord);
 
   // Find the start offset in the output for this pre-packed block
   auto bNbKL_to_offset = PrepackedLayoutB::bNbKL_to_offset(shape(B_in));
@@ -53,8 +56,10 @@ static __global__ void prepack_B_kernel(BInTensor B_in, ElementB* B_out_ptr) {
 
 template <typename PrepackedLayoutB, typename InLayout>
 static void prepack_B_template(
-    cudaStream_t stream, typename PrepackedLayoutB::ElementB const* B_in_ptr,
-    InLayout B_layout, typename PrepackedLayoutB::ElementB* B_out_ptr) {
+    cudaStream_t stream,
+    typename PrepackedLayoutB::ElementB const* B_in_ptr,
+    InLayout B_layout,
+    typename PrepackedLayoutB::ElementB* B_out_ptr) {
   using TileShapeNKL =
       decltype(append(typename PrepackedLayoutB::PPBlockShape_NK{}, _1{}));
   auto ilvd_NKbNbKL_to_offset =
diff --git a/custom_ops/gpu_ops/machete/machete_prepack_launcher.cuh b/custom_ops/gpu_ops/machete/machete_prepack_launcher.cuh
index 02264852a2..885097e82b 100644
--- a/custom_ops/gpu_ops/machete/machete_prepack_launcher.cuh
+++ b/custom_ops/gpu_ops/machete/machete_prepack_launcher.cuh
@@ -27,17 +27,17 @@ paddle::Tensor prepack_impl(paddle::Tensor const B) {
   auto eles_per_storage =
       (SizeOf(B.dtype()) * 8) / cute::sizeof_bits_v<ElementB>;
 
-  // paddle B passed in is/should be (packed_K,N), the kernel expects (N,K,L) (to
-  // match cutlass using (N,K,L) for B), so we transpose B to (N,packed_K,L)
+  // paddle B passed in is/should be (packed_K,N), the kernel expects (N,K,L)
+  // (to match cutlass using (N,K,L) for B), so we transpose B to (N,packed_K,L)
   // auto Bt_packed = B.transpose();
   auto Bt_packed = paddle::experimental::transpose(B, {1, 0});
 
-  PD_CHECK(
-      (B.shape()[0] * eles_per_storage) % size<1>(PPBlockShape_NK{}) == 0,
-      "B.shape[0] (in terms of unpacked elements) must be a multiple of ",
-      size<1>(PPBlockShape_NK{}));
+  PD_CHECK((B.shape()[0] * eles_per_storage) % size<1>(PPBlockShape_NK{}) == 0,
+           "B.shape[0] (in terms of unpacked elements) must be a multiple of ",
+           size<1>(PPBlockShape_NK{}));
   PD_CHECK(B.shape()[1] % size<0>(PPBlockShape_NK{}) == 0,
-              "B.shape[1] must be a multiple of ", size<0>(PPBlockShape_NK{}));
+           "B.shape[1] must be a multiple of ",
+           size<0>(PPBlockShape_NK{}));
 
   using StrideB = cutlass::detail::TagToStrideB_t<cutlass::layout::ColumnMajor>;
   auto const l_Bt_packed = make_cute_layout<StrideB>(Bt_packed, "B");
diff --git a/custom_ops/gpu_ops/machete/machete_prepacked_layout.cuh b/custom_ops/gpu_ops/machete/machete_prepacked_layout.cuh
index 62df504558..316c6578eb 100644
--- a/custom_ops/gpu_ops/machete/machete_prepacked_layout.cuh
+++ b/custom_ops/gpu_ops/machete/machete_prepacked_layout.cuh
@@ -37,8 +37,12 @@ struct IlvBlkLayoutAuto {};
 // The contract here is that the `TiledMma` determined below matches the one
 // ultimately used in the kernel. (this is also why the other element types are
 // required along with the kernel schedule)
-template <typename ElementA_, typename ElementB_, typename ElementConvert_,
-          typename AccumulatorT, class LayoutB, class KernelSchedule,
+template <typename ElementA_,
+          typename ElementB_,
+          typename ElementConvert_,
+          typename AccumulatorT,
+          class LayoutB,
+          class KernelSchedule,
           typename IlvBlkLayout_ = IlvBlkLayoutAuto>
 // clang-format on
 struct PrepackedLayoutBTemplate {
@@ -63,8 +67,9 @@ struct PrepackedLayoutBTemplate {
       std::is_same_v<IlvBlkLayout_, IlvBlkLayoutAuto>,
       std::conditional_t<
           should_interleave,
-          decltype(get_interleaved_blk_layout<
-                   ElementB, sizeof_bits_v<ElementConvert_>, 32>()),
+          decltype(get_interleaved_blk_layout<ElementB,
+                                              sizeof_bits_v<ElementConvert_>,
+                                              32>()),
           void>,
       IlvBlkLayout_>;
 
@@ -85,8 +90,8 @@ struct PrepackedLayoutBTemplate {
   //  registers.
   // The _128 here doesn't actually impact the shape of the stored tile directly
   //  but may impact the op selected by rs_op_selector
-  using GemmTileShape = decltype(make_shape(size<0>(PPBlockShape_NK{}), _128{},
-                                            size<1>(PPBlockShape_NK{})));
+  using GemmTileShape = decltype(make_shape(
+      size<0>(PPBlockShape_NK{}), _128{}, size<1>(PPBlockShape_NK{})));
 
   static constexpr cute::GMMA::Major GmmaMajorB =
       gmma_rs_tag_to_major_B<LayoutB>();
@@ -95,11 +100,16 @@ struct PrepackedLayoutBTemplate {
   // instructions so we use 2 atoms along the M dim (one for each warpgroup)
   using AtomLayoutMNK = cute::conditional_t<
       cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedCooperative>,
-      Layout<Shape<_2, _1, _1>>, Layout<Shape<_1, _1, _1>>>;
+      Layout<Shape<_2, _1, _1>>,
+      Layout<Shape<_1, _1, _1>>>;
 
   using TiledMma = decltype(cute::make_tiled_mma(
-      cute::GMMA::rs_op_selector<ElementMma, ElementMma, ElementAccumulator,
-                                 GemmTileShape, GMMA::Major::K, GmmaMajorB>(),
+      cute::GMMA::rs_op_selector<ElementMma,
+                                 ElementMma,
+                                 ElementAccumulator,
+                                 GemmTileShape,
+                                 GMMA::Major::K,
+                                 GmmaMajorB>(),
       AtomLayoutMNK{}));
 
   // Prepacked block, (athrid, val) -> (N,K)
@@ -163,9 +173,9 @@ struct PrepackedLayoutBTemplate {
     constexpr auto block_layout = ppblock_TV_to_offset();
 
     // (BlocksN, BlocksK, L)
-    auto blocks_shape =
-        cute::transform(shape_mkl, append(PPBlockShape_NK{}, _1{}),
-                        [](auto x, auto y) { return x / y; });
+    auto blocks_shape = cute::transform(shape_mkl,
+                                        append(PPBlockShape_NK{}, _1{}),
+                                        [](auto x, auto y) { return x / y; });
 
     // ((athrid, val), (BlocksN, BlocksK, L)) -> (storage_idx)
     auto result = make_layout(
@@ -198,9 +208,9 @@ struct PrepackedLayoutBTemplate {
     constexpr auto block_layout = ppblock_ilvd_NK_to_offset();
 
     // (BlocksN, BlocksK, L)
-    auto blocks_shape =
-        cute::transform(shape_mkl, append(PPBlockShape_NK{}, _1{}),
-                        [](auto x, auto y) { return x / y; });
+    auto blocks_shape = cute::transform(shape_mkl,
+                                        append(PPBlockShape_NK{}, _1{}),
+                                        [](auto x, auto y) { return x / y; });
 
     // ((athrid, val), (BlocksN, BlocksK, L)) -> (storage_idx)
     auto result = make_layout(
@@ -217,9 +227,9 @@ struct PrepackedLayoutBTemplate {
   template <typename Shape_NKL>
   CUTE_HOST_DEVICE static constexpr auto bNbKL_to_offset(Shape_NKL shape_mkl) {
     // (BlocksN, BlocksK, L)
-    auto blocks_shape =
-        cute::transform(shape_mkl, append(PPBlockShape_NK{}, _1{}),
-                        [](auto x, auto y) { return x / y; });
+    auto blocks_shape = cute::transform(shape_mkl,
+                                        append(PPBlockShape_NK{}, _1{}),
+                                        [](auto x, auto y) { return x / y; });
     auto stride = size(PPBlockShape_NK{});
 
     // (BlocksN, BlocksK, L) -> (storage_idx)
diff --git a/custom_ops/gpu_ops/machete/machete_supported_schedules.cu b/custom_ops/gpu_ops/machete/machete_supported_schedules.cu
index 79f1fc571b..92c81cc361 100644
--- a/custom_ops/gpu_ops/machete/machete_supported_schedules.cu
+++ b/custom_ops/gpu_ops/machete/machete_supported_schedules.cu
@@ -17,26 +17,26 @@
 
 template <typename T>
 std::optional<T> ConvertToStdOptional(const paddle::optional<T>& paddle_opt) {
-    return paddle_opt ? std::optional<T>(paddle_opt.get()) : std::nullopt;
+  return paddle_opt ? std::optional<T>(paddle_opt.get()) : std::nullopt;
 }
 
 std::vector<std::string> supported_schedules(
-    paddle::DataType a_type, int64_t b_type_id,
+    paddle::DataType a_type,
+    int64_t b_type_id,
     std::optional<paddle::DataType> maybe_group_scales_type,
     std::optional<paddle::DataType> maybe_group_zeros_type,
     std::optional<paddle::DataType> maybe_channel_scales_type,
     std::optional<paddle::DataType> maybe_token_scales_type,
     std::optional<paddle::DataType> maybe_out_type) {
   machete::ScalarType const b_type = machete::ScalarType::from_id(b_type_id);
-  auto schedules = machete::supported_schedules_dispatch({
-      .a_type = a_type,
-      .b_type = b_type,
-      .maybe_group_scales_type = maybe_group_scales_type,
-      .maybe_group_zeros_type = maybe_group_zeros_type,
-      .maybe_channel_scales_type = maybe_channel_scales_type,
-      .maybe_token_scales_type = maybe_token_scales_type,
-      .maybe_out_type = maybe_out_type
-  });
+  auto schedules = machete::supported_schedules_dispatch(
+      {.a_type = a_type,
+       .b_type = b_type,
+       .maybe_group_scales_type = maybe_group_scales_type,
+       .maybe_group_zeros_type = maybe_group_zeros_type,
+       .maybe_channel_scales_type = maybe_channel_scales_type,
+       .maybe_token_scales_type = maybe_token_scales_type,
+       .maybe_out_type = maybe_out_type});
   return schedules;
 }
 
@@ -56,17 +56,20 @@ std::vector<std::string> MacheteSupportedSchedules(
   } else {
     PADDLE_ENFORCE(false, "a_type_str not supported!");
   }
-  std::optional<paddle::DataType> maybe_group_scales_type = std::optional<paddle::DataType>(a_type);
-  std::optional<paddle::DataType> maybe_out_type = std::optional<paddle::DataType>(a_type);
+  std::optional<paddle::DataType> maybe_group_scales_type =
+      std::optional<paddle::DataType>(a_type);
+  std::optional<paddle::DataType> maybe_out_type =
+      std::optional<paddle::DataType>(a_type);
   std::optional<paddle::DataType> maybe_group_zeros_type = std::nullopt;
   std::optional<paddle::DataType> maybe_channel_scales_type = std::nullopt;
   std::optional<paddle::DataType> maybe_token_scales_type = std::nullopt;
 
-  auto schedules =  supported_schedules(a_type, b_type_id,
-        maybe_group_scales_type,
-        maybe_group_zeros_type,
-        maybe_channel_scales_type,
-        maybe_token_scales_type,
-        maybe_out_type);
+  auto schedules = supported_schedules(a_type,
+                                       b_type_id,
+                                       maybe_group_scales_type,
+                                       maybe_group_zeros_type,
+                                       maybe_channel_scales_type,
+                                       maybe_token_scales_type,
+                                       maybe_out_type);
   return schedules;
 }
diff --git a/custom_ops/gpu_ops/machete/utils/cute_utils.cuh b/custom_ops/gpu_ops/machete/utils/cute_utils.cuh
index 18765a8e28..7bb94a582e 100644
--- a/custom_ops/gpu_ops/machete/utils/cute_utils.cuh
+++ b/custom_ops/gpu_ops/machete/utils/cute_utils.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/cute_utils.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/cute_utils.cuh
 #pragma once
 
 #include <cute/tensor.hpp>
diff --git a/custom_ops/gpu_ops/machete/utils/machete_collective_builder.cuh b/custom_ops/gpu_ops/machete/utils/machete_collective_builder.cuh
index 251710a060..040865596c 100644
--- a/custom_ops/gpu_ops/machete/utils/machete_collective_builder.cuh
+++ b/custom_ops/gpu_ops/machete/utils/machete_collective_builder.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/vllm_collective_builder.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/vllm_collective_builder.cuh
 #pragma once
 
 #include "cutlass_extensions/gemm/collective/collective_builder.hpp"
@@ -7,8 +8,8 @@ namespace cutlass::gemm::collective {
 using namespace cute;
 
 //
-// MacheteCollectiveBuilder is a wrapper around CollectiveBuilder that allows for
-// for custom kernel tags, allowing you to build custom collectives. Without
+// MacheteCollectiveBuilder is a wrapper around CollectiveBuilder that allows
+// for for custom kernel tags, allowing you to build custom collectives. Without
 // touching the cutlass library headers, using `CutlassKernelTag` will mean it
 // will resort to using the standard cutlass collective builder.
 //
@@ -17,28 +18,67 @@ using namespace cute;
 // collective
 struct CutlassKernelTag {};
 
-template <class KernelTag, class ArchTag, class OpClass, class ElementA,
-          class GmemLayoutA, int AlignmentA, class ElementB, class GmemLayoutB,
-          int AlignmentB, class ElementAccumulator, class TileShape_MNK,
-          class ClusterShape_MNK, class StageCountType,
-          class KernelScheduleType, class Enable = void>
+template <class KernelTag,
+          class ArchTag,
+          class OpClass,
+          class ElementA,
+          class GmemLayoutA,
+          int AlignmentA,
+          class ElementB,
+          class GmemLayoutB,
+          int AlignmentB,
+          class ElementAccumulator,
+          class TileShape_MNK,
+          class ClusterShape_MNK,
+          class StageCountType,
+          class KernelScheduleType,
+          class Enable = void>
 struct MacheteCollectiveBuilder {
   static_assert(sizeof(ElementA) == 0,
                 "Could not build a collective for given parameters.");
 };
 
-template <class ArchTag, class OpClass, class ElementA, class GmemLayoutA,
-          int AlignmentA, class ElementB, class GmemLayoutB, int AlignmentB,
-          class ElementAccumulator, class TileShape_MNK, class ClusterShape_MNK,
-          class StageCountType, class KernelScheduleType>
-struct MacheteCollectiveBuilder<
-    CutlassKernelTag, ArchTag, OpClass, ElementA, GmemLayoutA, AlignmentA,
-    ElementB, GmemLayoutB, AlignmentB, ElementAccumulator, TileShape_MNK,
-    ClusterShape_MNK, StageCountType, KernelScheduleType> {
-  using CollectiveOp = typename CollectiveBuilder<
-      ArchTag, OpClass, ElementA, GmemLayoutA, AlignmentA, ElementB,
-      GmemLayoutB, AlignmentB, ElementAccumulator, TileShape_MNK,
-      ClusterShape_MNK, StageCountType, KernelScheduleType>::CollectiveOp;
+template <class ArchTag,
+          class OpClass,
+          class ElementA,
+          class GmemLayoutA,
+          int AlignmentA,
+          class ElementB,
+          class GmemLayoutB,
+          int AlignmentB,
+          class ElementAccumulator,
+          class TileShape_MNK,
+          class ClusterShape_MNK,
+          class StageCountType,
+          class KernelScheduleType>
+struct MacheteCollectiveBuilder<CutlassKernelTag,
+                                ArchTag,
+                                OpClass,
+                                ElementA,
+                                GmemLayoutA,
+                                AlignmentA,
+                                ElementB,
+                                GmemLayoutB,
+                                AlignmentB,
+                                ElementAccumulator,
+                                TileShape_MNK,
+                                ClusterShape_MNK,
+                                StageCountType,
+                                KernelScheduleType> {
+  using CollectiveOp =
+      typename CollectiveBuilder<ArchTag,
+                                 OpClass,
+                                 ElementA,
+                                 GmemLayoutA,
+                                 AlignmentA,
+                                 ElementB,
+                                 GmemLayoutB,
+                                 AlignmentB,
+                                 ElementAccumulator,
+                                 TileShape_MNK,
+                                 ClusterShape_MNK,
+                                 StageCountType,
+                                 KernelScheduleType>::CollectiveOp;
 };
 
 };  // namespace cutlass::gemm::collective
diff --git a/custom_ops/gpu_ops/machete/utils/machete_custom_types.cuh b/custom_ops/gpu_ops/machete/utils/machete_custom_types.cuh
index 2fb39a318c..a3fee76dd3 100644
--- a/custom_ops/gpu_ops/machete/utils/machete_custom_types.cuh
+++ b/custom_ops/gpu_ops/machete/utils/machete_custom_types.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/vllm_custom_types.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/vllm_custom_types.cuh
 #pragma once
 
 #include "cutlass/integer_subbyte.h"
diff --git a/custom_ops/gpu_ops/machete/utils/machete_numeric_conversion.cuh b/custom_ops/gpu_ops/machete/utils/machete_numeric_conversion.cuh
index 4483ec4249..cc0da2ae5d 100644
--- a/custom_ops/gpu_ops/machete/utils/machete_numeric_conversion.cuh
+++ b/custom_ops/gpu_ops/machete/utils/machete_numeric_conversion.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/vllm_numeric_conversion.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/vllm_numeric_conversion.cuh
 #pragma once
 
 #include "cutlass/numeric_conversion.h"
@@ -8,10 +9,10 @@
 
 // this file extends:
 //   https://github.com/NVIDIA/cutlass/blob/cutlass-3.5.0/include/cutlass/numeric_conversion.h
-// with vllm specific type conversions, namely: machete_uint4b8_t, machete_uint8b128_t
-// as well as adds interleaved numeric array converters for specific types.
-// (interleaved numeric array converters can be more efficient for subbyte
-// types)
+// with vllm specific type conversions, namely: machete_uint4b8_t,
+// machete_uint8b128_t as well as adds interleaved numeric array converters for
+// specific types. (interleaved numeric array converters can be more efficient
+// for subbyte types)
 
 namespace cutlass {
 
@@ -19,7 +20,10 @@ namespace cutlass {
 // deinterleaves converted elements based on IlvBlkLayout, interleaving can
 // make subbyte converts more efficient by allowing for efficient extraction
 // of subbyte elements from a 32bit register.
-template <typename IlvBlkLayout, typename T, typename S, int N,
+template <typename IlvBlkLayout,
+          typename T,
+          typename S,
+          int N,
           FloatRoundStyle Round = FloatRoundStyle::round_to_nearest,
           class Enable = void>
 struct InterleavedNumericArrayConverter {
@@ -34,12 +38,17 @@ struct InterleavedNumericArrayConverter {
       if constexpr (std::is_same_v<IlvBlkLayout, void>) {
         printf(
             "Convert %s <= %s (N = %d, IlvBlkLayout = void), not implemented\n",
-            nameof_v<T>, nameof_v<S>, N);
+            nameof_v<T>,
+            nameof_v<S>,
+            N);
       } else {
         printf(
             "Convert %s <= %s (N = %d, size(IlvBlkLayout{}) = %d), not "
             "implemented\n",
-            nameof_v<T>, nameof_v<S>, N, size(IlvBlkLayout{}));
+            nameof_v<T>,
+            nameof_v<S>,
+            N,
+            size(IlvBlkLayout{}));
       }
       __brkpt();
     }
@@ -50,10 +59,17 @@ struct InterleavedNumericArrayConverter {
   result_type operator()(source_type const& s) const { return convert(s); }
 };
 
-template <typename IlvBlkLayout, typename T, typename S, int N,
+template <typename IlvBlkLayout,
+          typename T,
+          typename S,
+          int N,
           FloatRoundStyle Round>
 struct InterleavedNumericArrayConverter<
-    IlvBlkLayout, T, S, N, Round,
+    IlvBlkLayout,
+    T,
+    S,
+    N,
+    Round,
     std::enable_if_t<is_identity_layout<IlvBlkLayout>()>> {
   using Converter = NumericArrayConverter<T, S, N, Round>;
 
@@ -131,18 +147,26 @@ struct ArrayConverterPacked32Bit {
     using ConverterType =
         ArrayConverterPacked32Bit<RegConvert32bit,
                                   typename result_type::Element,
-                                  typename source_type::Element, N>;
+                                  typename source_type::Element,
+                                  N>;
 
     if constexpr (src_elems_per_32bit_reg >= 8) {
-      detail::VectorizedConverter::convert<
-          ConverterType, result_packed_8_t, src_packed_8_t, result_packed_4_t,
-          src_packed_4_t, result_packed_2_t, src_packed_2_t>(result, source);
+      detail::VectorizedConverter::convert<ConverterType,
+                                           result_packed_8_t,
+                                           src_packed_8_t,
+                                           result_packed_4_t,
+                                           src_packed_4_t,
+                                           result_packed_2_t,
+                                           src_packed_2_t>(result, source);
     } else if constexpr (src_elems_per_32bit_reg >= 4) {
-      detail::VectorizedConverter::convert<ConverterType, result_packed_4_t,
-                                           src_packed_4_t, result_packed_2_t,
+      detail::VectorizedConverter::convert<ConverterType,
+                                           result_packed_4_t,
+                                           src_packed_4_t,
+                                           result_packed_2_t,
                                            src_packed_2_t>(result, source);
     } else {
-      detail::VectorizedConverter::convert<ConverterType, result_packed_2_t,
+      detail::VectorizedConverter::convert<ConverterType,
+                                           result_packed_2_t,
                                            src_packed_2_t>(result, source);
     }
 
@@ -152,10 +176,22 @@ struct ArrayConverterPacked32Bit {
 
 // Convert 8 4bit values packed into a 32bit register to 8 8bit values packed
 // into 2 32bit register.
-template <uint8_t LUT0, uint8_t LUT1, uint8_t LUT2, uint8_t LUT3,    //
-          uint8_t LUT4, uint8_t LUT5, uint8_t LUT6, uint8_t LUT7,    //
-          uint8_t LUT8, uint8_t LUT9, uint8_t LUT10, uint8_t LUT11,  //
-          uint8_t LUT12, uint8_t LUT13, uint8_t LUT14, uint8_t LUT15>
+template <uint8_t LUT0,
+          uint8_t LUT1,
+          uint8_t LUT2,
+          uint8_t LUT3,  //
+          uint8_t LUT4,
+          uint8_t LUT5,
+          uint8_t LUT6,
+          uint8_t LUT7,  //
+          uint8_t LUT8,
+          uint8_t LUT9,
+          uint8_t LUT10,
+          uint8_t LUT11,  //
+          uint8_t LUT12,
+          uint8_t LUT13,
+          uint8_t LUT14,
+          uint8_t LUT15>
 CUTLASS_DEVICE cutlass::AlignedArray<uint32_t, 2> lut_4bit_to_8bit_convert(
     uint32_t src) {
   cutlass::AlignedArray<uint32_t, 2> r;
@@ -202,7 +238,11 @@ CUTLASS_DEVICE cutlass::AlignedArray<uint32_t, 2> lut_4bit_to_8bit_convert(
         "  prmt.b32 %0, low, high, %6;\n"
         "}\n"
         : "=r"(r[ii])
-        : "n"(LOW_0), "n"(LOW_1), "n"(HIGH_0), "n"(HIGH_1), "r"(lut_idx),
+        : "n"(LOW_0),
+          "n"(LOW_1),
+          "n"(HIGH_0),
+          "n"(HIGH_1),
+          "r"(lut_idx),
           "r"(final_prmt_idx));
   }
 
@@ -222,10 +262,22 @@ struct NumericArrayConverter<int8_t, machete_uint4b8_t, N, Round> {
     template <typename PackedResultType>
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       // [-8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7] as int8s
-      auto r = lut_4bit_to_8bit_convert<0xF8, 0xF9, 0xFA, 0xFB,  //
-                                        0xFC, 0xFD, 0xFE, 0xFF,  //
-                                        0x00, 0x01, 0x02, 0x03,  //
-                                        0x04, 0x05, 0x06, 0x07>(src_[0]);
+      auto r = lut_4bit_to_8bit_convert<0xF8,
+                                        0xF9,
+                                        0xFA,
+                                        0xFB,  //
+                                        0xFC,
+                                        0xFD,
+                                        0xFE,
+                                        0xFF,  //
+                                        0x00,
+                                        0x01,
+                                        0x02,
+                                        0x03,  //
+                                        0x04,
+                                        0x05,
+                                        0x06,
+                                        0x07>(src_[0]);
       return reinterpret_cast<PackedResultType&>(r);
     };
   };
@@ -233,7 +285,8 @@ struct NumericArrayConverter<int8_t, machete_uint4b8_t, N, Round> {
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -244,7 +297,10 @@ struct NumericArrayConverter<int8_t, machete_uint4b8_t, N, Round> {
 
 // for Array<cutlass::float_e4m3_t, N> <= Array<machete_uint4b8_t, N>
 template <FloatRoundStyle Round, int N>
-struct NumericArrayConverter<cutlass::float_e4m3_t, machete_uint4b8_t, N, Round> {
+struct NumericArrayConverter<cutlass::float_e4m3_t,
+                             machete_uint4b8_t,
+                             N,
+                             Round> {
   using result_type = Array<cutlass::float_e4m3_t, N>;
   using source_type = Array<machete_uint4b8_t, N>;
 
@@ -255,10 +311,22 @@ struct NumericArrayConverter<cutlass::float_e4m3_t, machete_uint4b8_t, N, Round>
     template <typename PackedResultType>
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       // [-8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7] as fp8s
-      auto r = lut_4bit_to_8bit_convert<0xD0, 0xCE, 0xCC, 0xCA,  //
-                                        0xC8, 0xC4, 0xC0, 0xB8,  //
-                                        0x00, 0x38, 0x40, 0x44,  //
-                                        0x48, 0x4A, 0x4C, 0x4E>(src_[0]);
+      auto r = lut_4bit_to_8bit_convert<0xD0,
+                                        0xCE,
+                                        0xCC,
+                                        0xCA,  //
+                                        0xC8,
+                                        0xC4,
+                                        0xC0,
+                                        0xB8,  //
+                                        0x00,
+                                        0x38,
+                                        0x40,
+                                        0x44,  //
+                                        0x48,
+                                        0x4A,
+                                        0x4C,
+                                        0x4E>(src_[0]);
       return reinterpret_cast<PackedResultType&>(r);
     };
   };
@@ -266,7 +334,8 @@ struct NumericArrayConverter<cutlass::float_e4m3_t, machete_uint4b8_t, N, Round>
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -285,9 +354,9 @@ struct NumericArrayConverter<cutlass::half_t, machete_uint4b8_t, N, Round> {
     template <typename PackedResultType>
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       uint32_t src = src_[0];
-      using RegArray =
-          cutlass::AlignedArray<uint32_t, PackedResultType::kElements / 2,
-                                sizeof(PackedResultType)>;
+      using RegArray = cutlass::AlignedArray<uint32_t,
+                                             PackedResultType::kElements / 2,
+                                             sizeof(PackedResultType)>;
       RegArray r;
 
       // Below constructs the following temporary:
@@ -359,7 +428,8 @@ struct NumericArrayConverter<cutlass::half_t, machete_uint4b8_t, N, Round> {
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -372,8 +442,11 @@ struct NumericArrayConverter<cutlass::half_t, machete_uint4b8_t, N, Round> {
 //   for IlvdLayout: (2, 4):(4, 1)
 template <FloatRoundStyle Round, int N>
 struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
-                                        cutlass::half_t, machete_uint4b8_t, N,
-                                        Round, void> {
+                                        cutlass::half_t,
+                                        machete_uint4b8_t,
+                                        N,
+                                        Round,
+                                        void> {
   using IlvdLayout = Layout<Shape<_2, _4>, Stride<_4, _1>>;
   static_assert(N % size(IlvdLayout{}) == 0);
 
@@ -387,9 +460,9 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
     template <typename PackedResultType>
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       uint32_t src = src_[0];
-      using RegArray =
-          cutlass::AlignedArray<uint32_t, PackedResultType::kElements / 2,
-                                sizeof(PackedResultType)>;
+      using RegArray = cutlass::AlignedArray<uint32_t,
+                                             PackedResultType::kElements / 2,
+                                             sizeof(PackedResultType)>;
       RegArray r;
 
       static_assert(PackedResultType::kElements <= size(IlvdLayout{}));
@@ -449,7 +522,8 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -462,7 +536,10 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
 //   for IlvdLayout: (2, 4):(4, 1)
 template <FloatRoundStyle Round, int N>
 struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
-                                        cutlass::half_t, uint4_t, N, Round,
+                                        cutlass::half_t,
+                                        uint4_t,
+                                        N,
+                                        Round,
                                         void> {
   using IlvdLayout = Layout<Shape<_2, _4>, Stride<_4, _1>>;
   static_assert(N % size(IlvdLayout{}) == 0);
@@ -477,9 +554,9 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
     template <typename PackedResultType>
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       uint32_t src = src_[0];
-      using RegArray =
-          cutlass::AlignedArray<uint32_t, PackedResultType::kElements / 2,
-                                sizeof(PackedResultType)>;
+      using RegArray = cutlass::AlignedArray<uint32_t,
+                                             PackedResultType::kElements / 2,
+                                             sizeof(PackedResultType)>;
       RegArray r;
 
       static_assert(PackedResultType::kElements <= size(IlvdLayout{}));
@@ -538,7 +615,8 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -558,19 +636,19 @@ struct NumericArrayConverter<cutlass::half_t, machete_uint8b128_t, N, Round> {
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       uint32_t src = src_[0];
       // Hold output FP16s in reg. We need 1 reg for every 2 elements
-      using RegArray =
-          cutlass::AlignedArray<uint32_t, PackedResultType::kElements / 2,
-                                sizeof(PackedResultType)>;
+      using RegArray = cutlass::AlignedArray<uint32_t,
+                                             PackedResultType::kElements / 2,
+                                             sizeof(PackedResultType)>;
       RegArray r;
 
       uint32_t const prmt_indices[2] = {0x5150, 0x5352};
       static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
 
       for (int ii = 0; ii < RegArray::kElements; ++ii) {
-        asm volatile("prmt.b32 %0,%1,%2,%3;\n"
-                     : "=r"(r[ii])
-                     : "r"(src), "n"(start_byte_for_fp16),
-                       "r"(prmt_indices[ii]));
+        asm volatile(
+            "prmt.b32 %0,%1,%2,%3;\n"
+            : "=r"(r[ii])
+            : "r"(src), "n"(start_byte_for_fp16), "r"(prmt_indices[ii]));
       }
 
       // -128 is folded into bias subtraction, i.e. the 0x80 in the low bytes
@@ -589,7 +667,8 @@ struct NumericArrayConverter<cutlass::half_t, machete_uint8b128_t, N, Round> {
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -631,7 +710,8 @@ struct NumericArrayConverter<float, machete_uint8b128_t, N, Round> {
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -656,9 +736,9 @@ struct NumericArrayConverter<cutlass::bfloat16_t, machete_uint4b8_t, N, Round> {
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       uint32_t src_reg = src_[0];
       // Hold output BF16s in reg. We need 1 reg for every 2 elements
-      using RegArray =
-          cutlass::AlignedArray<uint32_t, PackedResultType::kElements / 2,
-                                sizeof(PackedResultType)>;
+      using RegArray = cutlass::AlignedArray<uint32_t,
+                                             PackedResultType::kElements / 2,
+                                             sizeof(PackedResultType)>;
       RegArray r;
       uint32_t src_reg_shifted = src_reg >> 4;
 
@@ -720,7 +800,8 @@ struct NumericArrayConverter<cutlass::bfloat16_t, machete_uint4b8_t, N, Round> {
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -733,8 +814,11 @@ struct NumericArrayConverter<cutlass::bfloat16_t, machete_uint4b8_t, N, Round> {
 //   for IlvdLayout: (2, 4):(4, 1)
 template <FloatRoundStyle Round, int N>
 struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
-                                        cutlass::bfloat16_t, machete_uint4b8_t, N,
-                                        Round, void> {
+                                        cutlass::bfloat16_t,
+                                        machete_uint4b8_t,
+                                        N,
+                                        Round,
+                                        void> {
   using IlvdLayout = Layout<Shape<_2, _4>, Stride<_4, _1>>;
   static_assert(N % size(IlvdLayout{}) == 0);
 
@@ -746,9 +830,9 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
     template <typename PackedResultType>
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       uint32_t src = src_[0];
-      using RegArray =
-          cutlass::AlignedArray<uint32_t, PackedResultType::kElements / 2,
-                                sizeof(PackedResultType)>;
+      using RegArray = cutlass::AlignedArray<uint32_t,
+                                             PackedResultType::kElements / 2,
+                                             sizeof(PackedResultType)>;
       RegArray r;
 
       static_assert(PackedResultType::kElements <= size(IlvdLayout{}));
@@ -789,7 +873,8 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -802,7 +887,10 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
 //   for IlvdLayout: (2, 4):(4, 1)
 template <FloatRoundStyle Round, int N>
 struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
-                                        cutlass::bfloat16_t, uint4_t, N, Round,
+                                        cutlass::bfloat16_t,
+                                        uint4_t,
+                                        N,
+                                        Round,
                                         void> {
   using IlvdLayout = Layout<Shape<_2, _4>, Stride<_4, _1>>;
   static_assert(N % size(IlvdLayout{}) == 0);
@@ -815,9 +903,9 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
     template <typename PackedResultType>
     CUTLASS_DEVICE static PackedResultType convert(Array<uint32_t, 1> src_) {
       uint32_t src = src_[0];
-      using RegArray =
-          cutlass::AlignedArray<uint32_t, PackedResultType::kElements / 2,
-                                sizeof(PackedResultType)>;
+      using RegArray = cutlass::AlignedArray<uint32_t,
+                                             PackedResultType::kElements / 2,
+                                             sizeof(PackedResultType)>;
       RegArray r;
 
       static_assert(PackedResultType::kElements <= size(IlvdLayout{}));
@@ -858,7 +946,8 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
@@ -869,7 +958,10 @@ struct InterleavedNumericArrayConverter<Layout<Shape<_2, _4>, Stride<_4, _1>>,
 
 // for Array<cutlass::bfloat16_t, N> <= Array<machete_uint8b128_t, N>
 template <FloatRoundStyle Round, int N>
-struct NumericArrayConverter<cutlass::bfloat16_t, machete_uint8b128_t, N, Round> {
+struct NumericArrayConverter<cutlass::bfloat16_t,
+                             machete_uint8b128_t,
+                             N,
+                             Round> {
   using result_type = Array<cutlass::bfloat16_t, N>;
   using source_type = Array<machete_uint8b128_t, N>;
   static FloatRoundStyle const round_style = Round;
@@ -898,13 +990,17 @@ struct NumericArrayConverter<cutlass::bfloat16_t, machete_uint8b128_t, N, Round>
         "Invalid PackedSrcType/PackedResultType must be 2 or 4 to use private "
         "convert dispatch.");
 
-    NumericArrayConverter<float, machete_uint8b128_t, PackedResultType::kElements,
+    NumericArrayConverter<float,
+                          machete_uint8b128_t,
+                          PackedResultType::kElements,
                           Round>
         convert_uint8_to_f32;
     Array<float, PackedResultType::kElements> tmp =
         convert_uint8_to_f32(source);
-    NumericArrayConverter<cutlass::bfloat16_t, float,
-                          PackedResultType::kElements, Round>
+    NumericArrayConverter<cutlass::bfloat16_t,
+                          float,
+                          PackedResultType::kElements,
+                          Round>
         convert_f32_to_bf16_;
     return convert_f32_to_bf16_(tmp);
   }
@@ -915,11 +1011,14 @@ struct NumericArrayConverter<cutlass::bfloat16_t, machete_uint8b128_t, N, Round>
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
     result_type result;
-    using ConverterType =
-        NumericArrayConverter<typename result_type::Element,
-                              typename source_type::Element, N, Round>;
-    detail::VectorizedConverter::convert<ConverterType, result_packed_4_t,
-                                         src_packed_4_t, result_packed_2_t,
+    using ConverterType = NumericArrayConverter<typename result_type::Element,
+                                                typename source_type::Element,
+                                                N,
+                                                Round>;
+    detail::VectorizedConverter::convert<ConverterType,
+                                         result_packed_4_t,
+                                         src_packed_4_t,
+                                         result_packed_2_t,
                                          src_packed_2_t>(result, source);
 
     return result;
@@ -944,8 +1043,10 @@ struct NumericArrayConverter<int8_t, cutlass::half_t, N, Round> {
     CUTLASS_DEVICE static PackedResultType convert(
         Array<uint32_t, src_regs> src) {
       // Hold output int8s in reg. We need 1 reg for every 4 elements
-      using RegArray = cutlass::AlignedArray<
-          uint32_t, std::max(PackedResultType::kElements / 4, size_t(1))>;
+      using RegArray =
+          cutlass::AlignedArray<uint32_t,
+                                std::max(PackedResultType::kElements / 4,
+                                         size_t(1))>;
       RegArray r;
 
       static constexpr uint32_t MAGIC_BIAS_ = 0x64806480;
@@ -963,10 +1064,10 @@ struct NumericArrayConverter<int8_t, cutlass::half_t, N, Round> {
       uint32_t uint8s;
       static constexpr uint32_t MASK_0246 = 0x6420;
       static constexpr uint32_t UINT8s_TO_INT8s_MASK = 0x80808080;
-      asm volatile("prmt.b32 %0,%1,%2,%3;\n"
-                   : "=r"(uint8s)
-                   : "r"(src[0]), "r"((src_regs > 1) ? src[1] : src[0]),
-                     "n"(MASK_0246));
+      asm volatile(
+          "prmt.b32 %0,%1,%2,%3;\n"
+          : "=r"(uint8s)
+          : "r"(src[0]), "r"((src_regs > 1) ? src[1] : src[0]), "n"(MASK_0246));
 
       uint32_t int8s = (uint8s ^ UINT8s_TO_INT8s_MASK);
 
@@ -977,7 +1078,8 @@ struct NumericArrayConverter<int8_t, cutlass::half_t, N, Round> {
  public:
   CUTLASS_DEVICE
   static result_type convert(source_type const& source) {
-    return ArrayConverterPacked32Bit<RegConvert, typename result_type::Element,
+    return ArrayConverterPacked32Bit<RegConvert,
+                                     typename result_type::Element,
                                      typename source_type::Element,
                                      N>::convert(source);
   }
diff --git a/custom_ops/gpu_ops/machete/utils/machete_type_utils.cuh b/custom_ops/gpu_ops/machete/utils/machete_type_utils.cuh
index f97bfdaaf1..34f65ceac9 100644
--- a/custom_ops/gpu_ops/machete/utils/machete_type_utils.cuh
+++ b/custom_ops/gpu_ops/machete/utils/machete_type_utils.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/vllm_numeric_conversion.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/vllm_numeric_conversion.cuh
 #include "cutlass/bfloat16.h"
 #include "cutlass/half.h"
 #include "cuda_bf16.h"
diff --git a/custom_ops/gpu_ops/machete/utils/paddle_utils.hpp b/custom_ops/gpu_ops/machete/utils/paddle_utils.hpp
index 31c04e2404..230452debe 100644
--- a/custom_ops/gpu_ops/machete/utils/paddle_utils.hpp
+++ b/custom_ops/gpu_ops/machete/utils/paddle_utils.hpp
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/torch_utils.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/torch_utils.hpp
 #pragma once
 
 #include "helper.h"
@@ -16,7 +17,9 @@ namespace cute {
 namespace detail {
 
 template <class T, class F, class G, int... I>
-CUTE_HOST_DEVICE constexpr auto tapply_with_idx(T&& t, F&& f, G&& g,
+CUTE_HOST_DEVICE constexpr auto tapply_with_idx(T&& t,
+                                                F&& f,
+                                                G&& g,
                                                 seq<I...>) {
   return g(f(cute::get<I>(static_cast<T&&>(t)), I)...);
 }
@@ -32,7 +35,9 @@ template <class T, class F>
 CUTE_HOST_DEVICE constexpr auto transform_with_idx(T const& t, F&& f) {
   if constexpr (cute::is_tuple<T>::value) {
     return detail::tapply_with_idx(
-        t, f, [](auto const&... a) { return cute::make_tuple(a...); },
+        t,
+        f,
+        [](auto const&... a) { return cute::make_tuple(a...); },
         tuple_seq<T>{});
   } else {
     return f(t);
@@ -53,8 +58,8 @@ CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) {
 // shape of the passed in tensor and the strides are of type `Stride` and
 // contain the strides of the passed in tensor, checking that any static strides
 // in `Stride{}` match the strides of the passed in tensor.
-// If `tensor.shape().size() < rank(Stride{})`, the shape is padded with 1s and the extra
-// strides are set to be 0 or 1.
+// If `tensor.shape().size() < rank(Stride{})`, the shape is padded with 1s and
+// the extra strides are set to be 0 or 1.
 template <typename Stride>
 static inline auto make_cute_layout(paddle::Tensor const& tensor,
                                     std::string_view name = "tensor") {
@@ -65,8 +70,16 @@ static inline auto make_cute_layout(paddle::Tensor const& tensor,
 
         if (idx < tensor.shape().size()) {
           if constexpr (cute::is_static_v<StrideEle>) {
-            PD_CHECK(StrideEle::value == tensor.strides()[idx], "Expected ",
-                        name, ".strides()[", idx, "] to be ", StrideEle::value, ", but got ", tensor.strides()[idx], ". ");
+            PD_CHECK(StrideEle::value == tensor.strides()[idx],
+                     "Expected ",
+                     name,
+                     ".strides()[",
+                     idx,
+                     "] to be ",
+                     StrideEle::value,
+                     ", but got ",
+                     tensor.strides()[idx],
+                     ". ");
             return StrideEle{};
           } else {
             if (tensor.shape()[idx] == 1) {
diff --git a/custom_ops/gpu_ops/machete/utils/scalar_type.h b/custom_ops/gpu_ops/machete/utils/scalar_type.h
index 019d8486c6..57054eb273 100644
--- a/custom_ops/gpu_ops/machete/utils/scalar_type.h
+++ b/custom_ops/gpu_ops/machete/utils/scalar_type.h
@@ -41,15 +41,18 @@ class ScalarType {
     NAN_REPR_ID_MAX
   };
 
-  constexpr ScalarType(uint8_t exponent, uint8_t mantissa, bool signed_,
-                       int32_t bias, bool finite_values_only = false,
+  constexpr ScalarType(uint8_t exponent,
+                       uint8_t mantissa,
+                       bool signed_,
+                       int32_t bias,
+                       bool finite_values_only = false,
                        NanRepr nan_repr = NAN_IEEE_754)
       : exponent(exponent),
         mantissa(mantissa),
         signed_(signed_),
         bias(bias),
         finite_values_only(finite_values_only),
-        nan_repr(nan_repr) {};
+        nan_repr(nan_repr){};
 
   static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) {
     return ScalarType(0, size_bits - 1, true, bias);
@@ -67,16 +70,17 @@ class ScalarType {
   }
 
   // IEEE 754 non-compliant floating point type
-  static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
+  static constexpr ScalarType float_(uint8_t exponent,
+                                     uint8_t mantissa,
                                      bool finite_values_only,
                                      NanRepr nan_repr) {
     // PADDLE_ENFORCE(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
     // PADDLE_ENFORCE(mantissa > 0 && exponent > 0);
     // PADDLE_ENFORCE(nan_repr != NAN_IEEE_754,
-    //             "use `float_IEEE754` constructor for floating point types that "
-    //             "follow IEEE 754 conventions");
-    return ScalarType(exponent, mantissa, true, 0, finite_values_only,
-                      nan_repr);
+    //             "use `float_IEEE754` constructor for floating point types
+    //             that " "follow IEEE 754 conventions");
+    return ScalarType(
+        exponent, mantissa, true, 0, finite_values_only, nan_repr);
   }
 
   uint8_t const exponent;  // size of the exponent field (0 for integer types)
@@ -103,7 +107,9 @@ class ScalarType {
   }
 
   template <typename Fn, typename Init, typename Member, typename... Rest>
-  static constexpr auto reduce_members_helper(Fn f, Init val, Member member,
+  static constexpr auto reduce_members_helper(Fn f,
+                                              Init val,
+                                              Member member,
                                               Rest... rest) {
     auto new_val = f(val, member);
     if constexpr (sizeof...(rest) > 0) {
@@ -116,8 +122,14 @@ class ScalarType {
   template <typename Fn, typename Init>
   constexpr auto reduce_members(Fn f, Init init) const {
     // Should be in constructor order for `from_id`
-    return reduce_members_helper(f, init, exponent, mantissa, signed_, bias,
-                                 finite_values_only, nan_repr);
+    return reduce_members_helper(f,
+                                 init,
+                                 exponent,
+                                 mantissa,
+                                 signed_,
+                                 bias,
+                                 finite_values_only,
+                                 nan_repr);
   };
 
   template <typename Fn, typename Init>
@@ -194,7 +206,8 @@ class ScalarType {
  private:
   double _floating_point_max() const {
     PADDLE_ENFORCE(mantissa <= 52 && exponent <= 11,
-                "Cannot represent max/min as a double for type ", str());
+                   "Cannot represent max/min as a double for type ",
+                   str());
 
     uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
     if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
@@ -204,7 +217,8 @@ class ScalarType {
     uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
     if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
       PADDLE_ENFORCE(exponent < 11,
-                  "Cannot represent max/min as a double for type ", str());
+                     "Cannot represent max/min as a double for type ",
+                     str());
       max_exponent += 1;
     }
 
@@ -336,14 +350,16 @@ static inline constexpr auto kFE2M1f =
     machete::ScalarType::float_(2, 1, true, machete::ScalarType::NAN_NONE);
 static inline constexpr auto kFE3M2f =
     machete::ScalarType::float_(3, 2, true, machete::ScalarType::NAN_NONE);
-static inline constexpr auto kFE4M3fn =
-    machete::ScalarType::float_(4, 3, true, machete::ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE4M3fn = machete::ScalarType::float_(
+    4, 3, true, machete::ScalarType::NAN_EXTD_RANGE_MAX_MIN);
 static inline constexpr auto kFE5M2 = machete::ScalarType::float_IEEE754(5, 2);
 static inline constexpr auto kFE8M7 = machete::ScalarType::float_IEEE754(8, 7);
-static inline constexpr auto kFE5M10 = machete::ScalarType::float_IEEE754(5, 10);
+static inline constexpr auto kFE5M10 =
+    machete::ScalarType::float_IEEE754(5, 10);
 
 // // Fixed width style names, generally following:
-// //  https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57
+// //
+// https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57
 constexpr auto kInt4 = kS4;
 constexpr auto kUint4 = kU4;
 constexpr auto kUint4b8 = kU4B8;
diff --git a/custom_ops/gpu_ops/merge_prefill_decode_output.cu b/custom_ops/gpu_ops/merge_prefill_decode_output.cu
index 6902b72505..a57d72bdf3 100644
--- a/custom_ops/gpu_ops/merge_prefill_decode_output.cu
+++ b/custom_ops/gpu_ops/merge_prefill_decode_output.cu
@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-
 #include "paddle/extension.h"
 
 #ifndef PD_BUILD_STATIC_OP
@@ -20,86 +19,83 @@
 #endif
 
 template <int warps, typename T>
-__global__ void FillEncoderDecoderResKernel(
-        T * encoder_res_data,
-        T * decoder_res_data,
-        const int * seq_lens_encoder,
-        const int * seq_lens_decoder,
-        const int * seq_lens_this_time,
-        const int * cu_seq_q,
-        const int head_num,
-        const int head_dim) {
+__global__ void FillEncoderDecoderResKernel(T *encoder_res_data,
+                                            T *decoder_res_data,
+                                            const int *seq_lens_encoder,
+                                            const int *seq_lens_decoder,
+                                            const int *seq_lens_this_time,
+                                            const int *cu_seq_q,
+                                            const int head_num,
+                                            const int head_dim) {
+  const int bidb = blockIdx.x;
+  const int bidh = blockIdx.y;
+  const int bidt = blockIdx.z * warps;
+  const int tid = threadIdx.x;
+  const int warp_id = tid / 32;
+  const int land_id = tid % 32;
+  const int token_id = bidt + warp_id;
 
-    const int bidb = blockIdx.x;
-    const int bidh = blockIdx.y;
-    const int bidt = blockIdx.z * warps;
-    const int tid = threadIdx.x;
-    const int warp_id = tid / 32;
-    const int land_id = tid % 32;
-    const int token_id = bidt + warp_id;
+  const int seq_len_encoder = seq_lens_encoder[bidb];
+  const int seq_len_decoder = seq_lens_decoder[bidb];
+  const int seq_len_this_time = seq_lens_this_time[bidb];
 
-    const int seq_len_encoder = seq_lens_encoder[bidb];
-    const int seq_len_decoder = seq_lens_decoder[bidb];
-    const int seq_len_this_time = seq_lens_this_time[bidb];
+  if (seq_len_encoder > 0 || seq_len_decoder == 0 ||
+      token_id >= seq_len_this_time) {
+    return;
+  }
 
-    if (seq_len_encoder > 0 || seq_len_decoder == 0 || token_id >= seq_len_this_time) {
-        return;
-    }
+  const int load_idx =
+      ((cu_seq_q[bidb] + token_id) * head_num + bidh) * head_dim + land_id * 4;
 
-    const int load_idx = ((cu_seq_q[bidb] + token_id) * head_num + bidh) * head_dim + land_id * 4;
-
-    *reinterpret_cast<float2*>(encoder_res_data + load_idx) = *reinterpret_cast<float2*>(decoder_res_data + load_idx);
+  *reinterpret_cast<float2 *>(encoder_res_data + load_idx) =
+      *reinterpret_cast<float2 *>(decoder_res_data + load_idx);
 }
 
-void MergePrefillDecodeOutput(
-        const paddle::Tensor &encoder_res,
-        const paddle::Tensor &decoder_res,
-        const paddle::Tensor &seq_lens_encoder,
-        const paddle::Tensor &seq_lens_decoder,
-        const paddle::Tensor &seq_lens_this_time,
-        const paddle::Tensor &cu_seq_q,
-        const int head_num,
-        const int head_dim,
-        const int max_token) {
+void MergePrefillDecodeOutput(const paddle::Tensor &encoder_res,
+                              const paddle::Tensor &decoder_res,
+                              const paddle::Tensor &seq_lens_encoder,
+                              const paddle::Tensor &seq_lens_decoder,
+                              const paddle::Tensor &seq_lens_this_time,
+                              const paddle::Tensor &cu_seq_q,
+                              const int head_num,
+                              const int head_dim,
+                              const int max_token) {
+  if (head_dim != 128) {
+    PD_THROW("Only supported head_dim = 128");
+  }
+  const int batch_size = seq_lens_encoder.shape()[0];
+  constexpr int warps = 4;
+  const int tokens_block = (max_token + warps - 1) / warps;
+  dim3 grid_dims;
+  grid_dims.x = batch_size;
+  grid_dims.y = head_num;
+  grid_dims.z = tokens_block;
 
-    if (head_dim != 128) {
-        PD_THROW("Only supported head_dim = 128");
-    }
-    const int batch_size = seq_lens_encoder.shape()[0];
-    constexpr int warps = 4;
-    const int tokens_block = (max_token + warps - 1) / warps;
-    dim3 grid_dims;
-    grid_dims.x = batch_size;
-    grid_dims.y = head_num;
-    grid_dims.z = tokens_block;
-
-    if (encoder_res.dtype() == paddle::DataType::FLOAT16) {
-        using T = phi::dtype::float16;
-        FillEncoderDecoderResKernel<warps>
-            <<<grid_dims, 128, 0, encoder_res.stream()>>>(
-            const_cast<T*>(encoder_res.data<T>()),
-            const_cast<T*>(decoder_res.data<T>()),
+  if (encoder_res.dtype() == paddle::DataType::FLOAT16) {
+    using T = phi::dtype::float16;
+    FillEncoderDecoderResKernel<warps>
+        <<<grid_dims, 128, 0, encoder_res.stream()>>>(
+            const_cast<T *>(encoder_res.data<T>()),
+            const_cast<T *>(decoder_res.data<T>()),
             seq_lens_encoder.data<int>(),
             seq_lens_decoder.data<int>(),
             seq_lens_this_time.data<int>(),
             cu_seq_q.data<int>(),
             head_num,
-            head_dim
-            );
-    } else if (encoder_res.dtype() == paddle::DataType::BFLOAT16) {
-        using T = phi::dtype::bfloat16;
-        FillEncoderDecoderResKernel<warps>
-            <<<grid_dims, 128, 0, encoder_res.stream()>>>(
-            const_cast<T*>(encoder_res.data<T>()),
-            const_cast<T*>(decoder_res.data<T>()),
+            head_dim);
+  } else if (encoder_res.dtype() == paddle::DataType::BFLOAT16) {
+    using T = phi::dtype::bfloat16;
+    FillEncoderDecoderResKernel<warps>
+        <<<grid_dims, 128, 0, encoder_res.stream()>>>(
+            const_cast<T *>(encoder_res.data<T>()),
+            const_cast<T *>(decoder_res.data<T>()),
             seq_lens_encoder.data<int>(),
             seq_lens_decoder.data<int>(),
             seq_lens_this_time.data<int>(),
             cu_seq_q.data<int>(),
             head_num,
-            head_dim
-            );
-    }
+            head_dim);
+  }
 }
 
 PD_BUILD_STATIC_OP(merge_prefill_decode_output)
@@ -110,8 +106,6 @@ PD_BUILD_STATIC_OP(merge_prefill_decode_output)
              "seq_lens_this_time",
              "cu_seq_q"})
     .Outputs({"res"})
-    .Attrs({"head_num: int",
-            "head_dim: int",
-            "max_token: int"})
+    .Attrs({"head_num: int", "head_dim: int", "max_token: int"})
     .SetInplaceMap({{"encoder_res", "res"}})
     .SetKernelFn(PD_KERNEL(MergePrefillDecodeOutput));
diff --git a/custom_ops/gpu_ops/mla_attn/attention_updater.cuh b/custom_ops/gpu_ops/mla_attn/attention_updater.cuh
index 49f8089d3d..8607fbec08 100644
--- a/custom_ops/gpu_ops/mla_attn/attention_updater.cuh
+++ b/custom_ops/gpu_ops/mla_attn/attention_updater.cuh
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 /*
- * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
- * Dao. Licensed under the BSD 3-Clause.
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar,
+ * Pradeep Ramani, Tri Dao. Licensed under the BSD 3-Clause.
  *
  * Modified by the FlashInfer team.
  */
@@ -30,18 +30,24 @@ using namespace cute;
 
 template <typename T>
 struct MaxOp {
-  __device__ __forceinline__ T operator()(T const& x, T const& y) { return x > y ? x : y; }
+  __device__ __forceinline__ T operator()(T const& x, T const& y) {
+    return x > y ? x : y;
+  }
 };
 
 template <>
 struct MaxOp<float> {
   // This is slightly faster
-  __device__ __forceinline__ float operator()(float const& x, float const& y) { return max(x, y); }
+  __device__ __forceinline__ float operator()(float const& x, float const& y) {
+    return max(x, y);
+  }
 };
 
 template <typename T>
 struct SumOp {
-  __device__ __forceinline__ T operator()(T const& x, T const& y) { return x + y; }
+  __device__ __forceinline__ T operator()(T const& x, T const& y) {
+    return x + y;
+  }
 };
 
 template <int THREADS>
@@ -64,10 +70,16 @@ struct Allreduce<2> {
   }
 };
 
-template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
+template <bool init,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1,
           typename Operator>
-__device__ __forceinline__ void thread_reduce_(Tensor<Engine0, Layout0> const& tensor,
-                                               Tensor<Engine1, Layout1>& summary, Operator& op) {
+__device__ __forceinline__ void thread_reduce_(
+    Tensor<Engine0, Layout0> const& tensor,
+    Tensor<Engine1, Layout1>& summary,
+    Operator& op) {
   static_assert(Layout0::rank == 2, "Only support 2D Tensor");
   static_assert(Layout1::rank == 1, "Only support 1D Tensor");
   CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
@@ -81,9 +93,14 @@ __device__ __forceinline__ void thread_reduce_(Tensor<Engine0, Layout0> const& t
   }
 }
 
-template <typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
+template <typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1,
+          typename Operator>
 __device__ __forceinline__ void quad_allreduce_(Tensor<Engine0, Layout0>& dst,
-                                                Tensor<Engine1, Layout1>& src, Operator& op) {
+                                                Tensor<Engine1, Layout1>& src,
+                                                Operator& op) {
   CUTE_STATIC_ASSERT_V(size(dst) == size(src));
 #pragma unroll
   for (int i = 0; i < size(dst); i++) {
@@ -91,25 +108,38 @@ __device__ __forceinline__ void quad_allreduce_(Tensor<Engine0, Layout0>& dst,
   }
 }
 
-template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
+template <bool init,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1,
           typename Operator>
 __device__ __forceinline__ void reduce_(Tensor<Engine0, Layout0> const& tensor,
-                                        Tensor<Engine1, Layout1>& summary, Operator& op) {
+                                        Tensor<Engine1, Layout1>& summary,
+                                        Operator& op) {
   thread_reduce_<init>(tensor, summary, op);
   quad_allreduce_(summary, summary, op);
 }
 
-template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__device__ __forceinline__ void reduce_max(Tensor<Engine0, Layout0> const& tensor,
-                                           Tensor<Engine1, Layout1>& max) {
+template <bool init,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1>
+__device__ __forceinline__ void reduce_max(
+    Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1>& max) {
   MaxOp<float> max_op;
   reduce_<init>(tensor, max, max_op);
 }
 
-template <bool init, bool warp_reduce = true, typename Engine0, typename Layout0, typename Engine1,
+template <bool init,
+          bool warp_reduce = true,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
           typename Layout1>
-__device__ __forceinline__ void reduce_sum(Tensor<Engine0, Layout0> const& tensor,
-                                           Tensor<Engine1, Layout1>& sum) {
+__device__ __forceinline__ void reduce_sum(
+    Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1>& sum) {
   SumOp<float> sum_op;
   thread_reduce_<init>(tensor, sum, sum_op);
   if constexpr (warp_reduce) {
@@ -117,9 +147,12 @@ __device__ __forceinline__ void reduce_sum(Tensor<Engine0, Layout0> const& tenso
   }
 }
 
-template <typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__forceinline__ __device__ void apply_exp2(Tensor<Engine0, Layout0>& tensor,
-                                           Tensor<Engine1, Layout1> const& max) {
+template <typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1>
+__forceinline__ __device__ void apply_exp2(
+    Tensor<Engine0, Layout0>& tensor, Tensor<Engine1, Layout1> const& max) {
   static_assert(Layout0::rank == 2, "Only support 2D Tensor");
   static_assert(Layout1::rank == 1, "Only support 1D Tensor");
   CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
@@ -133,10 +166,14 @@ __forceinline__ __device__ void apply_exp2(Tensor<Engine0, Layout0>& tensor,
   }
 }
 
-template <typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__forceinline__ __device__ void scale_apply_exp2(Tensor<Engine0, Layout0>& tensor,
-                                                 Tensor<Engine1, Layout1> const& max,
-                                                 const float scale) {
+template <typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1>
+__forceinline__ __device__ void scale_apply_exp2(
+    Tensor<Engine0, Layout0>& tensor,
+    Tensor<Engine1, Layout1> const& max,
+    const float scale) {
   static_assert(Layout0::rank == 2, "Only support 2D Tensor");
   static_assert(Layout1::rank == 1, "Only support 1D Tensor");
   CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
@@ -154,11 +191,13 @@ __forceinline__ __device__ void scale_apply_exp2(Tensor<Engine0, Layout0>& tenso
 template <int NUM_ROWS_PER_THREAD, bool WITH_SCALE>
 struct OnlineSoftmax {
   constexpr static float fill_value = -5e4;
-  using TensorT = decltype(make_tensor<float>(Shape<Int<NUM_ROWS_PER_THREAD>>{}));
+  using TensorT =
+      decltype(make_tensor<float>(Shape<Int<NUM_ROWS_PER_THREAD>>{}));
   TensorT row_max, row_sum, scores_scale;
   float sm_scale_log2;
 
-  CUTLASS_DEVICE OnlineSoftmax(float sm_scale_log2) : sm_scale_log2(sm_scale_log2) {
+  CUTLASS_DEVICE OnlineSoftmax(float sm_scale_log2)
+      : sm_scale_log2(sm_scale_log2) {
     clear(scores_scale);
   };
 
@@ -166,8 +205,10 @@ struct OnlineSoftmax {
 
   template <bool init, typename Tensor0>
   __forceinline__ __device__ TensorT update(Tensor0& acc_s) {
-    // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))
-    Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+    // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M),
+    // ncol=(2, V, MMA_N))
+    Tensor scores =
+        make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
 
     static_assert(decltype(size<0>(scores))::value == NUM_ROWS_PER_THREAD);
     if constexpr (init) {
@@ -188,7 +229,8 @@ struct OnlineSoftmax {
       for (int mi = 0; mi < size(row_max); ++mi) {
         float scores_max_cur = row_max(mi);
         if constexpr (WITH_SCALE) {
-          scores_scale(mi) = __expf((scores_max_prev(mi) - scores_max_cur) * sm_scale_log2);
+          scores_scale(mi) =
+              __expf((scores_max_prev(mi) - scores_max_cur) * sm_scale_log2);
         } else {
           scores_scale(mi) = __expf(scores_max_prev(mi) - scores_max_cur);
         }
@@ -208,8 +250,10 @@ struct OnlineSoftmax {
 
   template <typename Tensor0>
   __forceinline__ __device__ TensorT finalize(Tensor0& acc_s) {
-    // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))
-    Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+    // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M),
+    // ncol=(2, V, MMA_N))
+    Tensor scores =
+        make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
     static_assert(decltype(size<0>(scores))::value == NUM_ROWS_PER_THREAD);
     SumOp<float> sum_op;
     quad_allreduce_(row_sum, row_sum, sum_op);
@@ -225,9 +269,12 @@ struct OnlineSoftmax {
 
   template <typename Tensor1>
   __forceinline__ __device__ void rescale_o(Tensor1& acc_o) {
-    // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))
-    Tensor acc_o_rowcol = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
-    static_assert(decltype(size<0>(acc_o_rowcol))::value == NUM_ROWS_PER_THREAD);
+    // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2,
+    // MMA_K))
+    Tensor acc_o_rowcol =
+        make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+    static_assert(decltype(size<0>(acc_o_rowcol))::value ==
+                  NUM_ROWS_PER_THREAD);
 #pragma unroll
     for (int mi = 0; mi < size(row_max); ++mi) {
 #pragma unroll
@@ -238,10 +285,14 @@ struct OnlineSoftmax {
   };
 
   template <typename Tensor1, typename Tensor2>
-  __forceinline__ __device__ void rescale_o(Tensor1& acc_o, Tensor2& scores_scale_input) {
-    // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))
-    Tensor acc_o_rowcol = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
-    static_assert(decltype(size<0>(acc_o_rowcol))::value == NUM_ROWS_PER_THREAD);
+  __forceinline__ __device__ void rescale_o(Tensor1& acc_o,
+                                            Tensor2& scores_scale_input) {
+    // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2,
+    // MMA_K))
+    Tensor acc_o_rowcol =
+        make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+    static_assert(decltype(size<0>(acc_o_rowcol))::value ==
+                  NUM_ROWS_PER_THREAD);
 #pragma unroll
     for (int mi = 0; mi < size(row_max); ++mi) {
 #pragma unroll
diff --git a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h
index 62f5d5e9f9..72d13a4974 100644
--- a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h
+++ b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h
@@ -37,14 +37,21 @@ template <typename T>
 void BatchMLAWithPagedKVCacheKernel(
     const AppendAttnMetaData& meta_data,
     const paddle::Tensor& q,  // [token_num, q_head_num, head_dim]
-    const paddle::Tensor& latent_cache,  // [max_block_num, q_head_num, block_size, head_dim]
+    const paddle::Tensor&
+        latent_cache,  // [max_block_num, q_head_num, block_size, head_dim]
     const paddle::optional<paddle::Tensor>& attn_mask,
-    const paddle::optional<paddle::Tensor>& cache_k_scale,  // [num_kv_heads, head_dim]
-    const paddle::optional<paddle::Tensor>& cache_v_scale,  // [num_kv_heads, head_dim]
-    const paddle::optional<paddle::Tensor>& cache_k_zp,  // [num_kv_heads, head_dim]
-    const paddle::optional<paddle::Tensor>& cache_v_zp,  // [num_kv_heads, head_dim]
-    const paddle::optional<paddle::Tensor>& shift_bias,  // [num_kv_heads, head_dim]
-    const paddle::optional<paddle::Tensor>& smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>&
+        smooth_weight,  // [num_kv_heads, head_dim]
     const paddle::Tensor& seq_lens_this_time,
     const paddle::Tensor& seq_lens_decoder,
     const paddle::Tensor& cu_seqlens_q,
diff --git a/custom_ops/gpu_ops/mla_attn/epilogue.cuh b/custom_ops/gpu_ops/mla_attn/epilogue.cuh
index 72d1b55704..959150b917 100644
--- a/custom_ops/gpu_ops/mla_attn/epilogue.cuh
+++ b/custom_ops/gpu_ops/mla_attn/epilogue.cuh
@@ -13,13 +13,12 @@
 // limitations under the License.
 
 /*
- * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
- * Dao. Licensed under the BSD 3-Clause.
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar,
+ * Pradeep Ramani, Tri Dao. Licensed under the BSD 3-Clause.
  *
  * Modified by the FlashInfer team.
  */
 
-
 #ifndef ATTENTION_HOPPER_EPILOGUE_CUH_
 #define ATTENTION_HOPPER_EPILOGUE_CUH_
 
@@ -45,7 +44,8 @@ struct CollectiveEpilogue {
   static constexpr int BLOCK_SHAPE_Q = Ktraits::BLOCK_SHAPE_Q;
   static constexpr int BLOCK_SHAPE_KV = Ktraits::BLOCK_SHAPE_KV;
   static constexpr int HEAD_DIM_VO = Ktraits::HEAD_DIM_VO;
-  using TileShape_PDV = Shape<Int<BLOCK_SHAPE_Q>, Int<HEAD_DIM_VO>, Int<BLOCK_SHAPE_KV>>;
+  using TileShape_PDV =
+      Shape<Int<BLOCK_SHAPE_Q>, Int<HEAD_DIM_VO>, Int<BLOCK_SHAPE_KV>>;
 
   static constexpr int NUM_WARPS = Ktraits::NUM_WARPS;
   static constexpr int NUM_THREADS = NUM_WARPS * cutlass::NumThreadsPerWarp;
@@ -53,13 +53,18 @@ struct CollectiveEpilogue {
   static constexpr int NUM_COPY_THREADS = cutlass::NumThreadsPerWarpGroup;
   static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS;
 
-  using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_PDV{})),
-                                   decltype(cute::get<1>(TileShape_PDV{}))>());
-  using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{})));
+  using SmemLayoutAtomO =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               DTypeO,
+               decltype(cute::get<0>(TileShape_PDV{})),
+               decltype(cute::get<1>(TileShape_PDV{}))>());
+  using SmemLayoutO =
+      decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{})));
 
   using SmemCopyAtomO = Copy_Atom<cute::SM90_U32x4_STSM_N, DTypeO>;
-  using SharedStorage = cute::array_aligned<DTypeO, cute::cosize_v<SmemLayoutO>>;
+  using SharedStorage =
+      cute::array_aligned<DTypeO, cute::cosize_v<SmemLayoutO>>;
 
   using ShapeT = cute::Shape<int32_t, int32_t, int32_t>;
   using StrideT = cute::Shape<int32_t, _1, int32_t>;
@@ -76,22 +81,29 @@ struct CollectiveEpilogue {
   using GmemTiledCopyOTMA = cute::SM90_TMA_STORE;
   using TMA_O = decltype(make_tma_copy(
       GmemTiledCopyOTMA{},
-      make_tensor(make_gmem_ptr(static_cast<DTypeO*>(nullptr)), ShapeT{}, StrideT{}), SmemLayoutO{},
-      select<0, 1>(TileShape_PDV{}), _1{}));  // no mcast for O
+      make_tensor(
+          make_gmem_ptr(static_cast<DTypeO*>(nullptr)), ShapeT{}, StrideT{}),
+      SmemLayoutO{},
+      select<0, 1>(TileShape_PDV{}),
+      _1{}));  // no mcast for O
 
-  static constexpr int VEC_SIZE = cute::ceil_div(128, sizeof_bits_v<DTypeO>); // 8
+  static constexpr int VEC_SIZE =
+      cute::ceil_div(128, sizeof_bits_v<DTypeO>);  // 8
   static_assert(HEAD_DIM_VO % VEC_SIZE == 0);
-  static constexpr int NUM_THREADS_PER_ROW = HEAD_DIM_VO / VEC_SIZE; // 64
+  static constexpr int NUM_THREADS_PER_ROW = HEAD_DIM_VO / VEC_SIZE;  // 64
   static_assert(NUM_MMA_THREADS % NUM_THREADS_PER_ROW == 0);
   static constexpr int NUM_ROWS = NUM_MMA_THREADS / NUM_THREADS_PER_ROW;
-  using TiledCopyOAtom = cute::Copy_Atom<cute::UniversalCopy<cutlass::uint128_t>, DTypeO>;
+  using TiledCopyOAtom =
+      cute::Copy_Atom<cute::UniversalCopy<cutlass::uint128_t>, DTypeO>;
   using TiledCopyOThrLayout = decltype(cute::make_layout(
-      cute::make_shape(Int<NUM_ROWS>{}, Int<NUM_THREADS_PER_ROW>{}), LayoutRight{}));
-  using TiledCopyOValLayout =
-      decltype(cute::make_layout(cute::make_shape(_1{}, Int<VEC_SIZE>{}), LayoutRight{}));
+      cute::make_shape(Int<NUM_ROWS>{}, Int<NUM_THREADS_PER_ROW>{}),
+      LayoutRight{}));
+  using TiledCopyOValLayout = decltype(cute::make_layout(
+      cute::make_shape(_1{}, Int<VEC_SIZE>{}), LayoutRight{}));
   using TiledCopyO =
-      decltype(make_tiled_copy(TiledCopyOAtom{}, TiledCopyOThrLayout{},  // Thr layout
-                               TiledCopyOValLayout{}                     // Val layout
+      decltype(make_tiled_copy(TiledCopyOAtom{},
+                               TiledCopyOThrLayout{},  // Thr layout
+                               TiledCopyOValLayout{}   // Val layout
                                ));
   struct Arguments {
     DTypeO* O_ptr;
@@ -115,7 +127,9 @@ struct CollectiveEpilogue {
   CUTLASS_DEVICE
   static void prefetch_tma_descriptors(Params const& epilogue_params) {}
 
-  template <typename SharedStorage, typename FrgTensorO, typename FrgTensorLSE,
+  template <typename SharedStorage,
+            typename FrgTensorO,
+            typename FrgTensorLSE,
             typename TiledMma>
   CUTLASS_DEVICE void store(Params const& epilogue_params,
                             FrgTensorO const& tOrO,
@@ -133,25 +147,35 @@ struct CollectiveEpilogue {
                             const int max_draft_token_num,
                             const int o_stride_bsz) {
     const int num_chunks = cute::ceil_div(kv_len, chunk_size);
-    Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{});
+    Tensor sO =
+        make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{});
     auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma);
     auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx);
 
     Tensor tOrO_out = convert_type<DTypeO>(tOrO);
-    Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out);  // ((Atom,AtomNum), MMA_M, MMA_N)
-    Tensor taccOsO = smem_thr_copy_O.partition_D(sO);     // ((Atom,AtomNum),PIPE_M,PIPE_N)
+    Tensor taccOrO =
+        smem_thr_copy_O.retile_S(tOrO_out);  // ((Atom,AtomNum), MMA_M, MMA_N)
+    Tensor taccOsO =
+        smem_thr_copy_O.partition_D(sO);  // ((Atom,AtomNum),PIPE_M,PIPE_N)
     // make sure gemm done
-    cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS,
-                                      /*id=*/static_cast<int>(NamedBarriers::kValueEmpty));
+    cutlass::arch::NamedBarrier::sync(
+        NUM_MMA_THREADS,
+        /*id=*/static_cast<int>(NamedBarriers::kValueEmpty));
     // r2s
     cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);
     // make sure r2s done
-    cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS,
-                                      /*id=*/static_cast<int>(NamedBarriers::kValueEmpty));
+    cutlass::arch::NamedBarrier::sync(
+        NUM_MMA_THREADS,
+        /*id=*/static_cast<int>(NamedBarriers::kValueEmpty));
     TiledCopyO gmem_tiled_copy_O;
-    auto O_ptr = num_chunks == 1 ? epilogue_params.O_ptr + start_token_idx * o_stride_bsz : epilogue_params.O_ptr_tmp + (tile_idx * bsz + bid) * max_draft_token_num * o_stride_bsz;
+    auto O_ptr = num_chunks == 1
+                     ? epilogue_params.O_ptr + start_token_idx * o_stride_bsz
+                     : epilogue_params.O_ptr_tmp + (tile_idx * bsz + bid) *
+                                                       max_draft_token_num *
+                                                       o_stride_bsz;
     Tensor mO = make_tensor(make_gmem_ptr(O_ptr), epilogue_params.layout_O);
-    Tensor gO = local_tile(mO, select<0, 1>(TileShape_PDV{}), make_coord(_, _0{}))(_, _, _0{});
+    Tensor gO = local_tile(
+        mO, select<0, 1>(TileShape_PDV{}), make_coord(_, _0{}))(_, _, _0{});
     Tensor cO = make_identity_tensor(gO.shape());  // (O, D) -> (o_idx, d_idx)
     ThrCopy thr_copy_O = gmem_tiled_copy_O.get_slice(thread_idx);
     Tensor tOgO = thr_copy_O.partition_D(gO);  // (CPY, CPY_O, CPY_D)
diff --git a/custom_ops/gpu_ops/mla_attn/kernel_traits.cuh b/custom_ops/gpu_ops/mla_attn/kernel_traits.cuh
index 116ccb7c88..dcbdcc7474 100644
--- a/custom_ops/gpu_ops/mla_attn/kernel_traits.cuh
+++ b/custom_ops/gpu_ops/mla_attn/kernel_traits.cuh
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 /*
- * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
- * Dao. Licensed under the BSD 3-Clause.
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar,
+ * Pradeep Ramani, Tri Dao. Licensed under the BSD 3-Clause.
  *
  * Modified by the FlashInfer team.
  */
@@ -36,15 +36,29 @@ namespace mla_attn {
 
 using namespace cute;
 
-template <typename MainloopPipeline, typename MainloopPipelineQ, class DTypeQ, class DTypeKV, class DTypeQKAccum, class DTypeOut, class IdType,
-          int BLOCK_SHAPE_KV, class SmemLayoutQ, class SmemLayoutK, class SmemLayoutP, class SmemLayoutRow, class SmemLayoutO>
+template <typename MainloopPipeline,
+          typename MainloopPipelineQ,
+          class DTypeQ,
+          class DTypeKV,
+          class DTypeQKAccum,
+          class DTypeOut,
+          class IdType,
+          int BLOCK_SHAPE_KV,
+          class SmemLayoutQ,
+          class SmemLayoutK,
+          class SmemLayoutP,
+          class SmemLayoutRow,
+          class SmemLayoutO>
 struct alignas(16) SharedStorageQKVO {
   alignas(16) cute::array_aligned<DTypeQ, cute::cosize_v<SmemLayoutQ>> smem_q;
   alignas(16) cute::array_aligned<DTypeQ, cute::cosize_v<SmemLayoutP>> smem_p;
-  alignas(16) cute::array_aligned<DTypeQKAccum, cute::cosize_v<SmemLayoutRow>> smem_scale;
+  alignas(16) cute::array_aligned<DTypeQKAccum,
+                                  cute::cosize_v<SmemLayoutRow>> smem_scale;
   union {
-    alignas(16) cute::array_aligned<DTypeKV, cute::cosize_v<SmemLayoutK>> smem_kv;
-    alignas(16) cute::array_aligned<DTypeOut, cute::cosize_v<SmemLayoutO>> smem_o;
+    alignas(
+        16) cute::array_aligned<DTypeKV, cute::cosize_v<SmemLayoutK>> smem_kv;
+    alignas(
+        16) cute::array_aligned<DTypeOut, cute::cosize_v<SmemLayoutO>> smem_o;
   };
   struct {
     alignas(16) typename MainloopPipelineQ::SharedStorage pipeline_q;
@@ -52,10 +66,19 @@ struct alignas(16) SharedStorageQKVO {
   };
 };
 
-template <bool USE_TMA_LOAD_KV_, int HEAD_DIM_QK_, int HEAD_DIM_VO_, int GROUP_SIZE_, int BLOCK_SHAPE_Q_, int BLOCK_SHAPE_KV_,
-          int NUM_STAGES_, typename DTypeQ_, typename DTypeKV_, typename DTypeO_, typename IdType_, typename NV_TYPE_>
+template <bool USE_TMA_LOAD_KV_,
+          int HEAD_DIM_QK_,
+          int HEAD_DIM_VO_,
+          int GROUP_SIZE_,
+          int BLOCK_SHAPE_Q_,
+          int BLOCK_SHAPE_KV_,
+          int NUM_STAGES_,
+          typename DTypeQ_,
+          typename DTypeKV_,
+          typename DTypeO_,
+          typename IdType_,
+          typename NV_TYPE_>
 struct AttentionKernelTraits {
-
   using DTypeQ = DTypeQ_;
   using DTypeKV = DTypeKV_;
   using DTypeO = DTypeO_;
@@ -64,7 +87,6 @@ struct AttentionKernelTraits {
   using DTypePVAccum = float;
   using NV_TYPE = NV_TYPE_;
 
-
   static constexpr bool USE_TMA_LOAD_KV = USE_TMA_LOAD_KV_;
   static constexpr int GROUP_SIZE = GROUP_SIZE_;
   static constexpr int BLOCK_SHAPE_Q = BLOCK_SHAPE_Q_;
@@ -80,82 +102,135 @@ struct AttentionKernelTraits {
   static constexpr int NUM_THREADS = 384;
   static constexpr int NUM_PRODUCER_THREADS = 128;
 
-  using TileShape_QKD = Shape<Int<BLOCK_SHAPE_Q>, Int<BLOCK_SHAPE_KV>, Int<HEAD_DIM_QK>>;
-  using TileShape_PDV = Shape<Int<BLOCK_SHAPE_Q>, Int<HEAD_DIM_VO>, Int<BLOCK_SHAPE_KV>>;
+  using TileShape_QKD =
+      Shape<Int<BLOCK_SHAPE_Q>, Int<BLOCK_SHAPE_KV>, Int<HEAD_DIM_QK>>;
+  using TileShape_PDV =
+      Shape<Int<BLOCK_SHAPE_Q>, Int<HEAD_DIM_VO>, Int<BLOCK_SHAPE_KV>>;
 
   static constexpr int NUM_STAGES = NUM_STAGES_;
 
   using AtomLayoutQKD = Layout<Shape<Int<BLOCK_SHAPE_Q / 64>, _1, _1>>;
   using AtomLayoutPV = Layout<Shape<Int<BLOCK_SHAPE_Q / 64>, _2, _1>>;
   using TiledMmaQK = decltype(cute::make_tiled_mma(
-      cute::GMMA::ss_op_selector<DTypeQ, DTypeKV, DTypeQKAccum, TileShape_QKD>(), AtomLayoutQKD{}));
+      cute::GMMA::
+          ss_op_selector<DTypeQ, DTypeKV, DTypeQKAccum, TileShape_QKD>(),
+      AtomLayoutQKD{}));
   using TiledMmaPV = decltype(cute::make_tiled_mma(
-      cute::GMMA::rs_op_selector<DTypeKV, DTypeKV, /*ElementAccum=*/DTypePVAccum, TileShape_PDV,
-                                 GMMA::Major::K, GMMA::Major::MN>(),
+      cute::GMMA::rs_op_selector<DTypeKV,
+                                 DTypeKV,
+                                 /*ElementAccum=*/DTypePVAccum,
+                                 TileShape_PDV,
+                                 GMMA::Major::K,
+                                 GMMA::Major::MN>(),
       AtomLayoutPV{}));
   using TiledMmaPVSS = decltype(cute::make_tiled_mma(
-      cute::GMMA::ss_op_selector<DTypeKV, DTypeKV, /*ElementAccum=*/DTypePVAccum, TileShape_PDV,
-                                 GMMA::Major::K, GMMA::Major::MN>(),
+      cute::GMMA::ss_op_selector<DTypeKV,
+                                 DTypeKV,
+                                 /*ElementAccum=*/DTypePVAccum,
+                                 TileShape_PDV,
+                                 GMMA::Major::K,
+                                 GMMA::Major::MN>(),
       AtomLayoutPV{}));
 
   static constexpr int NUM_MMA_THREADS = size(TiledMmaPV{});
 
-  using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})),
-                                   decltype(cute::get<2>(TileShape_QKD{}))>());
-  using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_QKD{})));
+  using SmemLayoutAtomQ =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               DTypeQ,
+               decltype(cute::get<0>(TileShape_QKD{})),
+               decltype(cute::get<2>(TileShape_QKD{}))>());
+  using SmemLayoutQ =
+      decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_QKD{})));
 
-  using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeKV, decltype(cute::get<1>(TileShape_QKD{})),
-                                   decltype(cute::get<2>(TileShape_QKD{}))>());
-  using SmemLayoutK = decltype(tile_to_shape(
-      SmemLayoutAtomK{},
-      make_shape(shape<1>(TileShape_QKD{}), shape<2>(TileShape_QKD{}), Int<NUM_STAGES>{})));
+  using SmemLayoutAtomK =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               DTypeKV,
+               decltype(cute::get<1>(TileShape_QKD{})),
+               decltype(cute::get<2>(TileShape_QKD{}))>());
+  using SmemLayoutK =
+      decltype(tile_to_shape(SmemLayoutAtomK{},
+                             make_shape(shape<1>(TileShape_QKD{}),
+                                        shape<2>(TileShape_QKD{}),
+                                        Int<NUM_STAGES>{})));
   using SmemLayoutVt = decltype(composition(
-      SmemLayoutK{}, make_ordered_layout(make_shape(get<2>(TileShape_QKD{}),
-                                                    get<1>(TileShape_QKD{}), Int<NUM_STAGES>{}),
-                                         Step<_2, _1, _3>{})));
-  using SmemLayoutAtomV = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeKV, decltype(cute::get<2>(TileShape_PDV{})),
-                                   decltype(cute::get<1>(TileShape_PDV{}))>());
+      SmemLayoutK{},
+      make_ordered_layout(make_shape(get<2>(TileShape_QKD{}),
+                                     get<1>(TileShape_QKD{}),
+                                     Int<NUM_STAGES>{}),
+                          Step<_2, _1, _3>{})));
+  using SmemLayoutAtomV =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               DTypeKV,
+               decltype(cute::get<2>(TileShape_PDV{})),
+               decltype(cute::get<1>(TileShape_PDV{}))>());
   using SmemLayoutV = decltype(tile_to_shape(
       SmemLayoutAtomV{},
       make_shape(get<2>(TileShape_PDV{}), get<1>(TileShape_PDV{}), Int<1>{})));
 
   // Note this is the transpose in terms of the view, not in terms of memory.
   using SmemLayoutVtOneStage = decltype(composition(
-      SmemLayoutV{}, make_ordered_layout(make_shape(get<1>(TileShape_PDV{}),
-                                                    get<2>(TileShape_PDV{}), Int<1>{}),
-                                         Step<_2, _1, _3>{})));
+      SmemLayoutV{},
+      make_ordered_layout(
+          make_shape(
+              get<1>(TileShape_PDV{}), get<2>(TileShape_PDV{}), Int<1>{}),
+          Step<_2, _1, _3>{})));
 
-  using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_PDV{})),
-                                   decltype(cute::get<1>(TileShape_PDV{}))>());
-  using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{})));
+  using SmemLayoutAtomO =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               DTypeO,
+               decltype(cute::get<0>(TileShape_PDV{})),
+               decltype(cute::get<1>(TileShape_PDV{}))>());
+  using SmemLayoutO =
+      decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{})));
 
   using SmemCopyAtom = Copy_Atom<cute::SM90_U32x4_STSM_N, DTypeQ>;
 
   static constexpr bool IS_CTA_32 = (BLOCK_SHAPE_KV == 32);
   using SmemLayoutRowOneStage = Layout<Shape<_2, Int<128>>, Stride<_1, _2>>;
-  using SmemLayoutRowTwoStage = Layout<Shape<_2, Int<128>, _2>, Stride<_1, _2, _256>>;
-  using SmemLayoutRow = std::conditional_t<IS_CTA_32, SmemLayoutRowTwoStage, SmemLayoutRowOneStage>;
+  using SmemLayoutRowTwoStage =
+      Layout<Shape<_2, Int<128>, _2>, Stride<_1, _2, _256>>;
+  using SmemLayoutRow = std::
+      conditional_t<IS_CTA_32, SmemLayoutRowTwoStage, SmemLayoutRowOneStage>;
 
-  using SmemLayoutAtomP = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
-                                   GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})),
-                                   decltype(cute::get<1>(TileShape_QKD{}))>());
-  using SmemLayoutPSSOneStage = decltype(tile_to_shape(SmemLayoutAtomP{}, select<0, 1>(TileShape_QKD{})));
-  using SmemLayoutPSSTwoStage = decltype(tile_to_shape(SmemLayoutAtomP{}, make_shape(Int<BLOCK_SHAPE_Q>{}, Int<BLOCK_SHAPE_KV>{}, Int<2>{})));
-  using SmemLayoutP = std::conditional_t<IS_CTA_32, SmemLayoutPSSTwoStage, SmemLayoutPSSOneStage>;
+  using SmemLayoutAtomP =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               DTypeQ,
+               decltype(cute::get<0>(TileShape_QKD{})),
+               decltype(cute::get<1>(TileShape_QKD{}))>());
+  using SmemLayoutPSSOneStage =
+      decltype(tile_to_shape(SmemLayoutAtomP{}, select<0, 1>(TileShape_QKD{})));
+  using SmemLayoutPSSTwoStage = decltype(tile_to_shape(
+      SmemLayoutAtomP{},
+      make_shape(Int<BLOCK_SHAPE_Q>{}, Int<BLOCK_SHAPE_KV>{}, Int<2>{})));
+  using SmemLayoutP = std::
+      conditional_t<IS_CTA_32, SmemLayoutPSSTwoStage, SmemLayoutPSSOneStage>;
 
   using MainloopPipelineQ = typename cutlass::PipelineAsync<1>;
   using PipelineStateQ = typename cutlass::PipelineState<1>;
   using MainloopPipeline =
-      std::conditional_t<USE_TMA_LOAD_KV, typename cutlass::PipelineTmaAsync<NUM_STAGES>,
+      std::conditional_t<USE_TMA_LOAD_KV,
+                         typename cutlass::PipelineTmaAsync<NUM_STAGES>,
                          typename cutlass::PipelineAsync<NUM_STAGES>>;
   using PipelineState = typename cutlass::PipelineState<NUM_STAGES>;
 
-  using SharedStorage = SharedStorageQKVO<MainloopPipeline, MainloopPipelineQ, DTypeQ, DTypeKV, DTypeQKAccum, DTypeO, IdType, BLOCK_SHAPE_KV,
-                                          SmemLayoutQ, SmemLayoutK, SmemLayoutP, SmemLayoutRow, SmemLayoutO>;
+  using SharedStorage = SharedStorageQKVO<MainloopPipeline,
+                                          MainloopPipelineQ,
+                                          DTypeQ,
+                                          DTypeKV,
+                                          DTypeQKAccum,
+                                          DTypeO,
+                                          IdType,
+                                          BLOCK_SHAPE_KV,
+                                          SmemLayoutQ,
+                                          SmemLayoutK,
+                                          SmemLayoutP,
+                                          SmemLayoutRow,
+                                          SmemLayoutO>;
 };
 
 }  // namespace mla_attn
diff --git a/custom_ops/gpu_ops/mla_attn/mainloop_load.cuh b/custom_ops/gpu_ops/mla_attn/mainloop_load.cuh
index 4ee350f766..20fa31e5d5 100644
--- a/custom_ops/gpu_ops/mla_attn/mainloop_load.cuh
+++ b/custom_ops/gpu_ops/mla_attn/mainloop_load.cuh
@@ -51,28 +51,34 @@ struct CollectiveMainloop {
   static constexpr int HEAD_DIM_VO = Ktraits::HEAD_DIM_VO;
 
   using GmemTiledCopyKV = cute::SM90_TMA_LOAD;
-  static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(DTypeQ); // 8
-  static_assert(HEAD_DIM_QK % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad"); // 576 512
-  static constexpr int kGmemThreadsPerRow = 64 / kGmemElemsPerLoad; // 8
-  using AlignmentTypeQ = cute::uint_byte_t<static_cast<int>(sizeof(DTypeQ)) * kGmemElemsPerLoad>;
-  using GmemCopyAtomQ = cute::Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<AlignmentTypeQ>, DTypeQ>;
+  static constexpr int kGmemElemsPerLoad =
+      sizeof(cute::uint128_t) / sizeof(DTypeQ);  // 8
+  static_assert(HEAD_DIM_QK % kGmemElemsPerLoad == 0,
+                "kHeadDim must be a multiple of kGmemElemsPerLoad");  // 576 512
+  static constexpr int kGmemThreadsPerRow = 64 / kGmemElemsPerLoad;   // 8
+  using AlignmentTypeQ =
+      cute::uint_byte_t<static_cast<int>(sizeof(DTypeQ)) * kGmemElemsPerLoad>;
+  using GmemCopyAtomQ =
+      cute::Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<AlignmentTypeQ>, DTypeQ>;
   static constexpr int kNThreadsLoad = Ktraits::NUM_PRODUCER_THREADS;
-  static_assert(kNThreadsLoad % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow");
-  using GmemLayoutAtom = Layout<
-            Shape<Int<kNThreadsLoad / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>, // 32, 8
-            Stride<Int<kGmemThreadsPerRow>, _1>>;
+  static_assert(kNThreadsLoad % kGmemThreadsPerRow == 0,
+                "kNThreads must be a multiple of kGmemThreadsPerRow");
+  using GmemLayoutAtom = Layout<Shape<Int<kNThreadsLoad / kGmemThreadsPerRow>,
+                                      Int<kGmemThreadsPerRow>>,  // 32, 8
+                                Stride<Int<kGmemThreadsPerRow>, _1>>;
   using GmemTiledCopy = decltype(make_tiled_copy(
-          GmemCopyAtomQ{},
-          GmemLayoutAtom{},
-          Layout<Shape<_1, _8>>{}));  // Val layout, 8 vals per read
+      GmemCopyAtomQ{},
+      GmemLayoutAtom{},
+      Layout<Shape<_1, _8>>{}));  // Val layout, 8 vals per read
 
-  using GmemLayoutAtomQ = Layout<
-            Shape<Int<Ktraits::NUM_PRODUCER_THREADS / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>, // 32, 8
-            Stride<Int<kGmemThreadsPerRow>, _1>>;
+  using GmemLayoutAtomQ =
+      Layout<Shape<Int<Ktraits::NUM_PRODUCER_THREADS / kGmemThreadsPerRow>,
+                   Int<kGmemThreadsPerRow>>,  // 32, 8
+             Stride<Int<kGmemThreadsPerRow>, _1>>;
   using GmemTiledCopyQ = decltype(make_tiled_copy(
-          GmemCopyAtomQ{},
-          GmemLayoutAtomQ{},
-          Layout<Shape<_1, _8>>{}));  // Val layout, 8 vals per read
+      GmemCopyAtomQ{},
+      GmemLayoutAtomQ{},
+      Layout<Shape<_1, _8>>{}));  // Val layout, 8 vals per read
 
   using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
   using SmemLayoutAtomQ = typename Ktraits::SmemLayoutAtomQ;
@@ -95,13 +101,12 @@ struct CollectiveMainloop {
 
   using TMA_KV = decltype(make_tma_copy(
       GmemTiledCopyKV{},
-      make_tensor(
-          make_gmem_ptr(static_cast<DTypeKV const*>(nullptr)),
-          repeat_like(StrideT{}, int32_t(0)), StrideT{}
-      ),
+      make_tensor(make_gmem_ptr(static_cast<DTypeKV const*>(nullptr)),
+                  repeat_like(StrideT{}, int32_t(0)),
+                  StrideT{}),
       take<0, 2>(SmemLayoutK{}),
       select<1, 2>(TileShape_QKD{}),
-      _1{})); // no mcast for KV
+      _1{}));  // no mcast for KV
 
   static constexpr bool USE_TMA_LOAD_KV = Ktraits::USE_TMA_LOAD_KV;
   using MainloopPipeline = typename Ktraits::MainloopPipeline;
@@ -112,10 +117,10 @@ struct CollectiveMainloop {
   using PipelineParamsQ = typename MainloopPipelineQ::Params;
   using PipelineStateQ = typename MainloopPipelineQ::PipelineState;
 
-  static constexpr uint32_t TmaTransactionBytesQ =
-      static_cast<uint32_t>(size(SmemLayoutQ{}) * cutlass::sizeof_bits_v<DTypeQ> / 8);
-  static constexpr uint32_t TmaTransactionBytesKV =
-      static_cast<uint32_t>(size(take<0, 2>(SmemLayoutK{})) * cutlass::sizeof_bits_v<DTypeKV> / 8);
+  static constexpr uint32_t TmaTransactionBytesQ = static_cast<uint32_t>(
+      size(SmemLayoutQ{}) * cutlass::sizeof_bits_v<DTypeQ> / 8);
+  static constexpr uint32_t TmaTransactionBytesKV = static_cast<uint32_t>(
+      size(take<0, 2>(SmemLayoutK{})) * cutlass::sizeof_bits_v<DTypeKV> / 8);
 
   // Host side kernel arguments
   struct Arguments {
@@ -155,7 +160,7 @@ struct CollectiveMainloop {
     LayoutQT layout_Q;
     LayoutT layout_KV;
     LayoutMDT layout_MD;
-    DTypeQ *Q_ptr;
+    DTypeQ* Q_ptr;
     DTypeKV* KV_ptr;
     DTypeMD* m_ptr;
     DTypeMD* d_ptr;
@@ -188,8 +193,11 @@ struct CollectiveMainloop {
     TMA_KV tma_load_KV;
     if constexpr (USE_TMA_LOAD_KV) {
       Tensor mKV = make_tensor(make_gmem_ptr(args.KV_ptr), args.layout_KV);
-      tma_load_KV =
-          make_tma_copy(GmemTiledCopyKV{}, mKV, SmemLayoutK{}(_, _, _0{}), select<1, 2>(TileShape_QKD{}), _1{});
+      tma_load_KV = make_tma_copy(GmemTiledCopyKV{},
+                                  mKV,
+                                  SmemLayoutK{}(_, _, _0{}),
+                                  select<1, 2>(TileShape_QKD{}),
+                                  _1{});
     }
     return {args.layout_Q,
             args.layout_KV,
@@ -220,14 +228,14 @@ struct CollectiveMainloop {
             // args.chunk_size,
             args.chunk_num,
             args.max_draft_token_num,
-            tma_load_KV
-            };
+            tma_load_KV};
   }
 
   CUTLASS_DEVICE
   static void prefetch_tma_descriptors(Params const& mainloop_params) {
     if constexpr (USE_TMA_LOAD_KV) {
-      cute::prefetch_tma_descriptor(mainloop_params.tma_load_KV.get_tma_descriptor());
+      cute::prefetch_tma_descriptor(
+          mainloop_params.tma_load_KV.get_tma_descriptor());
     }
   }
 
@@ -240,10 +248,12 @@ struct CollectiveMainloop {
                              const int bid) {
     int start_q_token_idx = mainloop_params.cumsum_q_seqlens[bid];
     int offset_Q = mainloop_params.q_stride_bsz * start_q_token_idx;
-    Tensor mQ = make_tensor(make_gmem_ptr(mainloop_params.Q_ptr + offset_Q), mainloop_params.layout_Q);
-    Tensor gQ =
-        local_tile(mQ, select<0, 2>(TileShape_QKD{}), make_coord(_, _0{}))(_, _, _0{});
-    Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+    Tensor mQ = make_tensor(make_gmem_ptr(mainloop_params.Q_ptr + offset_Q),
+                            mainloop_params.layout_Q);
+    Tensor gQ = local_tile(
+        mQ, select<0, 2>(TileShape_QKD{}), make_coord(_, _0{}))(_, _, _0{});
+    Tensor sQ =
+        make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
     Tensor cQ = cute::make_identity_tensor(gQ.shape());
 
     GmemTiledCopyQ gmem_tiled_copy_q;
@@ -263,7 +273,8 @@ struct CollectiveMainloop {
 
     pipeline_q.producer_acquire(smem_pipe_write_q);
     copy_if(gmem_tiled_copy_q, q_predicate_fn, tQgQiGroup, tQsQiGroup);
-    pipeline_q.producer_commit(smem_pipe_write_q, cutlass::arch::cpasync_barrier_arrive);
+    pipeline_q.producer_commit(smem_pipe_write_q,
+                               cutlass::arch::cpasync_barrier_arrive);
     ++smem_pipe_write_q;
   }
 
@@ -276,31 +287,46 @@ struct CollectiveMainloop {
                               const int kv_len,
                               const int tile_idx) {
     int thread_idx = threadIdx.x;
-    int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (thread_idx / 32) % 4, 0);
-    Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
-    Tensor mKV = make_tensor(make_gmem_ptr(mainloop_params.KV_ptr), mainloop_params.layout_KV);
-    Tensor gKV = local_tile(mKV, make_shape(get<1>(TileShape_QKD{}), get<2>(TileShape_QKD{})), make_coord(_, _))(_, _, _0{}, _0{}, _);
+    int warp_idx_in_warpgroup =
+        __shfl_sync(0xffffffff, (thread_idx / 32) % 4, 0);
+    Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()),
+                            SmemLayoutK{});
+    Tensor mKV = make_tensor(make_gmem_ptr(mainloop_params.KV_ptr),
+                             mainloop_params.layout_KV);
+    Tensor gKV =
+        local_tile(mKV,
+                   make_shape(get<1>(TileShape_QKD{}), get<2>(TileShape_QKD{})),
+                   make_coord(_, _))(_, _, _0{}, _0{}, _);
     GmemTiledCopy gmem_tiled_copy_kv;
     auto gmem_thr_copy_kv = gmem_tiled_copy_kv.get_slice(thread_idx);
 
     static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
     const int start_len = tile_idx * mainloop_params.chunk_size_device[0];
     const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
-    const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size_device[0], kv_len), BLOCK_SHAPE_KV) - 1;
+    const int end_tile_idx =
+        cute::ceil_div(
+            min(start_len + mainloop_params.chunk_size_device[0], kv_len),
+            BLOCK_SHAPE_KV) -
+        1;
 
-    auto kv_block_tables = make_tensor(make_gmem_ptr(mainloop_params.kv_block_tables), make_layout(make_shape(mainloop_params.bsz, mainloop_params.max_block_num_per_seq), make_stride(mainloop_params.max_block_num_per_seq, 1)));
+    auto kv_block_tables = make_tensor(
+        make_gmem_ptr(mainloop_params.kv_block_tables),
+        make_layout(make_shape(mainloop_params.bsz,
+                               mainloop_params.max_block_num_per_seq),
+                    make_stride(mainloop_params.max_block_num_per_seq, 1)));
 
     Tensor tKgK = gmem_thr_copy_kv.partition_S(gKV);
     Tensor tKsK = gmem_thr_copy_kv.partition_S(sK);
 
-    for (int kv_tile_idx = end_tile_idx; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
+    for (int kv_tile_idx = end_tile_idx; kv_tile_idx >= start_tile_idx;
+         --kv_tile_idx) {
       const int block_idx = kv_block_tables(bid, kv_tile_idx);
       pipeline_kv.producer_acquire(smem_pipe_write_kv);
       Tensor tKgKiGroup = flatten_1(tKgK(_, _, _, block_idx));
-      Tensor tKsKiGroup =
-          flatten_1(tKsK(_, _, _, smem_pipe_write_kv.index()));
+      Tensor tKsKiGroup = flatten_1(tKsK(_, _, _, smem_pipe_write_kv.index()));
       copy(gmem_tiled_copy_kv, tKgKiGroup, tKsKiGroup);
-      pipeline_kv.producer_commit(smem_pipe_write_kv, cutlass::arch::cpasync_barrier_arrive);
+      pipeline_kv.producer_commit(smem_pipe_write_kv,
+                                  cutlass::arch::cpasync_barrier_arrive);
       ++smem_pipe_write_kv;
     }
   }
@@ -314,32 +340,51 @@ struct CollectiveMainloop {
                                   const int kv_len,
                                   const int tile_idx) {
     int thread_idx = threadIdx.x;
-    Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+    Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()),
+                            SmemLayoutK{});
 
-    Tensor mKV = mainloop_params.tma_load_KV.get_tma_tensor(mainloop_params.layout_KV.shape());
+    Tensor mKV = mainloop_params.tma_load_KV.get_tma_tensor(
+        mainloop_params.layout_KV.shape());
 
     // Prepare the TMA loads
-    Tensor gKV = local_tile(mKV, make_shape(get<1>(TileShape_QKD{}), get<2>(TileShape_QKD{})), make_coord(_, _))(_, _, _0{}, _0{}, _);
-    auto [tKgK, tKsK] =
-        tma_partition(mainloop_params.tma_load_KV, _0{}, Layout<_1>{},
-                      group_modes<0, 2>(sK), group_modes<0, 2>(gKV));
+    Tensor gKV =
+        local_tile(mKV,
+                   make_shape(get<1>(TileShape_QKD{}), get<2>(TileShape_QKD{})),
+                   make_coord(_, _))(_, _, _0{}, _0{}, _);
+    auto [tKgK, tKsK] = tma_partition(mainloop_params.tma_load_KV,
+                                      _0{},
+                                      Layout<_1>{},
+                                      group_modes<0, 2>(sK),
+                                      group_modes<0, 2>(gKV));
 
     static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
     const int start_len = tile_idx * mainloop_params.chunk_size_device[0];
     const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
-    const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size_device[0], kv_len), BLOCK_SHAPE_KV) - 1;
+    const int end_tile_idx =
+        cute::ceil_div(
+            min(start_len + mainloop_params.chunk_size_device[0], kv_len),
+            BLOCK_SHAPE_KV) -
+        1;
 
-    auto kv_block_tables = make_tensor(make_gmem_ptr(mainloop_params.kv_block_tables), make_layout(make_shape(mainloop_params.bsz, mainloop_params.max_block_num_per_seq), make_stride(mainloop_params.max_block_num_per_seq, 1)));
+    auto kv_block_tables = make_tensor(
+        make_gmem_ptr(mainloop_params.kv_block_tables),
+        make_layout(make_shape(mainloop_params.bsz,
+                               mainloop_params.max_block_num_per_seq),
+                    make_stride(mainloop_params.max_block_num_per_seq, 1)));
 
     int lane_predicate = cute::elect_one_sync();
 
     if (lane_predicate) {
 #pragma unroll 2
-      for (int kv_tile_idx = end_tile_idx; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
+      for (int kv_tile_idx = end_tile_idx; kv_tile_idx >= start_tile_idx;
+           --kv_tile_idx) {
         const int block_idx = kv_block_tables(bid, kv_tile_idx);
         pipeline_kv.producer_acquire(smem_pipe_write_kv);
-        copy(mainloop_params.tma_load_KV.with(*pipeline_kv.producer_get_barrier(smem_pipe_write_kv), /*mcast_mask=*/0),
-             tKgK(_, block_idx), tKsK(_, smem_pipe_write_kv.index()));
+        copy(mainloop_params.tma_load_KV.with(
+                 *pipeline_kv.producer_get_barrier(smem_pipe_write_kv),
+                 /*mcast_mask=*/0),
+             tKgK(_, block_idx),
+             tKsK(_, smem_pipe_write_kv.index()));
         ++smem_pipe_write_kv;
       }
     }
diff --git a/custom_ops/gpu_ops/mla_attn/mainloop_mma.cuh b/custom_ops/gpu_ops/mla_attn/mainloop_mma.cuh
index b6cdabb8fd..f18559b9d9 100644
--- a/custom_ops/gpu_ops/mla_attn/mainloop_mma.cuh
+++ b/custom_ops/gpu_ops/mla_attn/mainloop_mma.cuh
@@ -25,8 +25,16 @@
 
 namespace mla_attn {
 
-template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQ,
-          typename PipelineState, typename PipelineStateQ, typename SharedStorage, typename FrgTensorO, typename AttentionUpdater>
+template <typename Ktraits,
+          bool CAUSAL,
+          typename Params,
+          typename MainloopPipeline,
+          typename MainloopPipelineQ,
+          typename PipelineState,
+          typename PipelineStateQ,
+          typename SharedStorage,
+          typename FrgTensorO,
+          typename AttentionUpdater>
 CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
                             MainloopPipelineQ pipeline_q,
                             PipelineStateQ& smem_pipe_read_q,
@@ -57,26 +65,38 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
   using SmemLayoutVtOneStage = typename Ktraits::SmemLayoutVtOneStage;
   static_assert(is_rmem<FrgTensorO>::value, "O tensor must be rmem resident.");
 
-  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size_device[0]);
+  const int chunk_num_this_seq =
+      cute::ceil_div(kv_len, mainloop_params.chunk_size_device[0]);
 
   static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
   static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
 
-  Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
-  Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
-  Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutVtOneStage{});
-  Tensor sVt_s2 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
-  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutP{});
-  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
-  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr), mainloop_params.layout_MD)(tile_idx, _); // (bsz * draft_token_num * num_head)
-  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr), mainloop_params.layout_MD)(tile_idx, _);
+  Tensor sQ =
+      make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+  Tensor sK =
+      make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+  Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()),
+                              SmemLayoutVtOneStage{});
+  Tensor sVt_s2 = make_tensor(
+      make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE),
+      SmemLayoutVtOneStage{});
+  Tensor sPSS =
+      make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutP{});
+  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()),
+                               SmemLayoutRow{});
+  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr),
+                          mainloop_params.layout_MD)(
+      tile_idx, _);  // (bsz * draft_token_num * num_head)
+  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr),
+                          mainloop_params.layout_MD)(tile_idx, _);
 
   typename Ktraits::TiledMmaQK tiled_mma_qk;
   auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx);
   auto smem_tiled_copy_P = make_tiled_copy_C(SmemCopyAtom{}, tiled_mma_qk);
   auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(thread_idx);
   Tensor tPsP = smem_thr_copy_P.partition_D(sPSS);
-  Tensor tScalesScale = s_scale(_, thread_idx % cutlass::NumThreadsPerWarpGroup);
+  Tensor tScalesScale =
+      s_scale(_, thread_idx % cutlass::NumThreadsPerWarpGroup);
 
   typename Ktraits::TiledMmaPVSS tiled_mma_pv_ss;
   auto threadMmaPVSS = tiled_mma_pv_ss.get_thread_slice(thread_idx);
@@ -86,7 +106,11 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
 
   const int start_len = tile_idx * mainloop_params.chunk_size_device[0];
   const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
-  const int end_tile_idx =cute::ceil_div(min(start_len + mainloop_params.chunk_size_device[0], kv_len), BLOCK_SHAPE_KV) - 1;
+  const int end_tile_idx =
+      cute::ceil_div(
+          min(start_len + mainloop_params.chunk_size_device[0], kv_len),
+          BLOCK_SHAPE_KV) -
+      1;
   int kv_tile_idx = end_tile_idx;
 
   auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
@@ -100,19 +124,24 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
     Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ);
     Tensor tSrK = threadMmaQK.partition_fragment_B(sK);
 
-    constexpr int n_masking_steps = !CAUSAL ? 1 : cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) + 1;
-    auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; };
+    constexpr int n_masking_steps =
+        !CAUSAL ? 1 : cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) + 1;
+    auto col_limit_right = [&](int qo_idx) {
+      return qo_idx + 1 + kv_len - qo_len;
+    };
     bool is_first_step = true;
     // wait q
     consumer_wait(pipeline_q, smem_pipe_read_q);
-    Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+    Tensor tSrS =
+        partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
 #pragma unroll 1
-    for (int masking_step = n_masking_steps; kv_tile_idx >= start_tile_idx; --masking_step, --kv_tile_idx) {
+    for (int masking_step = n_masking_steps; kv_tile_idx >= start_tile_idx;
+         --masking_step, --kv_tile_idx) {
       // wait kv
       consumer_wait(pipeline_kv, smem_pipe_read_kv);
       // gemm qk
-      gemm</*init=*/true, /*wg_wait=*/0>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
-                                         tSrS);
+      gemm</*init=*/true, /*wg_wait=*/0>(
+          tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()), tSrS);
       // mask
       if (masking_step > 0) {
         Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
@@ -134,7 +163,9 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
       }
 
       // update s (exp(s - m))
-      Tensor scale_o = is_first_step ? attention_updater.update</*init=*/true>(tSrS) : attention_updater.update</*init=*/false>(tSrS);
+      Tensor scale_o = is_first_step
+                           ? attention_updater.update</*init=*/true>(tSrS)
+                           : attention_updater.update</*init=*/false>(tSrS);
       is_first_step = false;
 
       Tensor convert_tSrS = convert_type<DTypeKV>(tSrS);
@@ -146,50 +177,61 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
       // r2s fence wgmma
       cutlass::arch::fence_view_async_shared();
       // make sure r2s all done
-      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+      cutlass::arch::NamedBarrier::sync(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
 
       attention_updater.rescale_o(tOrO, scale_o);
 
       // pv gemm
       if (smem_pipe_read_kv.index() == 0) {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2,
-                                            tOrV1(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/0>(
+            tiled_mma_pv_ss, tOrP_CS2, tOrV1(_, _, _, _0{}), tOrO);
       } else {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2,
-                                            tOrV2(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/0>(
+            tiled_mma_pv_ss, tOrP_CS2, tOrV2(_, _, _, _0{}), tOrO);
       }
 
       pipeline_kv.consumer_release(smem_pipe_read_kv);
       ++smem_pipe_read_kv;
       // sync WG1 WG2
-      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2Sync));
+      cutlass::arch::NamedBarrier::sync(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWG1WG2Sync));
     }
     // release q
     pipeline_q.consumer_release(smem_pipe_read_q);
     ++smem_pipe_read_q;
 
     // normalize
-    Tensor scale_o = attention_updater.finalize(tSrS); // warp reduce row sum
+    Tensor scale_o = attention_updater.finalize(tSrS);  // warp reduce row sum
     if (chunk_num_this_seq == 1) {
       // norm
       cute::copy(scale_o, tScalesScale);
 
-      cutlass::arch::NamedBarrier::arrive(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG2));
+      cutlass::arch::NamedBarrier::arrive(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWarpSchedulerWG2));
       attention_updater.rescale_o(tOrO, scale_o);
     }
 
     // WG1 write m,d back to gmem
-    if (chunk_num_this_seq > 1 && thread_idx % 4 == 0) { // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
+    if (chunk_num_this_seq > 1 &&
+        thread_idx % 4 ==
+            0) {  // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
       const int warp_idx = thread_idx / 32;
 #pragma unroll
       for (int w_i = 0; w_i < 2; ++w_i) {
-        const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i;
+        const int token_group_idx =
+            warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i;
         const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
 
         if (token_idx < qo_len) {
           const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
-          const int bid_offset = mainloop_params.max_draft_token_num * Ktraits::GROUP_SIZE;
-          const int write_idx = bid * bid_offset + token_idx * Ktraits::GROUP_SIZE + head_idx;
+          const int bid_offset =
+              mainloop_params.max_draft_token_num * Ktraits::GROUP_SIZE;
+          const int write_idx =
+              bid * bid_offset + token_idx * Ktraits::GROUP_SIZE + head_idx;
           mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
           mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
         }
@@ -201,7 +243,9 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
     for (; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
       // wait kv
       consumer_wait(pipeline_kv, smem_pipe_read_kv);
-      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+      cutlass::arch::NamedBarrier::sync(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
 
       // A: tPsP
       cute::copy(tScalesScale, scale_o);
@@ -209,21 +253,25 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
       // rescale
       attention_updater.rescale_o(tOrO, scale_o);
       if (smem_pipe_read_kv.index() == 0) {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2,
-                                            tOrV1(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/0>(
+            tiled_mma_pv_ss, tOrP_CS2, tOrV1(_, _, _, _0{}), tOrO);
       } else {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2,
-                                            tOrV2(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/0>(
+            tiled_mma_pv_ss, tOrP_CS2, tOrV2(_, _, _, _0{}), tOrO);
       }
 
       pipeline_kv.consumer_release(smem_pipe_read_kv);
       ++smem_pipe_read_kv;
       // sync WG1 WG2
-      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2Sync));
+      cutlass::arch::NamedBarrier::sync(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWG1WG2Sync));
     }
     if (chunk_num_this_seq == 1) {
       // norm
-      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG2));
+      cutlass::arch::NamedBarrier::sync(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWarpSchedulerWG2));
       cute::copy(tScalesScale, scale_o);
       attention_updater.rescale_o(tOrO, scale_o);
     }
@@ -231,8 +279,16 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
   return;
 }
 
-template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQ,
-          typename PipelineState, typename PipelineStateQ, typename SharedStorage, typename FrgTensorO, typename AttentionUpdater>
+template <typename Ktraits,
+          bool CAUSAL,
+          typename Params,
+          typename MainloopPipeline,
+          typename MainloopPipelineQ,
+          typename PipelineState,
+          typename PipelineStateQ,
+          typename SharedStorage,
+          typename FrgTensorO,
+          typename AttentionUpdater>
 CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
                                        MainloopPipelineQ pipeline_q,
                                        PipelineStateQ& smem_pipe_read_q,
@@ -248,7 +304,7 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
                                        SharedStorage& shared_storage) {
   using DTypeQ = typename Ktraits::DTypeQ;
   using DTypeKV = typename Ktraits::DTypeKV;
-  using DTypeMD = typename Ktraits::DTypeO; // !!! bf16
+  using DTypeMD = typename Ktraits::DTypeO;  // !!! bf16
   using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
   using IdType = typename Ktraits::IdType;
   using TileShape_QKD = typename Ktraits::TileShape_QKD;
@@ -263,29 +319,44 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
   using SmemLayoutVtOneStage = typename Ktraits::SmemLayoutVtOneStage;
   static_assert(is_rmem<FrgTensorO>::value, "O tensor must be rmem resident.");
 
-  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size_device[0]);
+  const int chunk_num_this_seq =
+      cute::ceil_div(kv_len, mainloop_params.chunk_size_device[0]);
 
   static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
   static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
 
-  Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
-  Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
-  Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutVtOneStage{});
-  Tensor sVt_s2 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
-  Tensor sVt_s3 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 2 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
-  Tensor sVt_s4 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 3 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
-  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutP{});
-  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr), mainloop_params.layout_MD)(tile_idx, _);
-  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr), mainloop_params.layout_MD)(tile_idx, _);
+  Tensor sQ =
+      make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+  Tensor sK =
+      make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+  Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()),
+                              SmemLayoutVtOneStage{});
+  Tensor sVt_s2 = make_tensor(
+      make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE),
+      SmemLayoutVtOneStage{});
+  Tensor sVt_s3 = make_tensor(
+      make_smem_ptr(shared_storage.smem_kv.data() + 2 * Ktraits::NUM_PER_STAGE),
+      SmemLayoutVtOneStage{});
+  Tensor sVt_s4 = make_tensor(
+      make_smem_ptr(shared_storage.smem_kv.data() + 3 * Ktraits::NUM_PER_STAGE),
+      SmemLayoutVtOneStage{});
+  Tensor sPSS =
+      make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutP{});
+  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr),
+                          mainloop_params.layout_MD)(tile_idx, _);
+  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr),
+                          mainloop_params.layout_MD)(tile_idx, _);
 
-  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
+  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()),
+                               SmemLayoutRow{});
 
   typename Ktraits::TiledMmaQK tiled_mma_qk;
   auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx);
   auto smem_tiled_copy_P = make_tiled_copy_C(SmemCopyAtom{}, tiled_mma_qk);
   auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(thread_idx);
   Tensor tPsP = smem_thr_copy_P.partition_D(sPSS);
-  Tensor tScalesScale = s_scale(_, thread_idx % cutlass::NumThreadsPerWarpGroup, _);
+  Tensor tScalesScale =
+      s_scale(_, thread_idx % cutlass::NumThreadsPerWarpGroup, _);
 
   typename Ktraits::TiledMmaPVSS tiled_mma_pv_ss;
   auto threadMmaPVSS = tiled_mma_pv_ss.get_thread_slice(thread_idx);
@@ -297,7 +368,11 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
 
   const int start_len = tile_idx * mainloop_params.chunk_size_device[0];
   const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
-  const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size_device[0], kv_len), BLOCK_SHAPE_KV) - 1;
+  const int end_tile_idx =
+      cute::ceil_div(
+          min(start_len + mainloop_params.chunk_size_device[0], kv_len),
+          BLOCK_SHAPE_KV) -
+      1;
   int kv_tile_idx = end_tile_idx;
 
   auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
@@ -310,15 +385,18 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
     // consumer 0, compute qk
     Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ);
     Tensor tSrK = threadMmaQK.partition_fragment_B(sK);
-    auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; };
+    auto col_limit_right = [&](int qo_idx) {
+      return qo_idx + 1 + kv_len - qo_len;
+    };
     // wait q
     consumer_wait(pipeline_q, smem_pipe_read_q);
-    Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+    Tensor tSrS =
+        partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
     // wait k
     consumer_wait(pipeline_kv, smem_pipe_read_kv);
     // first qk gemm
-    gemm</*init=*/true, /*wg_wait=*/0>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
-                                       tSrS);
+    gemm</*init=*/true, /*wg_wait=*/0>(
+        tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()), tSrS);
     // mask
     {
       Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
@@ -342,38 +420,56 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
     Tensor scale_o = attention_updater.update</*init=*/true>(tSrS);
     Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS));
     // gather qk gemm res
-    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_read_kv.index() % 2));
+    cute::copy(
+        smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_read_kv.index() % 2));
     cute::copy(scale_o, tScalesScale(_, smem_pipe_read_kv.index() % 2));
     // r2s fence wgmma
     cutlass::arch::fence_view_async_shared();
-    cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+    cutlass::arch::NamedBarrier::sync(
+        Ktraits::NUM_MMA_THREADS,
+        static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
 
-    constexpr int n_masking_steps = CAUSAL ? cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) : 0;
+    constexpr int n_masking_steps =
+        CAUSAL ? cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) : 0;
     --kv_tile_idx;
-    for (int masking_step = n_masking_steps; kv_tile_idx >= start_tile_idx; --masking_step, --kv_tile_idx) {
-      Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+    for (int masking_step = n_masking_steps; kv_tile_idx >= start_tile_idx;
+         --masking_step, --kv_tile_idx) {
+      Tensor tSrS =
+          partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
       PipelineState smem_pipe_read_kv_cur = smem_pipe_read_kv;
       ++smem_pipe_read_kv;
       // wait next kv
       consumer_wait(pipeline_kv, smem_pipe_read_kv);
 
       // gemm next qk
-      gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
-                                          tSrS);
+      gemm</*init=*/true, /*wg_wait=*/-1>(
+          tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()), tSrS);
       attention_updater.rescale_o(tOrO);
       // last pv gemm
       if (smem_pipe_read_kv_cur.index() == 0) {
-        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
-                                             tOrV1(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/-1>(
+            tiled_mma_pv_ss,
+            tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
+            tOrV1(_, _, _, _0{}),
+            tOrO);
       } else if (smem_pipe_read_kv_cur.index() == 1) {
-        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
-                                             tOrV2(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/-1>(
+            tiled_mma_pv_ss,
+            tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
+            tOrV2(_, _, _, _0{}),
+            tOrO);
       } else if (smem_pipe_read_kv_cur.index() == 2) {
-        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
-                                             tOrV3(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/-1>(
+            tiled_mma_pv_ss,
+            tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
+            tOrV3(_, _, _, _0{}),
+            tOrO);
       } else {
-        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
-                                             tOrV4(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/-1>(
+            tiled_mma_pv_ss,
+            tOrP_CS2(_, _, _, smem_pipe_read_kv_cur.index() % 2),
+            tOrV4(_, _, _, _0{}),
+            tOrO);
       }
       // wait cur qk gemm
       warpgroup_wait<1>();
@@ -401,12 +497,16 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
       Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS));
 
       // gather qk gemm res
-      cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_read_kv.index() % 2));
+      cute::copy(smem_tiled_copy_P,
+                 tPrP,
+                 tPsP(_, _, _, smem_pipe_read_kv.index() % 2));
       cute::copy(scale_o, tScalesScale(_, smem_pipe_read_kv.index() % 2));
       // r2s fence wgmma
       cutlass::arch::fence_view_async_shared();
       // make sure tSrS r2s done
-      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+      cutlass::arch::NamedBarrier::sync(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
       // wait last pv gemm
       warpgroup_wait<0>();
       // release last kv
@@ -418,17 +518,29 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
     // compute last pv
     attention_updater.rescale_o(tOrO);
     if (smem_pipe_read_kv.index() == 0) {
-      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
-                                           tOrV1(_, _, _, _0{}), tOrO);
+      gemm</*init=*/false, /*wg_wait=*/-1>(
+          tiled_mma_pv_ss,
+          tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
+          tOrV1(_, _, _, _0{}),
+          tOrO);
     } else if (smem_pipe_read_kv.index() == 1) {
-      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
-                                           tOrV2(_, _, _, _0{}), tOrO);
+      gemm</*init=*/false, /*wg_wait=*/-1>(
+          tiled_mma_pv_ss,
+          tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
+          tOrV2(_, _, _, _0{}),
+          tOrO);
     } else if (smem_pipe_read_kv.index() == 2) {
-      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
-                                           tOrV3(_, _, _, _0{}), tOrO);
+      gemm</*init=*/false, /*wg_wait=*/-1>(
+          tiled_mma_pv_ss,
+          tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
+          tOrV3(_, _, _, _0{}),
+          tOrO);
     } else {
-      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
-                                           tOrV4(_, _, _, _0{}), tOrO);
+      gemm</*init=*/false, /*wg_wait=*/-1>(
+          tiled_mma_pv_ss,
+          tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
+          tOrV4(_, _, _, _0{}),
+          tOrO);
     }
     scale_o = attention_updater.finalize(tSrS);
     warpgroup_wait<0>();
@@ -439,21 +551,28 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
       // norm
       cute::copy(scale_o, tScalesScale(_, smem_pipe_read_kv.index() % 2));
 
-      cutlass::arch::NamedBarrier::arrive(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2LastSync));
+      cutlass::arch::NamedBarrier::arrive(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWG1WG2LastSync));
       attention_updater.rescale_o(tOrO);
     }
     // WG1 write m,d back to gmem
-    if (chunk_num_this_seq > 1 && thread_idx % 4 == 0) { // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
+    if (chunk_num_this_seq > 1 &&
+        thread_idx % 4 ==
+            0) {  // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
       const int warp_idx = thread_idx / 32;
 #pragma unroll
       for (int w_i = 0; w_i < 2; ++w_i) {
-        const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i;
+        const int token_group_idx =
+            warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i;
         const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
 
         if (token_idx < qo_len) {
           const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
-          const int bid_offset = mainloop_params.max_draft_token_num * Ktraits::GROUP_SIZE;
-          const int write_idx = bid * bid_offset + token_idx * Ktraits::GROUP_SIZE + head_idx;
+          const int bid_offset =
+              mainloop_params.max_draft_token_num * Ktraits::GROUP_SIZE;
+          const int write_idx =
+              bid * bid_offset + token_idx * Ktraits::GROUP_SIZE + head_idx;
           mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
           mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
         }
@@ -464,30 +583,46 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
     Tensor scale_o = make_tensor<DTypeQKAccum>(Shape<_2>{});
     for (; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
       consumer_wait(pipeline_kv, smem_pipe_read_kv);
-      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+      cutlass::arch::NamedBarrier::sync(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
       // A: tPsP
       cute::copy(tScalesScale(_, smem_pipe_read_kv.index() % 2), scale_o);
       // rescale
       attention_updater.rescale_o(tOrO, scale_o);
       if (smem_pipe_read_kv.index() == 0) {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
-                                            tOrV1(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/0>(
+            tiled_mma_pv_ss,
+            tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
+            tOrV1(_, _, _, _0{}),
+            tOrO);
       } else if (smem_pipe_read_kv.index() == 1) {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
-                                            tOrV2(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/0>(
+            tiled_mma_pv_ss,
+            tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
+            tOrV2(_, _, _, _0{}),
+            tOrO);
       } else if (smem_pipe_read_kv.index() == 2) {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
-                                            tOrV3(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/0>(
+            tiled_mma_pv_ss,
+            tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
+            tOrV3(_, _, _, _0{}),
+            tOrO);
       } else {
-        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
-                                            tOrV4(_, _, _, _0{}), tOrO);
+        gemm</*init=*/false, /*wg_wait=*/0>(
+            tiled_mma_pv_ss,
+            tOrP_CS2(_, _, _, smem_pipe_read_kv.index() % 2),
+            tOrV4(_, _, _, _0{}),
+            tOrO);
       }
       pipeline_kv.consumer_release(smem_pipe_read_kv);
       ++smem_pipe_read_kv;
     }
     if (chunk_num_this_seq == 1) {
       // norm
-      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2LastSync));
+      cutlass::arch::NamedBarrier::sync(
+          Ktraits::NUM_MMA_THREADS,
+          static_cast<int>(NamedBarriers::kWG1WG2LastSync));
       cute::copy(tScalesScale(_, smem_pipe_read_kv.index() % 2), scale_o);
       attention_updater.rescale_o(tOrO, scale_o);
     }
diff --git a/custom_ops/gpu_ops/mla_attn/named_barrier.cuh b/custom_ops/gpu_ops/mla_attn/named_barrier.cuh
index bf2f8bf219..41b64b7a07 100644
--- a/custom_ops/gpu_ops/mla_attn/named_barrier.cuh
+++ b/custom_ops/gpu_ops/mla_attn/named_barrier.cuh
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 /*
- * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
- * Dao. Licensed under the BSD 3-Clause.
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar,
+ * Pradeep Ramani, Tri Dao. Licensed under the BSD 3-Clause.
  *
  * Modified by the FlashInfer team.
  */
diff --git a/custom_ops/gpu_ops/mla_attn/utils.cuh b/custom_ops/gpu_ops/mla_attn/utils.cuh
index 2ea1d7e110..15ca2bf423 100644
--- a/custom_ops/gpu_ops/mla_attn/utils.cuh
+++ b/custom_ops/gpu_ops/mla_attn/utils.cuh
@@ -68,41 +68,53 @@ CUTLASS_HOST_DEVICE auto flatten_1(TensorT tensor) {
   return cute::group_modes<1, rank(tensor_flatten)>(tensor_flatten);
 }
 
-CUTLASS_HOST_DEVICE auto get_gmem_layout(int nnz, int num_heads, int head_dim, int64_t n_stride,
-                                         int64_t h_stride) {
+CUTLASS_HOST_DEVICE auto get_gmem_layout(
+    int nnz, int num_heads, int head_dim, int64_t n_stride, int64_t h_stride) {
   return make_layout(make_shape(nnz, head_dim, num_heads),
                      make_stride(n_stride, cute::_1{}, h_stride));
 }
 
 CUTLASS_HOST_DEVICE auto get_lse_gmem_layout(int nnz, int num_heads) {
-  return make_layout(make_shape(num_heads, nnz), make_stride(cute::_1{}, int64_t(num_heads)));
+  return make_layout(make_shape(num_heads, nnz),
+                     make_stride(cute::_1{}, int64_t(num_heads)));
 }
 
 template <typename MTensor, typename Shape>
-CUTLASS_DEVICE auto get_local_tile_tensor(const MTensor& m_tensor, const Shape& tile_shape,
-                                          int head_idx, int offset, int seq_len) {
-  auto g_offset = local_tile(m_tensor(_, _, head_idx), cute::make_shape(1, get<1>(tile_shape)),
+CUTLASS_DEVICE auto get_local_tile_tensor(const MTensor& m_tensor,
+                                          const Shape& tile_shape,
+                                          int head_idx,
+                                          int offset,
+                                          int seq_len) {
+  auto g_offset = local_tile(m_tensor(_, _, head_idx),
+                             cute::make_shape(1, get<1>(tile_shape)),
                              make_coord(offset, _0{}));
   auto g_sequence =
       make_tensor(g_offset.data(),
-                  make_layout(cute::make_shape(seq_len, get<1>(tile_shape)), g_offset.stride()));
+                  make_layout(cute::make_shape(seq_len, get<1>(tile_shape)),
+                              g_offset.stride()));
   auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_, _0{}));
   return g_tensor;
 }
 
 template <typename MTensor, typename Shape>
-CUTLASS_DEVICE auto get_lse_local_tile_tensor(const MTensor& m_tensor, const Shape& tile_shape,
-                                              int head_idx, int offset, int seq_len) {
-  auto g_offset = local_tile(m_tensor(head_idx, _), cute::make_shape(_1{}), make_coord(offset));
+CUTLASS_DEVICE auto get_lse_local_tile_tensor(const MTensor& m_tensor,
+                                              const Shape& tile_shape,
+                                              int head_idx,
+                                              int offset,
+                                              int seq_len) {
+  auto g_offset = local_tile(
+      m_tensor(head_idx, _), cute::make_shape(_1{}), make_coord(offset));
 
-  auto g_sequence = make_tensor(g_offset.data(), make_layout(cute::make_shape(seq_len),
-                                                             cute::make_shape(shape<0>(m_tensor))));
+  auto g_sequence =
+      make_tensor(g_offset.data(),
+                  make_layout(cute::make_shape(seq_len),
+                              cute::make_shape(shape<0>(m_tensor))));
   auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_));
   return g_tensor;
 }
 
-// For SM90, convert acc_layout from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V,
-// MMA_N))
+// For SM90, convert acc_layout from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2,
+// MMA_M), ncol=(2, V, MMA_N))
 template <typename Layout>
 __forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
   static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
@@ -113,8 +125,8 @@ __forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
                      make_layout(get<0, 0>(l), get<0, 2>(l), get<2>(l)));
 };
 
-// For SM90, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2, 2), MMA_M, (N / 16,
-// MMA_N))
+// For SM90, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2,
+// 2), MMA_M, (N / 16, MMA_N))
 template <typename MMA_traits, typename Layout>
 __forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
   using X = Underscore;
@@ -122,27 +134,39 @@ __forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
   static_assert(decltype(size<0, 1>(acc_layout))::value == 2);
   static_assert(decltype(rank(acc_layout))::value == 3);
   static_assert(decltype(rank(get<0>(acc_layout)))::value == 3);
-  auto l = logical_divide(get<0>(acc_layout), Shape<X, X, _2>{});  // (2, 2, (2, N / 16)))
-  return make_layout(make_layout(get<0>(l), get<1>(l), get<2, 0>(l)), get<1>(acc_layout),
+  auto l = logical_divide(get<0>(acc_layout),
+                          Shape<X, X, _2>{});  // (2, 2, (2, N / 16)))
+  return make_layout(make_layout(get<0>(l), get<1>(l), get<2, 0>(l)),
+                     get<1>(acc_layout),
                      make_layout(get<2, 1>(l), get<2>(acc_layout)));
 };
 
 template <typename To_type, typename Engine, typename Layout>
-__forceinline__ __device__ auto convert_type(Tensor<Engine, Layout> const& tensor) {
+__forceinline__ __device__ auto convert_type(
+    Tensor<Engine, Layout> const& tensor) {
   using From_type = typename Engine::value_type;
   constexpr int numel = decltype(size(tensor))::value;
   cutlass::NumericArrayConverter<To_type, From_type, numel> convert_op;
-  auto frag = convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel>*>(tensor.data()));
+  auto frag =
+      convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel>*>(
+          tensor.data()));
   return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());
 }
 
-template <bool init = false, int wg_wait = 0, typename TensorA, typename TensorB, typename TensorC,
+template <bool init = false,
+          int wg_wait = 0,
+          typename TensorA,
+          typename TensorB,
+          typename TensorC,
           typename TiledMma>
-__forceinline__ __device__ void gemm(TiledMma& tiled_mma, TensorA const& tCrA, TensorB const& tCrB,
+__forceinline__ __device__ void gemm(TiledMma& tiled_mma,
+                                     TensorA const& tCrA,
+                                     TensorB const& tCrB,
                                      TensorC& tCrC) {
-  constexpr bool Is_RS =
-      !cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value;
-  // Need to cast away const on tCrA since warpgroup_fence_operand doesn't take const
+  constexpr bool Is_RS = !cute::is_base_of<cute::GMMA::DescriptorIterator,
+                                           typename TiledMma::FrgTypeA>::value;
+  // Need to cast away const on tCrA since warpgroup_fence_operand doesn't take
+  // const
   if constexpr (Is_RS) {
     warpgroup_fence_operand(const_cast<TensorA&>(tCrA));
   }
@@ -225,13 +249,15 @@ struct prefill_softmax_state_t {
     }
   }
 
-  __device__ __forceinline__ void merge(const AlignedVector<T, vec_size>& other_o,
-                                        const float other_m,
-                                        const float other_d) {
+  __device__ __forceinline__ void merge(
+      const AlignedVector<T, vec_size>& other_o,
+      const float other_m,
+      const float other_d) {
     float m_prev = m, d_prev = d;
     m = max(m_prev, other_m);
     const float scale1 = __expf(m_prev - m), scale2 = __expf(other_m - m);
-    const T scale1_T = static_cast<T>(scale1), scale2_T = static_cast<T>(scale2);
+    const T scale1_T = static_cast<T>(scale1),
+            scale2_T = static_cast<T>(scale2);
     d = d_prev * scale1 + other_d * scale2;
 #pragma unroll
     for (size_t i = 0; i < vec_size; ++i) {
@@ -249,27 +275,32 @@ struct prefill_softmax_state_t {
 };
 
 template <typename T, int vec_size, uint32_t bdy, uint32_t HEAD_DIM>
-__global__ void merge_multi_chunks_kernel(const T * __restrict__ multi_out, // [max_num_chunks, bsz, max_draft_token, num_heads, head_dim]
-                                          const float * __restrict__ multi_m, // [max_num_chunks, bsz, max_draft_token, num_heads]
-                                          const float * __restrict__ multi_d, // [max_num_chunks, bsz, max_draft_token, num_heads]
-                                          const int * __restrict__ seq_lens_this_time,
-                                          const int * __restrict__ seq_lens_decoder,
-                                          const int *__restrict__ cu_seqlens_q,
-                                          const int * __restrict__ batch_id_per_token,
-                                          const int * __restrict__ chunk_size_device,
-                                          T * __restrict__ out, // [token_num, num_heads, head_dim]
-                                          const int num_heads,
-                                          const int head_dim,
-                                          const int token_num,
-                                          const int bsz,
-                                          const int max_draft_token_num) {
+__global__ void merge_multi_chunks_kernel(
+    const T* __restrict__ multi_out,  // [max_num_chunks, bsz, max_draft_token,
+                                      // num_heads, head_dim]
+    const float* __restrict__ multi_m,  // [max_num_chunks, bsz,
+                                        // max_draft_token, num_heads]
+    const float* __restrict__ multi_d,  // [max_num_chunks, bsz,
+                                        // max_draft_token, num_heads]
+    const int* __restrict__ seq_lens_this_time,
+    const int* __restrict__ seq_lens_decoder,
+    const int* __restrict__ cu_seqlens_q,
+    const int* __restrict__ batch_id_per_token,
+    const int* __restrict__ chunk_size_device,
+    T* __restrict__ out,  // [token_num, num_heads, head_dim]
+    const int num_heads,
+    const int head_dim,
+    const int token_num,
+    const int bsz,
+    const int max_draft_token_num) {
   const int vid = threadIdx.x, ty = threadIdx.y;
   const int hid = blockIdx.y;
   __shared__ T smem[bdy * HEAD_DIM];
   __shared__ float md_smem[bdy * 2];
   for (int qid = blockIdx.x; qid < token_num; qid += gridDim.x) {
     const uint32_t bid = batch_id_per_token[qid];
-    // NOTE : (changwenbin) Batch_id_per_token is initialized to [:]=-1, Marking meaningless batch IDs.
+    // NOTE : (changwenbin) Batch_id_per_token is initialized to [:]=-1, Marking
+    // meaningless batch IDs.
     if (bid == -1) continue;
     const int seq_len_q = seq_lens_this_time[bid];
     if (seq_len_q == 0) continue;
@@ -277,7 +308,8 @@ __global__ void merge_multi_chunks_kernel(const T * __restrict__ multi_out, // [
     int seq_len_kv = seq_lens_decoder[bid];
     if (seq_len_kv == 0) continue;
     seq_len_kv += seq_len_q;
-    const int num_chunks_this_seq = cute::ceil_div(seq_len_kv, chunk_size_device[0]);
+    const int num_chunks_this_seq =
+        cute::ceil_div(seq_len_kv, chunk_size_device[0]);
     if (num_chunks_this_seq <= 1) {
       // not need merge
       continue;
@@ -307,16 +339,23 @@ __global__ void merge_multi_chunks_kernel(const T * __restrict__ multi_out, // [
 
     for (int i = ty; i < num_chunks_this_seq; i += bdy) {
       uint32_t offset;
-      offset = ((i * bsz + bid) * max_draft_token_num + local_seq_id) * num_heads + hid;
+      offset =
+          ((i * bsz + bid) * max_draft_token_num + local_seq_id) * num_heads +
+          hid;
       float m_prev = m;
       float d_prev = d;
       const float m_now = multi_m[offset];
       const float d_now = multi_d[offset];
       m = max(m_prev, m_now);
-      offset = (((i * bsz + bid) * max_draft_token_num + local_seq_id) * num_heads + hid) * head_dim + vid * vec_size;
+      offset =
+          (((i * bsz + bid) * max_draft_token_num + local_seq_id) * num_heads +
+           hid) *
+              head_dim +
+          vid * vec_size;
       Load<T, vec_size>(&multi_out[offset], &load_vec);
       const float scale1 = __expf(m_prev - m), scale2 = __expf(m_now - m);
-      const T scale1_T = static_cast<T>(scale1), scale2_T = static_cast<T>(scale2);
+      const T scale1_T = static_cast<T>(scale1),
+              scale2_T = static_cast<T>(scale2);
       d = d * scale1 + d_now * scale2;
 #pragma unroll
       for (int j = 0; j < vec_size; j++) {
@@ -339,7 +378,8 @@ __global__ void merge_multi_chunks_kernel(const T * __restrict__ multi_out, // [
         st.merge(load_vec, m_tmp, d_tmp);
       }
       st.normalize();
-      Store<T, vec_size>(st.o, &out[(qid * num_heads + hid) * head_dim + vid * vec_size]);
+      Store<T, vec_size>(
+          st.o, &out[(qid * num_heads + hid) * head_dim + vid * vec_size]);
     }
     __syncthreads();
   }
diff --git a/custom_ops/gpu_ops/moba_attn/moba_attn.cu b/custom_ops/gpu_ops/moba_attn/moba_attn.cu
index bf0bd38150..c269205bf1 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_attn.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_attn.cu
@@ -20,314 +20,291 @@
 #endif
 
 std::vector<paddle::Tensor> MobaAttention(
-        const paddle::Tensor& qkv,
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& k_input,
-        const paddle::Tensor& v_input,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::Tensor& cu_seq_q_pack,
-        const paddle::Tensor& q_pack_tokens,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& key_cache,
-        const paddle::Tensor& value_cache,
-        const paddle::Tensor& block_tables,
-        const paddle::Tensor& rope_sin_cos,
-        const paddle::Tensor& k_block_means,
-        const paddle::optional<paddle::Tensor>& attn_gate_weight,
-        const paddle::optional<paddle::Tensor>& qkv_bias,
-        const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_zero_points,
-        const paddle::optional<paddle::Tensor>& cache_v_zero_points,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_seq_len,
-        const int max_enc_len_this_time,
-        const int max_dec_len_this_time,
-        const int moba_encoder_top_k_left,
-        const int moba_encoder_top_k_right,
-        const int moba_use_encoder_seq_limit,
-        const int moba_decoder_top_k_left,
-        const int moba_decoder_top_k_right,
-        const int moba_use_decoder_seq_limit,
-        const bool moba_use_mlp,
-        const std::string &cache_quant_type_str) {
+    const paddle::Tensor& qkv,
+    const paddle::Tensor& q_input,
+    const paddle::Tensor& k_input,
+    const paddle::Tensor& v_input,
+    const paddle::Tensor& cu_seq_q,
+    const paddle::Tensor& cu_seq_k,
+    const paddle::Tensor& cu_seq_q_pack,
+    const paddle::Tensor& q_pack_tokens,
+    const paddle::Tensor& seq_len_encoder,
+    const paddle::Tensor& seq_len_decoder,
+    const paddle::Tensor& key_cache,
+    const paddle::Tensor& value_cache,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& rope_sin_cos,
+    const paddle::Tensor& k_block_means,
+    const paddle::optional<paddle::Tensor>& attn_gate_weight,
+    const paddle::optional<paddle::Tensor>& qkv_bias,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zero_points,
+    const paddle::optional<paddle::Tensor>& cache_v_zero_points,
+    const int head_num,
+    const int kv_head_num,
+    const int head_dim,
+    const int max_seq_len,
+    const int max_enc_len_this_time,
+    const int max_dec_len_this_time,
+    const int moba_encoder_top_k_left,
+    const int moba_encoder_top_k_right,
+    const int moba_use_encoder_seq_limit,
+    const int moba_decoder_top_k_left,
+    const int moba_decoder_top_k_right,
+    const int moba_use_decoder_seq_limit,
+    const bool moba_use_mlp,
+    const std::string& cache_quant_type_str) {
+  paddle::Tensor out = paddle::empty(
+      {qkv.dims()[0], head_num * head_dim}, qkv.dtype(), qkv.place());
+  if (max_dec_len_this_time > 0) {
+    MobaDecoderAttnWriteCacheKv(qkv,
+                                q_input,
+                                cu_seq_q,
+                                cu_seq_k,
+                                seq_len_encoder,
+                                seq_len_decoder,
+                                key_cache,
+                                value_cache,
+                                block_tables,
+                                rope_sin_cos,
+                                k_block_means,
+                                qkv_bias,
+                                cache_k_quant_scale,
+                                cache_v_quant_scale,
+                                cache_k_dequant_scale,
+                                cache_v_dequant_scale,
+                                cache_k_zero_points,
+                                cache_v_zero_points,
+                                head_num,
+                                kv_head_num,
+                                head_dim,
+                                max_seq_len,
+                                cache_quant_type_str);
 
-    paddle::Tensor out = paddle::empty({qkv.dims()[0], head_num * head_dim}, qkv.dtype(), qkv.place());
-    if (max_dec_len_this_time > 0) {
-        MobaDecoderAttnWriteCacheKv(
-            qkv,
-            q_input,
-            cu_seq_q,
-            cu_seq_k,
-            seq_len_encoder,
-            seq_len_decoder,
-            key_cache,
-            value_cache,
-            block_tables,
-            rope_sin_cos,
-            k_block_means,
-            qkv_bias,
-            cache_k_quant_scale,
-            cache_v_quant_scale,
-            cache_k_dequant_scale,
-            cache_v_dequant_scale,
-            cache_k_zero_points,
-            cache_v_zero_points,
-            head_num,
-            kv_head_num,
-            head_dim,
-            max_seq_len,
-            cache_quant_type_str);
+    auto qk_gate_weight = MobaQKGemm(q_input,
+                                     k_block_means,
+                                     seq_len_encoder,
+                                     seq_len_decoder,
+                                     cu_seq_q,
+                                     cu_seq_k,
+                                     max_dec_len_this_time,
+                                     max_dec_len_this_time,
+                                     head_num,
+                                     kv_head_num,
+                                     true,
+                                     moba_use_decoder_seq_limit)[0];
 
-        auto qk_gate_weight = MobaQKGemm(
-            q_input,
-            k_block_means,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            cu_seq_k,
-            max_dec_len_this_time,
-            max_dec_len_this_time,
-            head_num,
-            kv_head_num,
-            true,
-            moba_use_decoder_seq_limit
-        )[0];
+    auto qk_gate_topk_idx = QkSortDecoder(qk_gate_weight,
+                                          seq_len_encoder,
+                                          seq_len_decoder,
+                                          head_num,
+                                          kv_head_num,
+                                          moba_decoder_top_k_left,
+                                          moba_decoder_top_k_right,
+                                          moba_use_decoder_seq_limit)[0];
 
-        auto qk_gate_topk_idx = QkSortDecoder(
-            qk_gate_weight,
-            seq_len_encoder,
-            seq_len_decoder,
-            head_num,
-            kv_head_num,
-            moba_decoder_top_k_left,
-            moba_decoder_top_k_right,
-            moba_use_decoder_seq_limit
-        )[0];
+    MobaDecoderAttn(q_input,
+                    seq_len_encoder,
+                    seq_len_decoder,
+                    cu_seq_q,
+                    key_cache,
+                    value_cache,
+                    block_tables,
+                    k_block_means,
+                    out,
+                    qk_gate_topk_idx,
+                    cache_k_quant_scale,
+                    cache_v_quant_scale,
+                    cache_k_dequant_scale,
+                    cache_v_dequant_scale,
+                    cache_k_zero_points,
+                    cache_v_zero_points,
+                    head_num,
+                    kv_head_num,
+                    head_dim,
+                    max_seq_len,
+                    moba_use_decoder_seq_limit,
+                    max_dec_len_this_time,
+                    max_dec_len_this_time,
+                    cache_quant_type_str);
+  }
 
-        MobaDecoderAttn(
-            q_input,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            key_cache,
-            value_cache,
-            block_tables,
-            k_block_means,
-            out,
-            qk_gate_topk_idx,
-            cache_k_quant_scale,
-            cache_v_quant_scale,
-            cache_k_dequant_scale,
-            cache_v_dequant_scale,
-            cache_k_zero_points,
-            cache_v_zero_points,
-            head_num,
-            kv_head_num,
-            head_dim,
-            max_seq_len,
-            moba_use_decoder_seq_limit,
-            max_dec_len_this_time,
-            max_dec_len_this_time,
-            cache_quant_type_str
-        );
+  if (max_enc_len_this_time > 0) {
+    FusedBlockMeanAndRope(qkv,
+                          k_block_means,
+                          q_input,
+                          k_input,
+                          v_input,
+                          rope_sin_cos,
+                          seq_len_encoder,
+                          seq_len_decoder,
+                          cu_seq_q,
+                          cu_seq_k,
+                          qkv_bias,
+                          head_num,
+                          kv_head_num,
+                          head_dim,
+                          max_seq_len,
+                          max_enc_len_this_time,
+                          max_enc_len_this_time,
+                          cache_quant_type_str);
+
+    MobaEncoderAttnWriteCacheKv(k_input,
+                                v_input,
+                                cu_seq_k,
+                                seq_len_encoder,
+                                seq_len_decoder,
+                                key_cache,
+                                value_cache,
+                                block_tables,
+                                cache_k_quant_scale,
+                                cache_v_quant_scale,
+                                cache_k_dequant_scale,
+                                cache_v_dequant_scale,
+                                cache_k_zero_points,
+                                cache_v_zero_points,
+                                head_num,
+                                kv_head_num,
+                                head_dim,
+                                max_enc_len_this_time,
+                                cache_quant_type_str);
+
+    GetKVFromCache(k_input,
+                   v_input,
+                   cu_seq_k,
+                   seq_len_encoder,
+                   seq_len_decoder,
+                   key_cache,
+                   value_cache,
+                   block_tables,
+                   cache_k_dequant_scale,
+                   cache_v_dequant_scale,
+                   cache_k_zero_points,
+                   cache_v_zero_points,
+                   head_num,
+                   kv_head_num,
+                   head_dim,
+                   max_seq_len,
+                   max_enc_len_this_time + max_dec_len_this_time,
+                   cache_quant_type_str);
+
+    paddle::Tensor* k_gate_weight = const_cast<paddle::Tensor*>(&k_block_means);
+
+    if (moba_use_mlp && attn_gate_weight) {
+      paddle::Tensor k_gate_mlp = MobaMlpEinsum(k_input,
+                                                attn_gate_weight.get(),
+                                                seq_len_encoder,
+                                                seq_len_decoder,
+                                                cu_seq_k,
+                                                max_seq_len,
+                                                kv_head_num)[0];
+      k_gate_weight = &k_gate_mlp;
     }
 
-    if (max_enc_len_this_time > 0) {
-        FusedBlockMeanAndRope(
-            qkv,
-            k_block_means,
-            q_input,
-            k_input,
-            v_input,
-            rope_sin_cos,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            cu_seq_k,
-            qkv_bias,
-            head_num,
-            kv_head_num,
-            head_dim,
-            max_seq_len,
-            max_enc_len_this_time,
-            max_enc_len_this_time,
-            cache_quant_type_str
-        );
+    auto qk_gate_weight =
+        MobaQKGemm(q_input,
+                   *k_gate_weight,
+                   seq_len_encoder,
+                   seq_len_decoder,
+                   cu_seq_q,
+                   cu_seq_k,
+                   max_enc_len_this_time,
+                   max_enc_len_this_time + max_dec_len_this_time,
+                   head_num,
+                   kv_head_num,
+                   false,
+                   moba_use_encoder_seq_limit)[0];
 
-        MobaEncoderAttnWriteCacheKv(
-            k_input,
-            v_input,
-            cu_seq_k,
-            seq_len_encoder,
-            seq_len_decoder,
-            key_cache,
-            value_cache,
-            block_tables,
-            cache_k_quant_scale,
-            cache_v_quant_scale,
-            cache_k_dequant_scale,
-            cache_v_dequant_scale,
-            cache_k_zero_points,
-            cache_v_zero_points,
-            head_num,
-            kv_head_num,
-            head_dim,
-            max_enc_len_this_time,
-            cache_quant_type_str
-        );
+    auto qk_gate_topk_idx = QkSortEncoder(
+        qk_gate_weight,
+        seq_len_encoder,
+        seq_len_decoder,
+        cu_seq_q,
+        cu_seq_k,
+        cu_seq_q_pack,
+        q_pack_tokens,
+        max_enc_len_this_time,
+        max_enc_len_this_time + max_dec_len_this_time,
+        head_num,
+        kv_head_num,
+        moba_encoder_top_k_left,
+        moba_encoder_top_k_right,
+        moba_use_mlp && !attn_gate_weight ? max_seq_len
+                                          : moba_use_encoder_seq_limit)[0];
 
-        GetKVFromCache(
-            k_input,
-            v_input,
-            cu_seq_k,
-            seq_len_encoder,
-            seq_len_decoder,
-            key_cache,
-            value_cache,
-            block_tables,
-            cache_k_dequant_scale,
-            cache_v_dequant_scale,
-            cache_k_zero_points,
-            cache_v_zero_points,
-            head_num,
-            kv_head_num,
-            head_dim,
-            max_seq_len,
-            max_enc_len_this_time + max_dec_len_this_time,
-            cache_quant_type_str
-        );
+    MobaEncoderAttn(q_input,
+                    k_input,
+                    v_input,
+                    qk_gate_topk_idx,
+                    cu_seq_q,
+                    cu_seq_k,
+                    cu_seq_q_pack,
+                    seq_len_encoder,
+                    seq_len_decoder,
+                    out,
+                    max_enc_len_this_time,
+                    max_enc_len_this_time + max_dec_len_this_time,
+                    head_num,
+                    kv_head_num,
+                    head_dim,
+                    max_seq_len);
+  }
 
-        paddle::Tensor *k_gate_weight = const_cast<paddle::Tensor*>(&k_block_means);
-
-        if (moba_use_mlp && attn_gate_weight) {
-            paddle::Tensor k_gate_mlp = MobaMlpEinsum(
-                k_input,
-                attn_gate_weight.get(),
-                seq_len_encoder,
-                seq_len_decoder,
-                cu_seq_k,
-                max_seq_len,
-                kv_head_num
-            )[0];
-            k_gate_weight = &k_gate_mlp;
-        }
-
-        auto qk_gate_weight = MobaQKGemm(
-            q_input,
-            *k_gate_weight,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            cu_seq_k,
-            max_enc_len_this_time,
-            max_enc_len_this_time + max_dec_len_this_time,
-            head_num,
-            kv_head_num,
-            false,
-            moba_use_encoder_seq_limit
-        )[0];
-
-
-        auto qk_gate_topk_idx = QkSortEncoder(
-            qk_gate_weight,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            cu_seq_k,
-            cu_seq_q_pack,
-            q_pack_tokens,
-            max_enc_len_this_time,
-            max_enc_len_this_time + max_dec_len_this_time,
-            head_num,
-            kv_head_num,
-            moba_encoder_top_k_left,
-            moba_encoder_top_k_right,
-            moba_use_mlp && !attn_gate_weight ? max_seq_len : moba_use_encoder_seq_limit)[0];
-
-        MobaEncoderAttn(
-            q_input,
-            k_input,
-            v_input,
-            qk_gate_topk_idx,
-            cu_seq_q,
-            cu_seq_k,
-            cu_seq_q_pack,
-            seq_len_encoder,
-            seq_len_decoder,
-            out,
-            max_enc_len_this_time,
-            max_enc_len_this_time + max_dec_len_this_time,
-            head_num,
-            kv_head_num,
-            head_dim,
-            max_seq_len
-        );
-    }
-
-    return {out};
+  return {out};
 }
 
-
 PD_BUILD_STATIC_OP(moba_attention)
-    .Inputs({
-        "qkv",
-        "q_input",
-        "k_input",
-        "v_input",
-        "cu_seq_q",
-        "cu_seq_k",
-        "cu_seq_q_pack",
-        "q_pack_tokens",
-        "seq_len_encoder",
-        "seq_len_decoder",
-        "key_cache",
-        "value_cache",
-        "block_tables",
-        "rope_sin_cos",
-        "k_block_means",
-        paddle::Optional("attn_gate_weight"),
-        paddle::Optional("qkv_bias"),
-        paddle::Optional("cache_k_quant_scale"),
-        paddle::Optional("cache_v_quant_scale"),
-        paddle::Optional("cache_k_dequant_scale"),
-        paddle::Optional("cache_v_dequant_scale"),
-        paddle::Optional("cache_k_zero_points"),
-        paddle::Optional("cache_v_zero_points")})
-    .Attrs({
-        "head_num: int",
-        "kv_head_num: int",
-        "head_dim: int",
-        "max_seq_len: int",
-        "max_enc_len_this_time: int",
-        "max_dec_len_this_time: int",
-        "moba_encoder_top_k_left: int",
-        "moba_encoder_top_k_right: int",
-        "moba_use_encoder_seq_limit: int",
-        "moba_decoder_top_k_left: int",
-        "moba_decoder_top_k_right: int",
-        "moba_use_decoder_seq_limit: int",
-        "moba_use_mlp: bool",
-        "cache_quant_type_str: std::string"})
-    .Outputs({
-        "out",
-        "q_input_out",
-        "k_input_out",
-        "v_input_out",
-        "key_cache_out",
-        "value_cache_out",
-        "k_block_means_out"})
-    .SetInplaceMap({{
-        "q_input", "q_input_out"},
-        {"k_input", "k_input_out"},
-        {"v_input", "v_input_out"},
-        {"key_cache", "key_cache_out"},
-        {"value_cache", "value_cache_out"},
-        {"k_block_means", "k_block_means_out"}})
+    .Inputs({"qkv",
+             "q_input",
+             "k_input",
+             "v_input",
+             "cu_seq_q",
+             "cu_seq_k",
+             "cu_seq_q_pack",
+             "q_pack_tokens",
+             "seq_len_encoder",
+             "seq_len_decoder",
+             "key_cache",
+             "value_cache",
+             "block_tables",
+             "rope_sin_cos",
+             "k_block_means",
+             paddle::Optional("attn_gate_weight"),
+             paddle::Optional("qkv_bias"),
+             paddle::Optional("cache_k_quant_scale"),
+             paddle::Optional("cache_v_quant_scale"),
+             paddle::Optional("cache_k_dequant_scale"),
+             paddle::Optional("cache_v_dequant_scale"),
+             paddle::Optional("cache_k_zero_points"),
+             paddle::Optional("cache_v_zero_points")})
+    .Attrs({"head_num: int",
+            "kv_head_num: int",
+            "head_dim: int",
+            "max_seq_len: int",
+            "max_enc_len_this_time: int",
+            "max_dec_len_this_time: int",
+            "moba_encoder_top_k_left: int",
+            "moba_encoder_top_k_right: int",
+            "moba_use_encoder_seq_limit: int",
+            "moba_decoder_top_k_left: int",
+            "moba_decoder_top_k_right: int",
+            "moba_use_decoder_seq_limit: int",
+            "moba_use_mlp: bool",
+            "cache_quant_type_str: std::string"})
+    .Outputs({"out",
+              "q_input_out",
+              "k_input_out",
+              "v_input_out",
+              "key_cache_out",
+              "value_cache_out",
+              "k_block_means_out"})
+    .SetInplaceMap({{"q_input", "q_input_out"},
+                    {"k_input", "k_input_out"},
+                    {"v_input", "v_input_out"},
+                    {"key_cache", "key_cache_out"},
+                    {"value_cache", "value_cache_out"},
+                    {"k_block_means", "k_block_means_out"}})
     .SetKernelFn(PD_KERNEL(MobaAttention));
diff --git a/custom_ops/gpu_ops/moba_attn/moba_attn.h b/custom_ops/gpu_ops/moba_attn/moba_attn.h
index 1b4611022e..a3d81b64c7 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_attn.h
+++ b/custom_ops/gpu_ops/moba_attn/moba_attn.h
@@ -39,7 +39,7 @@ void MobaDecoderAttnWriteCacheKv(
     const int kv_head_num,
     const int head_dim,
     const int max_input_length,
-    const std::string &cache_quant_type_str);
+    const std::string& cache_quant_type_str);
 
 void MobaEncoderAttnWriteCacheKv(
     const paddle::Tensor& k_input,
@@ -60,7 +60,7 @@ void MobaEncoderAttnWriteCacheKv(
     const int kv_head_num,
     const int head_dim,
     const int max_seq_q,
-    const std::string &cache_quant_type_str);
+    const std::string& cache_quant_type_str);
 
 void MobaDecoderAttn(
     const paddle::Tensor& q_input,
@@ -86,28 +86,26 @@ void MobaDecoderAttn(
     const int use_moba_seq_limit,
     const int max_seq_q,
     const int max_seq_k,
-    const std::string &cache_quant_type_str);
+    const std::string& cache_quant_type_str);
 
-
-void FusedBlockMeanAndRope(
-    const paddle::Tensor& qkv_out,
-    const paddle::Tensor& k_block_means,
-    const paddle::Tensor& q_input,
-    const paddle::Tensor& k_input,
-    const paddle::Tensor& v_input,
-    const paddle::Tensor& rotary_embs,
-    const paddle::Tensor& seq_len_encoder,
-    const paddle::Tensor& seq_len_decoder,
-    const paddle::Tensor& cu_seq_q,
-    const paddle::Tensor& cu_seq_k,
-    const paddle::optional<paddle::Tensor>& qkv_bias,
-    const int head_num,
-    const int kv_head_num,
-    const int head_dim,
-    const int max_input_length,
-    const int max_seq_q,
-    const int max_seq_k,
-    const std::string &cache_quant_type_str);
+void FusedBlockMeanAndRope(const paddle::Tensor& qkv_out,
+                           const paddle::Tensor& k_block_means,
+                           const paddle::Tensor& q_input,
+                           const paddle::Tensor& k_input,
+                           const paddle::Tensor& v_input,
+                           const paddle::Tensor& rotary_embs,
+                           const paddle::Tensor& seq_len_encoder,
+                           const paddle::Tensor& seq_len_decoder,
+                           const paddle::Tensor& cu_seq_q,
+                           const paddle::Tensor& cu_seq_k,
+                           const paddle::optional<paddle::Tensor>& qkv_bias,
+                           const int head_num,
+                           const int kv_head_num,
+                           const int head_dim,
+                           const int max_input_length,
+                           const int max_seq_q,
+                           const int max_seq_k,
+                           const std::string& cache_quant_type_str);
 
 std::vector<paddle::Tensor> GetCurCuSeqLenk(
     const paddle::Tensor& seq_lens_encoder,
@@ -115,29 +113,27 @@ std::vector<paddle::Tensor> GetCurCuSeqLenk(
     const paddle::Tensor& seq_lens_this_time,
     const int pack_size);
 
-std::vector<paddle::Tensor> MobaQKGemm(
-    const paddle::Tensor& q_input,
-    const paddle::Tensor& k_block_means,
-    const paddle::Tensor& seq_len_encoder,
-    const paddle::Tensor& seq_len_decoder,
-    const paddle::Tensor& cu_seq_q,
-    const paddle::Tensor& cu_seq_k,
-    const int max_seq_q,
-    const int max_seq_k,
-    const int head_num,
-    const int kv_head_num,
-    const bool is_split_kv,
-    const int use_moba_seq_limit);
+std::vector<paddle::Tensor> MobaQKGemm(const paddle::Tensor& q_input,
+                                       const paddle::Tensor& k_block_means,
+                                       const paddle::Tensor& seq_len_encoder,
+                                       const paddle::Tensor& seq_len_decoder,
+                                       const paddle::Tensor& cu_seq_q,
+                                       const paddle::Tensor& cu_seq_k,
+                                       const int max_seq_q,
+                                       const int max_seq_k,
+                                       const int head_num,
+                                       const int kv_head_num,
+                                       const bool is_split_kv,
+                                       const int use_moba_seq_limit);
 
-std::vector<paddle::Tensor> QkSortDecoder(
-    const paddle::Tensor& qk_gate_weight,
-    const paddle::Tensor& seq_len_encoder,
-    const paddle::Tensor& seq_len_decoder,
-    const int head_num,
-    const int kv_head_num,
-    const int top_k_left,
-    const int top_k_right,
-    const int use_moba_seq_limit);
+std::vector<paddle::Tensor> QkSortDecoder(const paddle::Tensor& qk_gate_weight,
+                                          const paddle::Tensor& seq_len_encoder,
+                                          const paddle::Tensor& seq_len_decoder,
+                                          const int head_num,
+                                          const int kv_head_num,
+                                          const int top_k_left,
+                                          const int top_k_right,
+                                          const int use_moba_seq_limit);
 
 void GetKVFromCache(
     const paddle::Tensor& k_input,
@@ -157,42 +153,39 @@ void GetKVFromCache(
     const int head_dim,
     const int max_input_length,
     const int max_seq_k,
-    const std::string &cache_quant_type_str);
+    const std::string& cache_quant_type_str);
 
+void MobaEncoderAttn(const paddle::Tensor& q_input,
+                     const paddle::Tensor& k_input,
+                     const paddle::Tensor& v_input,
+                     const paddle::Tensor& qk_gate_topk_idx,
+                     const paddle::Tensor& cu_seq_q,
+                     const paddle::Tensor& cu_seq_k,
+                     const paddle::Tensor& cu_seq_q_pack,
+                     const paddle::Tensor& seq_len_encoder,
+                     const paddle::Tensor& seq_len_decoder,
+                     const paddle::Tensor& out,
+                     const int max_seq_q,
+                     const int max_seq_k,
+                     const int head_num,
+                     const int kv_head_num,
+                     const int head_dim,
+                     const int max_input_length);
 
-void MobaEncoderAttn(
-    const paddle::Tensor& q_input,
-    const paddle::Tensor& k_input,
-    const paddle::Tensor& v_input,
-    const paddle::Tensor& qk_gate_topk_idx,
-    const paddle::Tensor& cu_seq_q,
-    const paddle::Tensor& cu_seq_k,
-    const paddle::Tensor& cu_seq_q_pack,
-    const paddle::Tensor& seq_len_encoder,
-    const paddle::Tensor& seq_len_decoder,
-    const paddle::Tensor& out,
-    const int max_seq_q,
-    const int max_seq_k,
-    const int head_num,
-    const int kv_head_num,
-    const int head_dim,
-    const int max_input_length);
-
-std::vector<paddle::Tensor> QkSortEncoder(
-    const paddle::Tensor& qk_gate_weight,
-    const paddle::Tensor& seq_len_encoder,
-    const paddle::Tensor& seq_len_decoder,
-    const paddle::Tensor& cu_seq_q,
-    const paddle::Tensor& cu_seq_k,
-    const paddle::Tensor& cu_seq_q_pack,
-    const paddle::Tensor& q_pack_tokens,
-    const int max_seq_q,
-    const int max_seq_k,
-    const int head_num,
-    const int kv_head_num,
-    const int top_k_left,
-    const int top_k_right,
-    const int use_moba_seq_limit);
+std::vector<paddle::Tensor> QkSortEncoder(const paddle::Tensor& qk_gate_weight,
+                                          const paddle::Tensor& seq_len_encoder,
+                                          const paddle::Tensor& seq_len_decoder,
+                                          const paddle::Tensor& cu_seq_q,
+                                          const paddle::Tensor& cu_seq_k,
+                                          const paddle::Tensor& cu_seq_q_pack,
+                                          const paddle::Tensor& q_pack_tokens,
+                                          const int max_seq_q,
+                                          const int max_seq_k,
+                                          const int head_num,
+                                          const int kv_head_num,
+                                          const int top_k_left,
+                                          const int top_k_right,
+                                          const int use_moba_seq_limit);
 
 std::vector<paddle::Tensor> MobaMlpEinsum(
     const paddle::Tensor& k_input,
diff --git a/custom_ops/gpu_ops/moba_attn/moba_attn_utils.hpp b/custom_ops/gpu_ops/moba_attn/moba_attn_utils.hpp
index 824b2e0f90..0a77e4bc69 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_attn_utils.hpp
+++ b/custom_ops/gpu_ops/moba_attn/moba_attn_utils.hpp
@@ -41,708 +41,858 @@
 
 using namespace cute;
 
-template<typename T>
+template <typename T>
 struct PackedHalf;
 
-template<>
+template <>
 struct PackedHalf<cutlass::half_t> {
-    using Type = __half2;
+  using Type = __half2;
 };
 
-template<>
+template <>
 struct PackedHalf<cutlass::bfloat16_t> {
-    using Type = nv_bfloat162;
+  using Type = nv_bfloat162;
 };
 
-template<>
+template <>
 struct PackedHalf<phi::dtype::float16> {
-    using Type = __half2;
+  using Type = __half2;
 };
 
-template<>
+template <>
 struct PackedHalf<phi::dtype::bfloat16> {
-    using Type = nv_bfloat162;
+  using Type = nv_bfloat162;
 };
 
-
-template<typename T>
+template <typename T>
 struct HalfSub;
 
-template<>
+template <>
 struct HalfSub<cutlass::half_t> {
-    inline __device__ void operator()(uint32_t* result_ptr, const uint32_t magic_num) {
-        asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(*result_ptr) : "r"(*result_ptr), "r"(magic_num));
-    }
+  inline __device__ void operator()(uint32_t *result_ptr,
+                                    const uint32_t magic_num) {
+    asm volatile("sub.f16x2 %0, %1, %2;\n"
+                 : "=r"(*result_ptr)
+                 : "r"(*result_ptr), "r"(magic_num));
+  }
 };
 
-template<>
+template <>
 struct HalfSub<cutlass::bfloat16_t> {
-    inline __device__ void operator()(uint32_t* result_ptr, const uint32_t magic_num) {
-        *reinterpret_cast<nv_bfloat162*>(result_ptr) -= *reinterpret_cast<const nv_bfloat162*>(&magic_num);
-    }
+  inline __device__ void operator()(uint32_t *result_ptr,
+                                    const uint32_t magic_num) {
+    *reinterpret_cast<nv_bfloat162 *>(result_ptr) -=
+        *reinterpret_cast<const nv_bfloat162 *>(&magic_num);
+  }
 };
 
-template<typename T>
+template <typename T>
 struct HalfMul;
 
-template<>
+template <>
 struct HalfMul<cutlass::half_t> {
-    inline __device__ void operator()(uint32_t* result_ptr, const uint32_t magic_num) {
-        asm volatile("mul.f16x2 %0, %1, %2;\n" : "=r"(*result_ptr) : "r"(*result_ptr), "r"(magic_num));
-    }
+  inline __device__ void operator()(uint32_t *result_ptr,
+                                    const uint32_t magic_num) {
+    asm volatile("mul.f16x2 %0, %1, %2;\n"
+                 : "=r"(*result_ptr)
+                 : "r"(*result_ptr), "r"(magic_num));
+  }
 };
 
-template<>
+template <>
 struct HalfMul<cutlass::bfloat16_t> {
-    inline __device__ void operator()(uint32_t* result_ptr, const uint32_t magic_num) {
-        *reinterpret_cast<nv_bfloat162*>(result_ptr) *= *reinterpret_cast<const nv_bfloat162*>(&magic_num);
-    }
+  inline __device__ void operator()(uint32_t *result_ptr,
+                                    const uint32_t magic_num) {
+    *reinterpret_cast<nv_bfloat162 *>(result_ptr) *=
+        *reinterpret_cast<const nv_bfloat162 *>(&magic_num);
+  }
 };
 
-
-template<typename T>
+template <typename T>
 struct HalfMax;
-template<>
+template <>
 struct HalfMax<cutlass::half_t> {
-    inline __device__ __half2 operator()(const __half2 x, const __half2 y) {
-        __half2 res;
-        asm volatile("max.f16x2 %0, %1, %2;\n" :
-            "=r"(*reinterpret_cast<uint32_t*>(&res)) :
-            "r"(*reinterpret_cast<const uint32_t*>(&x)),
-            "r"(*reinterpret_cast<const uint32_t*>(&y)));
-        return res;
-    }
+  inline __device__ __half2 operator()(const __half2 x, const __half2 y) {
+    __half2 res;
+    asm volatile("max.f16x2 %0, %1, %2;\n"
+                 : "=r"(*reinterpret_cast<uint32_t *>(&res))
+                 : "r"(*reinterpret_cast<const uint32_t *>(&x)),
+                   "r"(*reinterpret_cast<const uint32_t *>(&y)));
+    return res;
+  }
 };
 
-template<>
+template <>
 struct HalfMax<cutlass::bfloat16_t> {
-    inline __device__ nv_bfloat162 operator()(const nv_bfloat162 x, const nv_bfloat162 y) {
-        nv_bfloat162 res;
-        asm volatile("max.bf16x2 %0, %1, %2;\n" :
-            "=r"(*reinterpret_cast<uint32_t*>(&res)) :
-            "r"(*reinterpret_cast<const uint32_t*>(&x)),
-            "r"(*reinterpret_cast<const uint32_t*>(&y)));
-        return res;
-    }
+  inline __device__ nv_bfloat162 operator()(const nv_bfloat162 x,
+                                            const nv_bfloat162 y) {
+    nv_bfloat162 res;
+    asm volatile("max.bf16x2 %0, %1, %2;\n"
+                 : "=r"(*reinterpret_cast<uint32_t *>(&res))
+                 : "r"(*reinterpret_cast<const uint32_t *>(&x)),
+                   "r"(*reinterpret_cast<const uint32_t *>(&y)));
+    return res;
+  }
 };
 
-
-template<typename T>
+template <typename T>
 struct HalfMin;
-template<>
+template <>
 struct HalfMin<cutlass::half_t> {
-    inline __device__ __half2 operator()(const __half2 x, const __half2 y) {
-        __half2 res;
-        asm volatile("min.f16x2 %0, %1, %2;\n" :
-            "=r"(*reinterpret_cast<uint32_t*>(&res)) :
-            "r"(*reinterpret_cast<const uint32_t*>(&x)),
-            "r"(*reinterpret_cast<const uint32_t*>(&y)));
-        return res;
-    }
+  inline __device__ __half2 operator()(const __half2 x, const __half2 y) {
+    __half2 res;
+    asm volatile("min.f16x2 %0, %1, %2;\n"
+                 : "=r"(*reinterpret_cast<uint32_t *>(&res))
+                 : "r"(*reinterpret_cast<const uint32_t *>(&x)),
+                   "r"(*reinterpret_cast<const uint32_t *>(&y)));
+    return res;
+  }
 };
 
-template<>
+template <>
 struct HalfMin<cutlass::bfloat16_t> {
-    inline __device__ nv_bfloat162 operator()(const nv_bfloat162 x, const nv_bfloat162 y) {
-        nv_bfloat162 res;
-        asm volatile("min.bf16x2 %0, %1, %2;\n" :
-            "=r"(*reinterpret_cast<uint32_t*>(&res)) :
-            "r"(*reinterpret_cast<const uint32_t*>(&x)),
-            "r"(*reinterpret_cast<const uint32_t*>(&y)));
-        return res;
-    }
+  inline __device__ nv_bfloat162 operator()(const nv_bfloat162 x,
+                                            const nv_bfloat162 y) {
+    nv_bfloat162 res;
+    asm volatile("min.bf16x2 %0, %1, %2;\n"
+                 : "=r"(*reinterpret_cast<uint32_t *>(&res))
+                 : "r"(*reinterpret_cast<const uint32_t *>(&x)),
+                   "r"(*reinterpret_cast<const uint32_t *>(&y)));
+    return res;
+  }
 };
 
-
-template<typename T>
+template <typename T>
 struct MaxOp {
-__device__ __forceinline__ T operator()(T const & x, T const & y) { return x > y ? x : y; }
+  __device__ __forceinline__ T operator()(T const &x, T const &y) {
+    return x > y ? x : y;
+  }
 };
 
 template <>
 struct MaxOp<float> {
-__device__ __forceinline__ float operator()(float const &x, float const &y) { return max(x, y); }
+  __device__ __forceinline__ float operator()(float const &x, float const &y) {
+    return max(x, y);
+  }
 };
 
-template<typename T>
+template <typename T>
 struct MinOp {
-__device__ __forceinline__ T operator()(T const & x, T const & y) { return x < y ? x : y; }
+  __device__ __forceinline__ T operator()(T const &x, T const &y) {
+    return x < y ? x : y;
+  }
 };
 
 template <>
 struct MinOp<float> {
-__device__ __forceinline__ float operator()(float const &x, float const &y) { return min(x, y); }
+  __device__ __forceinline__ float operator()(float const &x, float const &y) {
+    return min(x, y);
+  }
 };
 
-
-template<typename T>
+template <typename T>
 struct SumOp {
-__device__ __forceinline__ T operator()(T const & x, T const & y) { return x + y; }
+  __device__ __forceinline__ T operator()(T const &x, T const &y) {
+    return x + y;
+  }
 };
 
-template<typename T, bool Is_K>
-inline __device__ static void convert_c8_2_half(uint32_t *src, T *dst, const T *cache_scale, const T* cache_zp) {
-    uint32_t* half_result_ptr = reinterpret_cast<uint32_t*>(dst);
-    if constexpr (std::is_same_v<T, cutlass::bfloat16_t>) {
-        static constexpr uint32_t fp32_base = 0x4B000000;
-        float fp32_intermediates[4];
+template <typename T, bool Is_K>
+inline __device__ static void convert_c8_2_half(uint32_t *src,
+                                                T *dst,
+                                                const T *cache_scale,
+                                                const T *cache_zp) {
+  uint32_t *half_result_ptr = reinterpret_cast<uint32_t *>(dst);
+  if constexpr (std::is_same_v<T, cutlass::bfloat16_t>) {
+    static constexpr uint32_t fp32_base = 0x4B000000;
+    float fp32_intermediates[4];
 
-        uint32_t* fp32_intermediates_casted = reinterpret_cast<uint32_t*>(fp32_intermediates);
-        fp32_intermediates_casted[0] = __byte_perm(*src, fp32_base, 0x7650);
-        fp32_intermediates_casted[1] = __byte_perm(*src, fp32_base, 0x7651);
-        fp32_intermediates_casted[2] = __byte_perm(*src, fp32_base, 0x7652);
-        fp32_intermediates_casted[3] = __byte_perm(*src, fp32_base, 0x7653);
+    uint32_t *fp32_intermediates_casted =
+        reinterpret_cast<uint32_t *>(fp32_intermediates);
+    fp32_intermediates_casted[0] = __byte_perm(*src, fp32_base, 0x7650);
+    fp32_intermediates_casted[1] = __byte_perm(*src, fp32_base, 0x7651);
+    fp32_intermediates_casted[2] = __byte_perm(*src, fp32_base, 0x7652);
+    fp32_intermediates_casted[3] = __byte_perm(*src, fp32_base, 0x7653);
 
-        #pragma unroll
-        for (int ii = 0; ii < 4; ++ii) {
-            fp32_intermediates[ii] -= 8388608.f;
-        }
+#pragma unroll
+    for (int ii = 0; ii < 4; ++ii) {
+      fp32_intermediates[ii] -= 8388608.f;
+    }
 
-        #pragma unroll
-        for (int ii = 0; ii < 2; ++ii) {
-            half_result_ptr[ii] = __byte_perm(fp32_intermediates_casted[2 * ii + 0], fp32_intermediates_casted[2 * ii + 1], 0x7632);
-        }
+#pragma unroll
+    for (int ii = 0; ii < 2; ++ii) {
+      half_result_ptr[ii] = __byte_perm(fp32_intermediates_casted[2 * ii + 0],
+                                        fp32_intermediates_casted[2 * ii + 1],
+                                        0x7632);
+    }
+  } else {
+    static constexpr uint32_t head_for_fp16 = 0x64006400;
+    half_result_ptr[0] = __byte_perm(*src, head_for_fp16, 0x7150);
+    half_result_ptr[1] = __byte_perm(*src, head_for_fp16, 0x7352);
+  }
+
+  using pack_half = typename PackedHalf<T>::Type;
+#pragma unroll
+  for (int i = 0; i < 2; i++) {
+    if constexpr (Is_K) {
+      HalfSub<T>()(half_result_ptr + i,
+                   *reinterpret_cast<const uint32_t *>(cache_zp + i * 2));
+      HalfMul<T>()(half_result_ptr + i,
+                   *reinterpret_cast<const uint32_t *>(cache_scale + i * 2));
     } else {
-        static constexpr uint32_t head_for_fp16 = 0x64006400;
-        half_result_ptr[0] = __byte_perm(*src, head_for_fp16, 0x7150);
-        half_result_ptr[1] = __byte_perm(*src, head_for_fp16, 0x7352);
-    }
-
-    using pack_half = typename PackedHalf<T>::Type;
-    #pragma unroll
-    for (int i = 0; i < 2; i++){
-        if constexpr (Is_K) {
-            HalfSub<T>()(half_result_ptr + i, *reinterpret_cast<const uint32_t*>(cache_zp + i * 2));
-            HalfMul<T>()(half_result_ptr + i, *reinterpret_cast<const uint32_t*>(cache_scale + i * 2));
-        } else {
-            pack_half bias;
-            pack_half scale;
-            bias.x = cache_zp[0];
-            bias.y = cache_zp[0];
-            scale.x = cache_scale[0];
-            scale.y = cache_scale[0];
-            HalfSub<T>()(half_result_ptr + i, *reinterpret_cast<const uint32_t*>(&bias));
-            HalfMul<T>()(half_result_ptr + i, *reinterpret_cast<const uint32_t*>(&scale));
-        }
+      pack_half bias;
+      pack_half scale;
+      bias.x = cache_zp[0];
+      bias.y = cache_zp[0];
+      scale.x = cache_scale[0];
+      scale.y = cache_scale[0];
+      HalfSub<T>()(half_result_ptr + i,
+                   *reinterpret_cast<const uint32_t *>(&bias));
+      HalfMul<T>()(half_result_ptr + i,
+                   *reinterpret_cast<const uint32_t *>(&scale));
     }
+  }
 }
 
-template<typename T, bool Is_K>
-inline __device__ static void convert_c4_2_half(uint32_t *src, T *dst, const T *cache_scale, const T* cache_zp) {
-    using pack_half = typename PackedHalf<T>::Type;
-    static constexpr uint32_t MASK = 0x0f0f0f0f;
-    static constexpr uint32_t head_for_fp16 = std::is_same_v<T, cutlass::bfloat16_t> ? 0x43004300 : 0x64006400;
-    static constexpr uint32_t mask_for_c42fp16_one = 0x7253;
-    static constexpr uint32_t mask_for_c42fp16_two = 0x7051;
-    uint32_t* result_ptr = reinterpret_cast<uint32_t*>(dst);
-    uint32_t source = *reinterpret_cast<uint32_t const*>(src);
-    // source = {e0 e4 e1 e5 e2 e6 e3 e7}
-    uint32_t bottom_i4s = source & MASK;
-    // bottom_i4s = {0 e4 0 e5 0 e6 0 e7}
-    uint32_t top_i4s = (source >> 4) & MASK;
-    // top_i4s = {0 e0 0 e1 0 e2 0 e3}
-    asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(result_ptr[0]) : "r"(top_i4s), "n"(head_for_fp16), "n"(mask_for_c42fp16_one));
-    // result_ptr[0] = {e0 e1}
-    asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(result_ptr[1]) : "r"(top_i4s), "n"(head_for_fp16), "n"(mask_for_c42fp16_two));
-    asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(result_ptr[2]) : "r"(bottom_i4s), "n"(head_for_fp16), "n"(mask_for_c42fp16_one));
-    asm volatile("prmt.b32 %0,%1,%2,%3;\n" : "=r"(result_ptr[3]) : "r"(bottom_i4s), "n"(head_for_fp16), "n"(mask_for_c42fp16_two));
+template <typename T, bool Is_K>
+inline __device__ static void convert_c4_2_half(uint32_t *src,
+                                                T *dst,
+                                                const T *cache_scale,
+                                                const T *cache_zp) {
+  using pack_half = typename PackedHalf<T>::Type;
+  static constexpr uint32_t MASK = 0x0f0f0f0f;
+  static constexpr uint32_t head_for_fp16 =
+      std::is_same_v<T, cutlass::bfloat16_t> ? 0x43004300 : 0x64006400;
+  static constexpr uint32_t mask_for_c42fp16_one = 0x7253;
+  static constexpr uint32_t mask_for_c42fp16_two = 0x7051;
+  uint32_t *result_ptr = reinterpret_cast<uint32_t *>(dst);
+  uint32_t source = *reinterpret_cast<uint32_t const *>(src);
+  // source = {e0 e4 e1 e5 e2 e6 e3 e7}
+  uint32_t bottom_i4s = source & MASK;
+  // bottom_i4s = {0 e4 0 e5 0 e6 0 e7}
+  uint32_t top_i4s = (source >> 4) & MASK;
+  // top_i4s = {0 e0 0 e1 0 e2 0 e3}
+  asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+               : "=r"(result_ptr[0])
+               : "r"(top_i4s), "n"(head_for_fp16), "n"(mask_for_c42fp16_one));
+  // result_ptr[0] = {e0 e1}
+  asm volatile("prmt.b32 %0,%1,%2,%3;\n"
+               : "=r"(result_ptr[1])
+               : "r"(top_i4s), "n"(head_for_fp16), "n"(mask_for_c42fp16_two));
+  asm volatile(
+      "prmt.b32 %0,%1,%2,%3;\n"
+      : "=r"(result_ptr[2])
+      : "r"(bottom_i4s), "n"(head_for_fp16), "n"(mask_for_c42fp16_one));
+  asm volatile(
+      "prmt.b32 %0,%1,%2,%3;\n"
+      : "=r"(result_ptr[3])
+      : "r"(bottom_i4s), "n"(head_for_fp16), "n"(mask_for_c42fp16_two));
 
-    #pragma unroll
-    for (int i = 0; i < 4; ++i) {
-        if constexpr (Is_K) {
-            const int ith_col = i % 2 * 2;
-            HalfSub<T>()(result_ptr + i, *reinterpret_cast<const uint32_t*>(cache_zp + ith_col));
-            HalfMul<T>()(result_ptr + i, *reinterpret_cast<const uint32_t*>(cache_scale + ith_col));
-        } else {
-            const int ith_col = i / 2;
-            pack_half bias;
-            pack_half scale;
-            bias.x = cache_zp[ith_col];
-            bias.y = cache_zp[ith_col];
-            scale.x = cache_scale[ith_col];
-            scale.y = cache_scale[ith_col];
-            HalfSub<T>()(result_ptr + i, *reinterpret_cast<const uint32_t*>(&bias));
-            HalfMul<T>()(result_ptr + i, *reinterpret_cast<const uint32_t*>(&scale));
-        }
+#pragma unroll
+  for (int i = 0; i < 4; ++i) {
+    if constexpr (Is_K) {
+      const int ith_col = i % 2 * 2;
+      HalfSub<T>()(result_ptr + i,
+                   *reinterpret_cast<const uint32_t *>(cache_zp + ith_col));
+      HalfMul<T>()(result_ptr + i,
+                   *reinterpret_cast<const uint32_t *>(cache_scale + ith_col));
+    } else {
+      const int ith_col = i / 2;
+      pack_half bias;
+      pack_half scale;
+      bias.x = cache_zp[ith_col];
+      bias.y = cache_zp[ith_col];
+      scale.x = cache_scale[ith_col];
+      scale.y = cache_scale[ith_col];
+      HalfSub<T>()(result_ptr + i, *reinterpret_cast<const uint32_t *>(&bias));
+      HalfMul<T>()(result_ptr + i, *reinterpret_cast<const uint32_t *>(&scale));
     }
+  }
 }
 
-template<typename CacheKV_traits, typename T, int kHeadDim, int kDataNumPer2Byte, bool A_in_regs=false, typename Tensor0, typename Tensor1,
-         typename Tensor2, typename Tensor3, typename Tensor4,
-         typename TiledMma, typename ThrCopy0, typename TiledCopy0>
-inline __device__ void gemm_qk_quant(
-        Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCsA, Tensor3 &tCrB,
-        Tensor4 const& sB, TiledMma tiled_mma,
-        ThrCopy0 smem_thr_copy_A,
-        TiledCopy0 smem_tiled_copy_A,
-        const int32_t tidx,
-        const T * cache_scale, const T * cache_zp) {
-    CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
-    CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
-    Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
-    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
-    if (!A_in_regs) {
-        copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
-    }
-    uint32_t *sBdata = reinterpret_cast<uint32_t *>(sB.data().get()) + tidx * (kDataNumPer2Byte / 4);
+template <typename CacheKV_traits,
+          typename T,
+          int kHeadDim,
+          int kDataNumPer2Byte,
+          bool A_in_regs = false,
+          typename Tensor0,
+          typename Tensor1,
+          typename Tensor2,
+          typename Tensor3,
+          typename Tensor4,
+          typename TiledMma,
+          typename ThrCopy0,
+          typename TiledCopy0>
+inline __device__ void gemm_qk_quant(Tensor0 &acc,
+                                     Tensor1 &tCrA,
+                                     Tensor2 &tCsA,
+                                     Tensor3 &tCrB,
+                                     Tensor4 const &sB,
+                                     TiledMma tiled_mma,
+                                     ThrCopy0 smem_thr_copy_A,
+                                     TiledCopy0 smem_tiled_copy_A,
+                                     const int32_t tidx,
+                                     const T *cache_scale,
+                                     const T *cache_zp) {
+  CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+  CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+  Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
+  CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
+  if (!A_in_regs) {
+    copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
+  }
+  uint32_t *sBdata = reinterpret_cast<uint32_t *>(sB.data().get()) +
+                     tidx * (kDataNumPer2Byte / 4);
 
-    #pragma unroll
-    for (int i = 0; i < size<2>(tCrA); ++i) {
-        if (i < size<2>(tCrA) - 1) {
-            if (!A_in_regs) {
-                copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
-            }
-        }
-        if constexpr (kDataNumPer2Byte == 4) {
-            convert_c4_2_half<T, true>(sBdata + i * kHeadDim, tCrB.data(), cache_scale + i * 4, cache_zp + i * 4);
-        } else {
-            convert_c8_2_half<T, true>(sBdata + i * (kHeadDim * 2), tCrB.data(), cache_scale + i * 4, cache_zp + i * 4);
-            convert_c8_2_half<T, true>(sBdata + i * (kHeadDim * 2) + 1, tCrB.data() + 4, cache_scale + i * 4, cache_zp + i * 4);
-        }
-
-        cute::gemm(tiled_mma, tCrA(_, _, i), tCrB, acc);
+#pragma unroll
+  for (int i = 0; i < size<2>(tCrA); ++i) {
+    if (i < size<2>(tCrA) - 1) {
+      if (!A_in_regs) {
+        copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
+      }
     }
+    if constexpr (kDataNumPer2Byte == 4) {
+      convert_c4_2_half<T, true>(sBdata + i * kHeadDim,
+                                 tCrB.data(),
+                                 cache_scale + i * 4,
+                                 cache_zp + i * 4);
+    } else {
+      convert_c8_2_half<T, true>(sBdata + i * (kHeadDim * 2),
+                                 tCrB.data(),
+                                 cache_scale + i * 4,
+                                 cache_zp + i * 4);
+      convert_c8_2_half<T, true>(sBdata + i * (kHeadDim * 2) + 1,
+                                 tCrB.data() + 4,
+                                 cache_scale + i * 4,
+                                 cache_zp + i * 4);
+    }
+
+    cute::gemm(tiled_mma, tCrA(_, _, i), tCrB, acc);
+  }
 }
 
-template<typename CacheKV_traits, typename T, int kHeadDim, int kDataNumPer2Byte, bool A_in_regs=false, typename Tensor0, typename Tensor1,
-         typename Tensor2, typename Tensor3, typename Tensor4,
-         typename TiledMma, typename ThrCopy0, typename TiledCopy0>
-inline __device__ void gemm_value_quant(
-        Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCsA, Tensor3 &tCrB,
-        Tensor4 const& sB, TiledMma tiled_mma,
-        ThrCopy0 smem_thr_copy_A,
-        TiledCopy0 smem_tiled_copy_A,
-        int32_t tidx,
-        const T * cache_scale, const T * cache_zp) {
-    CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
-    CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
-    Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
-    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
-    if (!A_in_regs) {
-        copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
-    }
-    uint32_t *sBdata = reinterpret_cast<uint32_t *>(sB.data().get()) + tidx * (2 * kDataNumPer2Byte / 4);
+template <typename CacheKV_traits,
+          typename T,
+          int kHeadDim,
+          int kDataNumPer2Byte,
+          bool A_in_regs = false,
+          typename Tensor0,
+          typename Tensor1,
+          typename Tensor2,
+          typename Tensor3,
+          typename Tensor4,
+          typename TiledMma,
+          typename ThrCopy0,
+          typename TiledCopy0>
+inline __device__ void gemm_value_quant(Tensor0 &acc,
+                                        Tensor1 &tCrA,
+                                        Tensor2 &tCsA,
+                                        Tensor3 &tCrB,
+                                        Tensor4 const &sB,
+                                        TiledMma tiled_mma,
+                                        ThrCopy0 smem_thr_copy_A,
+                                        TiledCopy0 smem_tiled_copy_A,
+                                        int32_t tidx,
+                                        const T *cache_scale,
+                                        const T *cache_zp) {
+  CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+  CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+  Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
+  CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
+  if (!A_in_regs) {
+    copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
+  }
+  uint32_t *sBdata = reinterpret_cast<uint32_t *>(sB.data().get()) +
+                     tidx * (2 * kDataNumPer2Byte / 4);
 
-    #pragma unroll
-    for (int i = 0; i < size<2>(tCrA); ++i) {
-        const int cur_idx = i * kHeadDim * (2 * kDataNumPer2Byte / 4);
+#pragma unroll
+  for (int i = 0; i < size<2>(tCrA); ++i) {
+    const int cur_idx = i * kHeadDim * (2 * kDataNumPer2Byte / 4);
 
-        if (i < size<2>(tCrA) - 1) {
-            if (!A_in_regs) {
-                copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
-            }
-        }
-        if constexpr (kDataNumPer2Byte == 4) {
-            convert_c4_2_half<T, false>(sBdata + cur_idx, tCrB.data(), cache_scale, cache_zp);
-            convert_c4_2_half<T, false>(sBdata + cur_idx + 1, tCrB.data() + 8, cache_scale + 2, cache_zp + 2);
-        } else {
-            convert_c8_2_half<T, false>(sBdata + cur_idx, tCrB.data(), cache_scale, cache_zp);
-            convert_c8_2_half<T, false>(sBdata + cur_idx + 1, tCrB.data() + 4, cache_scale + 1, cache_zp + 1);
-            convert_c8_2_half<T, false>(sBdata + cur_idx + 2, tCrB.data() + 8, cache_scale + 2, cache_zp + 2);
-            convert_c8_2_half<T, false>(sBdata + cur_idx + 3, tCrB.data() + 12, cache_scale + 3, cache_zp + 3);
-        }
-        cute::gemm(tiled_mma, tCrA(_, _, i), tCrB, acc);
+    if (i < size<2>(tCrA) - 1) {
+      if (!A_in_regs) {
+        copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
+      }
     }
+    if constexpr (kDataNumPer2Byte == 4) {
+      convert_c4_2_half<T, false>(
+          sBdata + cur_idx, tCrB.data(), cache_scale, cache_zp);
+      convert_c4_2_half<T, false>(
+          sBdata + cur_idx + 1, tCrB.data() + 8, cache_scale + 2, cache_zp + 2);
+    } else {
+      convert_c8_2_half<T, false>(
+          sBdata + cur_idx, tCrB.data(), cache_scale, cache_zp);
+      convert_c8_2_half<T, false>(
+          sBdata + cur_idx + 1, tCrB.data() + 4, cache_scale + 1, cache_zp + 1);
+      convert_c8_2_half<T, false>(
+          sBdata + cur_idx + 2, tCrB.data() + 8, cache_scale + 2, cache_zp + 2);
+      convert_c8_2_half<T, false>(sBdata + cur_idx + 3,
+                                  tCrB.data() + 12,
+                                  cache_scale + 3,
+                                  cache_zp + 3);
+    }
+    cute::gemm(tiled_mma, tCrA(_, _, i), tCrB, acc);
+  }
 }
 
-
-template<int kMiLen, typename Engine, typename Layout>
-inline __device__ void apply_mask(Tensor<Engine, Layout> &scores, const uint32_t warp_id, const uint32_t col, const uint32_t reamin_seq_len) {
-    const int cols = size<1>(scores) / 2;
-    #pragma unroll
-    for (int mi = 0; mi < kMiLen; ++mi) {
-        #pragma unroll
-        for (int ni = 0; ni < cols; ++ni) {
-            const int col_index = warp_id * 8 + ni * 32 + col * 2;
-            if (col_index >= reamin_seq_len) {
-                scores(mi, ni * 2) = -INFINITY;
-            }
-            if (col_index + 1 >= reamin_seq_len) {
-                scores(mi, ni * 2 + 1) = -INFINITY;
-            }
-        }
+template <int kMiLen, typename Engine, typename Layout>
+inline __device__ void apply_mask(Tensor<Engine, Layout> &scores,
+                                  const uint32_t warp_id,
+                                  const uint32_t col,
+                                  const uint32_t reamin_seq_len) {
+  const int cols = size<1>(scores) / 2;
+#pragma unroll
+  for (int mi = 0; mi < kMiLen; ++mi) {
+#pragma unroll
+    for (int ni = 0; ni < cols; ++ni) {
+      const int col_index = warp_id * 8 + ni * 32 + col * 2;
+      if (col_index >= reamin_seq_len) {
+        scores(mi, ni * 2) = -INFINITY;
+      }
+      if (col_index + 1 >= reamin_seq_len) {
+        scores(mi, ni * 2 + 1) = -INFINITY;
+      }
     }
+  }
 }
 
-
-template<int THREADS>
+template <int THREADS>
 struct Allreduce {
-    static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
-    template<typename T, typename Operator>
-    static __device__ inline T run(T x, Operator &op) {
-        constexpr int OFFSET = THREADS / 2;
-        x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
-        return Allreduce<OFFSET>::run(x, op);
-    }
+  static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
+  template <typename T, typename Operator>
+  static __device__ inline T run(T x, Operator &op) {
+    constexpr int OFFSET = THREADS / 2;
+    x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
+    return Allreduce<OFFSET>::run(x, op);
+  }
 };
 
-template<>
+template <>
 struct Allreduce<2> {
-template<typename T, typename Operator>
-static __device__ inline T run(T x, Operator &op) {
+  template <typename T, typename Operator>
+  static __device__ inline T run(T x, Operator &op) {
     x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
     return x;
-}
+  }
 };
 
-template<int kMiLen, typename Engine0, typename Layout0, typename T>
-__device__ inline void reduce_max(Tensor<Engine0, Layout0> const& tensor, T *scores_max){
-    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
-    MaxOp<T> max_op;
-    #pragma unroll
-    for (int mi = 0; mi < kMiLen; ++mi) {
-        #pragma unroll
-        for (int ni = 0; ni < size<1>(tensor); ni++) {
-            scores_max[mi] = max_op(scores_max[mi], tensor(mi, ni));
-        }
-        scores_max[mi] = Allreduce<4>::run(scores_max[mi], max_op);
+template <int kMiLen, typename Engine0, typename Layout0, typename T>
+__device__ inline void reduce_max(Tensor<Engine0, Layout0> const &tensor,
+                                  T *scores_max) {
+  static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+  MaxOp<T> max_op;
+#pragma unroll
+  for (int mi = 0; mi < kMiLen; ++mi) {
+#pragma unroll
+    for (int ni = 0; ni < size<1>(tensor); ni++) {
+      scores_max[mi] = max_op(scores_max[mi], tensor(mi, ni));
     }
+    scores_max[mi] = Allreduce<4>::run(scores_max[mi], max_op);
+  }
 }
 
 template <int kMiLen, typename Engine0, typename Layout0, typename T>
-inline __device__ void scale_apply_exp2(Tensor<Engine0, Layout0> &tensor, T const *max, T *sum, const float scale) {
-    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
-    #pragma unroll
-    for (int mi = 0; mi < kMiLen; ++mi) {
-        const float max_scaled = max[mi] * scale;
-        #pragma unroll
-        for (int ni = 0; ni < size<1>(tensor); ++ni)  {
-            tensor(mi, ni) = expf(tensor(mi, ni) * scale - max_scaled);
-            sum[mi] += tensor(mi, ni);
-        }
+inline __device__ void scale_apply_exp2(Tensor<Engine0, Layout0> &tensor,
+                                        T const *max,
+                                        T *sum,
+                                        const float scale) {
+  static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+#pragma unroll
+  for (int mi = 0; mi < kMiLen; ++mi) {
+    const float max_scaled = max[mi] * scale;
+#pragma unroll
+    for (int ni = 0; ni < size<1>(tensor); ++ni) {
+      tensor(mi, ni) = expf(tensor(mi, ni) * scale - max_scaled);
+      sum[mi] += tensor(mi, ni);
     }
+  }
 }
 
-
 template <typename paddle_type>
 struct cuteType;
 
 template <>
 struct cuteType<phi::dtype::float16> {
-    using type = cutlass::half_t;
+  using type = cutlass::half_t;
 };
 
 template <>
 struct cuteType<phi::dtype::bfloat16> {
-    using type = cutlass::bfloat16_t;
+  using type = cutlass::bfloat16_t;
 };
 
-template<typename T>
+template <typename T>
 __forceinline__ __device__ auto float_2_half2(const float x) {
-    if constexpr (std::is_same<T, cutlass::half_t>::value) {
-        return __float2half2_rn(x);
-    } else {
-        return __float2bfloat162_rn(x);
-    }
+  if constexpr (std::is_same<T, cutlass::half_t>::value) {
+    return __float2half2_rn(x);
+  } else {
+    return __float2bfloat162_rn(x);
+  }
 }
 
-
 struct uint16 {
-    uint4 u;
-    uint4 v;
-    uint4 s;
-    uint4 t;
+  uint4 u;
+  uint4 v;
+  uint4 s;
+  uint4 t;
 };
 
-
 struct uint8 {
-    uint4 u;
-    uint4 v;
+  uint4 u;
+  uint4 v;
 };
 
-template<int BYTES>
+template <int BYTES>
 struct BytesToType {};
 
-template<>
+template <>
 struct BytesToType<64> {
-    using Type = uint16;
-    static_assert(sizeof(Type) == 64);
+  using Type = uint16;
+  static_assert(sizeof(Type) == 64);
 };
 
-template<>
+template <>
 struct BytesToType<32> {
-    using Type = uint8;
-    static_assert(sizeof(Type) == 32);
+  using Type = uint8;
+  static_assert(sizeof(Type) == 32);
 };
 
-template<>
+template <>
 struct BytesToType<16> {
-    using Type = uint4;
-    static_assert(sizeof(Type) == 16);
+  using Type = uint4;
+  static_assert(sizeof(Type) == 16);
 };
 
-template<>
+template <>
 struct BytesToType<8> {
-    using Type = uint64_t;
-    static_assert(sizeof(Type) == 8);
+  using Type = uint64_t;
+  static_assert(sizeof(Type) == 8);
 };
 
-template<>
+template <>
 struct BytesToType<4> {
-    using Type = uint32_t;
-    static_assert(sizeof(Type) == 4);
+  using Type = uint32_t;
+  static_assert(sizeof(Type) == 4);
 };
 
-template<>
+template <>
 struct BytesToType<2> {
-    using Type = uint16_t;
-    static_assert(sizeof(Type) == 2);
+  using Type = uint16_t;
+  static_assert(sizeof(Type) == 2);
 };
 
-template<>
+template <>
 struct BytesToType<1> {
-    using Type = uint8_t;
-    static_assert(sizeof(Type) == 1);
+  using Type = uint8_t;
+  static_assert(sizeof(Type) == 1);
 };
 
-template<typename Elt_type, uint32_t NUM_ELT>
+template <typename Elt_type, uint32_t NUM_ELT>
 struct Vec {
+  enum { BYTES = NUM_ELT * sizeof(Elt_type) };
 
-    enum { BYTES = NUM_ELT * sizeof(Elt_type) };
+  using Vec_type = typename BytesToType<BYTES>::Type;
 
-    using Vec_type = typename BytesToType<BYTES>::Type;
+  using Alias_type = union {
+    Vec_type vec;
+    Elt_type elt[NUM_ELT];
+  };
 
-    using Alias_type = union {
-        Vec_type vec;
-        Elt_type elt[NUM_ELT];
-    };
+  Alias_type data;
 
-    Alias_type data;
+  inline __device__ Vec() {}
 
-    inline __device__ Vec() {}
-
-    template<typename S>
-    inline __device__ void to(Vec<S, NUM_ELT> &other) {
-        #pragma unroll
-        for( int it = 0; it < NUM_ELT; it++ ) {
-            other.data.elt[it] = S(this->data.elt[it]);
-        }
+  template <typename S>
+  inline __device__ void to(Vec<S, NUM_ELT> &other) {
+#pragma unroll
+    for (int it = 0; it < NUM_ELT; it++) {
+      other.data.elt[it] = S(this->data.elt[it]);
     }
+  }
 
-    template<typename Op>
-    inline __device__ void assign(const Op &op) {
-        #pragma unroll
-        for( int it = 0; it < NUM_ELT; it++ ) {
-            this->data.elt[it] = op(it);
-        }
+  template <typename Op>
+  inline __device__ void assign(const Op &op) {
+#pragma unroll
+    for (int it = 0; it < NUM_ELT; it++) {
+      this->data.elt[it] = op(it);
     }
+  }
 
-    inline __device__ void load_from(const void *base_ptr) {
-        this->data.vec = *reinterpret_cast<const Vec_type *>(base_ptr);
+  inline __device__ void load_from(const void *base_ptr) {
+    this->data.vec = *reinterpret_cast<const Vec_type *>(base_ptr);
+  }
+
+  inline __device__ void store_to(void *base_ptr) {
+    *reinterpret_cast<Vec_type *>(base_ptr) = this->data.vec;
+  }
+
+  inline __device__ void add(const Vec<Elt_type, NUM_ELT> &other) {
+    static_assert(NUM_ELT % 2 == 0);
+    using type = typename PackedHalf<Elt_type>::Type;
+#pragma unroll
+    for (int it = 0; it < NUM_ELT / 2; it++) {
+      type b = *reinterpret_cast<const type *>(other.data.elt + it * 2);
+      *reinterpret_cast<type *>(this->data.elt + it * 2) += b;
     }
+  }
 
-
-    inline __device__ void store_to(void *base_ptr) {
-        *reinterpret_cast<Vec_type *>(base_ptr) = this->data.vec;
+  inline __device__ void set_zero() {
+    constexpr int size = sizeof(Vec_type) / sizeof(int);
+#pragma unroll
+    for (int i = 0; i < size; ++i) {
+      (reinterpret_cast<int *>(this->data.elt))[i] = 0;
     }
+  }
 
-    inline __device__ void add(const Vec<Elt_type, NUM_ELT> &other) {
-        static_assert(NUM_ELT % 2 == 0);
-        using type = typename PackedHalf<Elt_type>::Type;
-        #pragma unroll
-        for (int it = 0; it < NUM_ELT / 2; it++) {
-            type b = *reinterpret_cast<const type *>(other.data.elt + it * 2);
-            *reinterpret_cast<type *>(this->data.elt + it * 2) += b;
-        }
-    }
-
-    inline __device__ void set_zero() {
-        constexpr int size = sizeof(Vec_type) / sizeof(int);
-        #pragma unroll
-        for (int i = 0; i < size; ++i) {
-            (reinterpret_cast<int *>(this->data.elt))[i] = 0;
-        }
-    }
-
-    inline __device__ void fma(const Vec<Elt_type, NUM_ELT> &scale, const Vec<Elt_type, NUM_ELT> &bias) {
-        static_assert(NUM_ELT % 2 == 0);
-        using type = typename PackedHalf<Elt_type>::Type;
-        #pragma unroll
-        for (int it = 0; it < NUM_ELT / 2; it++) {
-            type a = *reinterpret_cast<const type *>(scale.data.elt + it * 2);
-            type b = *reinterpret_cast<const type *>(bias.data.elt + it * 2);
-            *reinterpret_cast<type *>(this->data.elt + it * 2) += a * b;
-        }
+  inline __device__ void fma(const Vec<Elt_type, NUM_ELT> &scale,
+                             const Vec<Elt_type, NUM_ELT> &bias) {
+    static_assert(NUM_ELT % 2 == 0);
+    using type = typename PackedHalf<Elt_type>::Type;
+#pragma unroll
+    for (int it = 0; it < NUM_ELT / 2; it++) {
+      type a = *reinterpret_cast<const type *>(scale.data.elt + it * 2);
+      type b = *reinterpret_cast<const type *>(bias.data.elt + it * 2);
+      *reinterpret_cast<type *>(this->data.elt + it * 2) += a * b;
     }
+  }
 };
 
-template<typename T, int PackSize>
-inline __device__ void apply_rotary_embedding(Vec<T, PackSize>& vec, Vec<float, PackSize / 2>& cos, Vec<float, PackSize / 2>& sin) {
-    static_assert(PackSize % 2 == 0);
-    #pragma unroll
-    for (int i = 0; i < PackSize / 2; i++) {
-        const float cos_inv_freq = cos.data.elt[i];
-        const float sin_inv_freq = sin.data.elt[i];
-        const float v1 = static_cast<float>(vec.data.elt[2 * i]);
-        const float v2 = static_cast<float>(vec.data.elt[2 * i + 1]);
-        vec.data.elt[2 * i] = static_cast<T>(cos_inv_freq * v1 - sin_inv_freq * v2);
-        vec.data.elt[2 * i + 1] = static_cast<T>(sin_inv_freq * v1 + cos_inv_freq * v2);
-    }
+template <typename T, int PackSize>
+inline __device__ void apply_rotary_embedding(Vec<T, PackSize> &vec,
+                                              Vec<float, PackSize / 2> &cos,
+                                              Vec<float, PackSize / 2> &sin) {
+  static_assert(PackSize % 2 == 0);
+#pragma unroll
+  for (int i = 0; i < PackSize / 2; i++) {
+    const float cos_inv_freq = cos.data.elt[i];
+    const float sin_inv_freq = sin.data.elt[i];
+    const float v1 = static_cast<float>(vec.data.elt[2 * i]);
+    const float v2 = static_cast<float>(vec.data.elt[2 * i + 1]);
+    vec.data.elt[2 * i] = static_cast<T>(cos_inv_freq * v1 - sin_inv_freq * v2);
+    vec.data.elt[2 * i + 1] =
+        static_cast<T>(sin_inv_freq * v1 + cos_inv_freq * v2);
+  }
 }
 
-template <bool Is_even_MN=true, typename TiledCopy, typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Engine2, typename Layout2>
+template <bool Is_even_MN = true,
+          typename TiledCopy,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1,
+          typename Engine2,
+          typename Layout2>
 __forceinline__ __device__ void copy(
-        TiledCopy tiled_copy, Tensor<Engine0, Layout0> const &S,
-        Tensor<Engine1, Layout1> &D,
-        Tensor<Engine2, Layout2> const &identity_MN,
-        const int max_MN = 0) {
-    CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});
-    CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});
-    CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D));                     // MMA
-    CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D));                     // MMA_M
-    CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D));                     // MMA_K
-    #pragma unroll
-    for (int m = 0; m < size<1>(S); ++m) {
-        if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) {
-            #pragma unroll
-            for (int k = 0; k < size<2>(S); ++k) {
-                cute::copy(tiled_copy, S(_, m, k), D(_, m, k));
-            }
-        }
+    TiledCopy tiled_copy,
+    Tensor<Engine0, Layout0> const &S,
+    Tensor<Engine1, Layout1> &D,
+    Tensor<Engine2, Layout2> const &identity_MN,
+    const int max_MN = 0) {
+  CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});
+  CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});
+  CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D));  // MMA
+  CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D));  // MMA_M
+  CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D));  // MMA_K
+#pragma unroll
+  for (int m = 0; m < size<1>(S); ++m) {
+    if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) {
+#pragma unroll
+      for (int k = 0; k < size<2>(S); ++k) {
+        cute::copy(tiled_copy, S(_, m, k), D(_, m, k));
+      }
     }
+  }
 }
 
-template<bool A_in_regs=false, bool B_in_regs=false, typename Tensor0, typename Tensor1,
-         typename Tensor2, typename Tensor3, typename Tensor4,
-         typename TiledMma, typename ThrCopy0, typename ThrCopy1,
-         typename TiledCopy0, typename TiledCopy1>
-inline __device__ void gemm(
-        Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const& tCsA,
-        Tensor4 const& tCsB, TiledMma tiled_mma,
-        ThrCopy0 &smem_thr_copy_A, ThrCopy1 &smem_thr_copy_B,
-        TiledCopy0 &smem_tiled_copy_A, TiledCopy1 &smem_tiled_copy_B) {
-    CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
-    CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
-    CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
-    Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
-    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
-    Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
-    CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
+template <bool A_in_regs = false,
+          bool B_in_regs = false,
+          typename Tensor0,
+          typename Tensor1,
+          typename Tensor2,
+          typename Tensor3,
+          typename Tensor4,
+          typename TiledMma,
+          typename ThrCopy0,
+          typename ThrCopy1,
+          typename TiledCopy0,
+          typename TiledCopy1>
+inline __device__ void gemm(Tensor0 &acc,
+                            Tensor1 &tCrA,
+                            Tensor2 &tCrB,
+                            Tensor3 const &tCsA,
+                            Tensor4 const &tCsB,
+                            TiledMma tiled_mma,
+                            ThrCopy0 &smem_thr_copy_A,
+                            ThrCopy1 &smem_thr_copy_B,
+                            TiledCopy0 &smem_tiled_copy_A,
+                            TiledCopy1 &smem_tiled_copy_B) {
+  CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+  CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+  CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
+  Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
+  CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
+  Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+  CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
 
-    if (!A_in_regs) { copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{})); }
-    if (!B_in_regs) { copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{})); }
+  if (!A_in_regs) {
+    copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
+  }
+  if (!B_in_regs) {
+    copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
+  }
 
-    #pragma unroll
-    for (int i = 0; i < size<2>(tCrA); ++i) {
-        if (i < size<2>(tCrA) - 1) {
-            if (!A_in_regs) { copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1)); }
-            if (!B_in_regs) { copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1)); }
-        }
-        cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
+#pragma unroll
+  for (int i = 0; i < size<2>(tCrA); ++i) {
+    if (i < size<2>(tCrA) - 1) {
+      if (!A_in_regs) {
+        copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
+      }
+      if (!B_in_regs) {
+        copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
+      }
     }
+    cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
+  }
 }
 
 template <typename To_type, typename Engine, typename Layout>
 inline __device__ auto convert_type(Tensor<Engine, Layout> const &tensor) {
-    using From_type = typename Engine::value_type;
-    constexpr int numel = decltype(size(tensor))::value;
-    cutlass::NumericArrayConverter<To_type, From_type, numel> convert_op;
-    auto frag = convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel> *>(tensor.data()));
-    return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());
+  using From_type = typename Engine::value_type;
+  constexpr int numel = decltype(size(tensor))::value;
+  cutlass::NumericArrayConverter<To_type, From_type, numel> convert_op;
+  auto frag =
+      convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel> *>(
+          tensor.data()));
+  return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());
 }
 
-template<typename T, typename ReductionOp, int block_size>
+template <typename T, typename ReductionOp, int block_size>
 __inline__ __device__ T BlockAllReduce(T val) {
-    typedef cub::BlockReduce<T, block_size> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
-    __shared__ T result_broadcast;
-    T result = BlockReduce(temp_storage).Reduce(val, ReductionOp());
-    if (threadIdx.x == 0) { result_broadcast = result; }
-    __syncthreads();
-    return result_broadcast;
+  typedef cub::BlockReduce<T, block_size> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ T result_broadcast;
+  T result = BlockReduce(temp_storage).Reduce(val, ReductionOp());
+  if (threadIdx.x == 0) {
+    result_broadcast = result;
+  }
+  __syncthreads();
+  return result_broadcast;
 }
 
-template<typename MMA_traits, typename Layout>
+template <typename MMA_traits, typename Layout>
 __forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
-    using X = Underscore;
-    if constexpr (decltype(rank<0>(acc_layout))::value == 3) {  // SM90
-        static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
-        static_assert(decltype(size<0, 1>(acc_layout))::value == 2);
-        static_assert(decltype(rank(acc_layout))::value == 3);
-        static_assert(decltype(rank(get<0>(acc_layout)))::value == 3);
-        auto l = logical_divide(get<0>(acc_layout), Shape<X, X, _2>{});  // (2, 2, (2, N / 16)))
-        return make_layout(make_layout(get<0>(l), get<1>(l), get<2, 0>(l)), get<1>(acc_layout), make_layout(get<2, 1>(l), get<2>(acc_layout)));
-    } else {  // SM80
-        static_assert(decltype(size<0>(acc_layout))::value == 4);
-        static_assert(decltype(rank(acc_layout))::value == 3);
-        constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});
-        static_assert(mma_shape_K == 8 || mma_shape_K == 16);
-        if constexpr (mma_shape_K == 8) {
-            return acc_layout;
-        } else {
-            auto l = logical_divide(acc_layout, Shape<X, X, _2>{});  // (4, MMA_M, (2, MMA_N / 2)))
-            return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
-        }
-    }
-};
-
-template <bool zero_init=false, int wg_wait=0, bool arrive=true, bool commit=true, typename Tensor0, typename Tensor1, typename Tensor2, typename TiledMma>
-__forceinline__ __device__ void gemm(TiledMma &tiled_mma, Tensor0 const &tCrA, Tensor1 const &tCrB, Tensor2 &tCrC) {
-    constexpr bool Is_RS = !cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value;
-    if constexpr (Is_RS) { warpgroup_fence_operand(const_cast<Tensor0 &>(tCrA)); }
-    warpgroup_fence_operand(tCrC);
-    if constexpr (arrive) {
-        warpgroup_arrive();
-    }
-    if constexpr (zero_init) {
-        tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
-        CUTLASS_PRAGMA_UNROLL
-        for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
-          cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
-          tiled_mma.accumulate_ = GMMA::ScaleOut::One;
-        }
+  using X = Underscore;
+  if constexpr (decltype(rank<0>(acc_layout))::value == 3) {  // SM90
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(size<0, 1>(acc_layout))::value == 2);
+    static_assert(decltype(rank(acc_layout))::value == 3);
+    static_assert(decltype(rank(get<0>(acc_layout)))::value == 3);
+    auto l = logical_divide(get<0>(acc_layout),
+                            Shape<X, X, _2>{});  // (2, 2, (2, N / 16)))
+    return make_layout(make_layout(get<0>(l), get<1>(l), get<2, 0>(l)),
+                       get<1>(acc_layout),
+                       make_layout(get<2, 1>(l), get<2>(acc_layout)));
+  } else {  // SM80
+    static_assert(decltype(size<0>(acc_layout))::value == 4);
+    static_assert(decltype(rank(acc_layout))::value == 3);
+    constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});
+    static_assert(mma_shape_K == 8 || mma_shape_K == 16);
+    if constexpr (mma_shape_K == 8) {
+      return acc_layout;
     } else {
-        CUTLASS_PRAGMA_UNROLL
-        for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
-          cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
-          tiled_mma.accumulate_ = GMMA::ScaleOut::One;
-        }
-    }
-    if constexpr (commit) {
-        warpgroup_commit_batch();
-    }
-    if constexpr (wg_wait >= 0) { warpgroup_wait<wg_wait>(); }
-    warpgroup_fence_operand(tCrC);
-    if constexpr (Is_RS) { warpgroup_fence_operand(const_cast<Tensor0 &>(tCrA)); }
-}
-
-
-template<typename Layout>
-__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
-    if constexpr (decltype(rank<0>(acc_layout))::value == 3) {  // SM90
-        static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
-        static_assert(decltype(size<0, 1>(acc_layout))::value == 2);
-        static_assert(decltype(rank(acc_layout))::value == 3);
-        auto l = acc_layout;
-        return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<0, 2>(l), get<2>(l)));
-    } else {  // SM80
-        static_assert(decltype(size<0>(acc_layout))::value == 4);
-        static_assert(decltype(rank(acc_layout))::value == 3);
-        auto l = logical_divide(acc_layout, Shape<_2>{});  // ((2, 2), MMA_M, MMA_N)
-        return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));
+      auto l = logical_divide(
+          acc_layout, Shape<X, X, _2>{});  // (4, MMA_M, (2, MMA_N / 2)))
+      return make_layout(
+          make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
     }
+  }
 };
 
-template<typename T, typename ReductionOp, int thread_group_width = 32>
-__inline__ __device__ T WarpAllReduce(T val) {
-    ReductionOp op;
-    #pragma unroll
-    for (int mask = thread_group_width / 2; mask > 0; mask /= 2) {
-        val = op(val, __shfl_xor_sync(0xffffffff, val, mask));
+template <bool zero_init = false,
+          int wg_wait = 0,
+          bool arrive = true,
+          bool commit = true,
+          typename Tensor0,
+          typename Tensor1,
+          typename Tensor2,
+          typename TiledMma>
+__forceinline__ __device__ void gemm(TiledMma &tiled_mma,
+                                     Tensor0 const &tCrA,
+                                     Tensor1 const &tCrB,
+                                     Tensor2 &tCrC) {
+  constexpr bool Is_RS = !cute::is_base_of<cute::GMMA::DescriptorIterator,
+                                           typename TiledMma::FrgTypeA>::value;
+  if constexpr (Is_RS) {
+    warpgroup_fence_operand(const_cast<Tensor0 &>(tCrA));
+  }
+  warpgroup_fence_operand(tCrC);
+  if constexpr (arrive) {
+    warpgroup_arrive();
+  }
+  if constexpr (zero_init) {
+    tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
+    CUTLASS_PRAGMA_UNROLL
+    for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+      cute::gemm(tiled_mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC);
+      tiled_mma.accumulate_ = GMMA::ScaleOut::One;
     }
-    return val;
+  } else {
+    CUTLASS_PRAGMA_UNROLL
+    for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+      cute::gemm(tiled_mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC);
+      tiled_mma.accumulate_ = GMMA::ScaleOut::One;
+    }
+  }
+  if constexpr (commit) {
+    warpgroup_commit_batch();
+  }
+  if constexpr (wg_wait >= 0) {
+    warpgroup_wait<wg_wait>();
+  }
+  warpgroup_fence_operand(tCrC);
+  if constexpr (Is_RS) {
+    warpgroup_fence_operand(const_cast<Tensor0 &>(tCrA));
+  }
 }
 
+template <typename Layout>
+__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
+  if constexpr (decltype(rank<0>(acc_layout))::value == 3) {  // SM90
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(size<0, 1>(acc_layout))::value == 2);
+    static_assert(decltype(rank(acc_layout))::value == 3);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0, 1>(l), get<1>(l)),
+                       make_layout(get<0, 0>(l), get<0, 2>(l), get<2>(l)));
+  } else {  // SM80
+    static_assert(decltype(size<0>(acc_layout))::value == 4);
+    static_assert(decltype(rank(acc_layout))::value == 3);
+    auto l = logical_divide(acc_layout, Shape<_2>{});  // ((2, 2), MMA_M, MMA_N)
+    return make_layout(make_layout(get<0, 1>(l), get<1>(l)),
+                       make_layout(get<0, 0>(l), get<2>(l)));
+  }
+};
+
+template <typename T, typename ReductionOp, int thread_group_width = 32>
+__inline__ __device__ T WarpAllReduce(T val) {
+  ReductionOp op;
+#pragma unroll
+  for (int mask = thread_group_width / 2; mask > 0; mask /= 2) {
+    val = op(val, __shfl_xor_sync(0xffffffff, val, mask));
+  }
+  return val;
+}
 
 template <int kPackSize, int knthreads>
-__device__ inline int get_data_count(const float * src, const float limit_value) {
-    int count = 0;
-    #pragma unroll
-    for (int i = 0; i < kPackSize; i++) {
-        if (src[i] >= limit_value) {
-            count++;
-        }
+__device__ inline int get_data_count(const float *src,
+                                     const float limit_value) {
+  int count = 0;
+#pragma unroll
+  for (int i = 0; i < kPackSize; i++) {
+    if (src[i] >= limit_value) {
+      count++;
     }
-    count = BlockAllReduce<int, SumOp<int>, knthreads>(count);
-    return count;
+  }
+  count = BlockAllReduce<int, SumOp<int>, knthreads>(count);
+  return count;
 }
diff --git a/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_attn.cu b/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_attn.cu
index 260e538602..0a6bf680e0 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_attn.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_attn.cu
@@ -17,786 +17,969 @@
 #include "moba_decoder_attn_kernel.h"
 #include "moba_attn/moba_attn.h"
 
-
-template<bool Is_first, int kMiLen, typename Tensor0, typename Tensor1, typename T>
-inline __device__ void softmax_rescale_o(Tensor0 &scores, Tensor1 &acc_o, const T *scores_max, const T *scores_max_prev, T * scores_sum, const float softmax_scale) {
-    if (Is_first) {
-        scale_apply_exp2<kMiLen>(scores, scores_max, scores_sum, softmax_scale);
-    } else {
-        Tensor acc_o_rowcol = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
-        #pragma unroll
-        for (int mi = 0; mi < kMiLen; ++mi) {
-            const float scores_scale = expf((scores_max_prev[mi] - scores_max[mi]) * softmax_scale);
-            scores_sum[mi] *= scores_scale;
-            #pragma unroll
-            for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
-                acc_o_rowcol(mi, ni) *= scores_scale;
-            }
-        }
-        scale_apply_exp2<kMiLen>(scores, scores_max, scores_sum, softmax_scale);
+template <bool Is_first,
+          int kMiLen,
+          typename Tensor0,
+          typename Tensor1,
+          typename T>
+inline __device__ void softmax_rescale_o(Tensor0& scores,
+                                         Tensor1& acc_o,
+                                         const T* scores_max,
+                                         const T* scores_max_prev,
+                                         T* scores_sum,
+                                         const float softmax_scale) {
+  if (Is_first) {
+    scale_apply_exp2<kMiLen>(scores, scores_max, scores_sum, softmax_scale);
+  } else {
+    Tensor acc_o_rowcol =
+        make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+#pragma unroll
+    for (int mi = 0; mi < kMiLen; ++mi) {
+      const float scores_scale =
+          expf((scores_max_prev[mi] - scores_max[mi]) * softmax_scale);
+      scores_sum[mi] *= scores_scale;
+#pragma unroll
+      for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
+        acc_o_rowcol(mi, ni) *= scores_scale;
+      }
     }
+    scale_apply_exp2<kMiLen>(scores, scores_max, scores_sum, softmax_scale);
+  }
 };
 
-template<typename Kernel_traits, typename ParamType>
-__global__ __launch_bounds__(Kernel_traits::kNThreads) void moba_decoder_attention_kernel(ParamType params) {
-    using cuteType = typename Kernel_traits::cuteType;
-    using ElementAccum = typename Kernel_traits::ElementAccum;
-    using CacheKV_traits = typename Kernel_traits::CacheKV_traits;
-    constexpr int32_t kHeadDim = Kernel_traits::kHeadDim;
-    constexpr int32_t kHeadDimKV = Kernel_traits::kHeadDimKV;
-    constexpr int32_t kBlockM = Kernel_traits::kBlockM;
-    constexpr int32_t kBlockSize = Kernel_traits::kBlockSize;
-    constexpr int32_t kGqaGroupSize = Kernel_traits::kGqaGroupSize;
-    constexpr int32_t kNWarps = Kernel_traits::kNWarps;
-    constexpr int32_t kTileN = Kernel_traits::kTileN;
-    constexpr int32_t kBlockN = kTileN * kBlockSize;
-    constexpr int32_t kDataBits = Kernel_traits::kDataBits;
-    constexpr int32_t kMiLen = (kGqaGroupSize + 7) / 8;
+template <typename Kernel_traits, typename ParamType>
+__global__
+__launch_bounds__(Kernel_traits::kNThreads) void moba_decoder_attention_kernel(
+    ParamType params) {
+  using cuteType = typename Kernel_traits::cuteType;
+  using ElementAccum = typename Kernel_traits::ElementAccum;
+  using CacheKV_traits = typename Kernel_traits::CacheKV_traits;
+  constexpr int32_t kHeadDim = Kernel_traits::kHeadDim;
+  constexpr int32_t kHeadDimKV = Kernel_traits::kHeadDimKV;
+  constexpr int32_t kBlockM = Kernel_traits::kBlockM;
+  constexpr int32_t kBlockSize = Kernel_traits::kBlockSize;
+  constexpr int32_t kGqaGroupSize = Kernel_traits::kGqaGroupSize;
+  constexpr int32_t kNWarps = Kernel_traits::kNWarps;
+  constexpr int32_t kTileN = Kernel_traits::kTileN;
+  constexpr int32_t kBlockN = kTileN * kBlockSize;
+  constexpr int32_t kDataBits = Kernel_traits::kDataBits;
+  constexpr int32_t kMiLen = (kGqaGroupSize + 7) / 8;
 
-    const int32_t bi = blockIdx.y;
-    const int32_t tidx = threadIdx.x;
-    const int32_t partition_idx = blockIdx.x;
-    const int32_t kv_head_idx = blockIdx.z;
-    const int32_t q_head_idx = kv_head_idx * kGqaGroupSize;
+  const int32_t bi = blockIdx.y;
+  const int32_t tidx = threadIdx.x;
+  const int32_t partition_idx = blockIdx.x;
+  const int32_t kv_head_idx = blockIdx.z;
+  const int32_t q_head_idx = kv_head_idx * kGqaGroupSize;
 
-    const int32_t seq_len = params.seq_lens_decoder[bi] == 0 ? 0 : params.seq_lens_decoder[bi] + 1;
+  const int32_t seq_len =
+      params.seq_lens_decoder[bi] == 0 ? 0 : params.seq_lens_decoder[bi] + 1;
 
-    const int32_t head_num = params.head_num;
-    const int32_t kv_head_num = params.kv_head_num;
+  const int32_t head_num = params.head_num;
+  const int32_t kv_head_num = params.kv_head_num;
 
-    const int32_t partition_num = (seq_len + kBlockN - 1) / kBlockN;
+  const int32_t partition_num = (seq_len + kBlockN - 1) / kBlockN;
 
-    if (seq_len == 0 || partition_idx >= partition_num) {
-        return;
+  if (seq_len == 0 || partition_idx >= partition_num) {
+    return;
+  }
+
+  if (seq_len >= params.use_moba_seq_limit &&
+      params.qk_gate_topk_idx_ptr[(bi * kv_head_num + kv_head_idx) *
+                                      Kernel_traits::kMaxN +
+                                  partition_idx] == 0) {
+    return;
+  }
+
+  const int q_bias_offset = q_head_idx * kHeadDim;
+
+  cuteType* q_input = reinterpret_cast<cuteType*>(params.q_input) +
+                      params.cu_seq_q[bi] * head_num * kHeadDim;
+
+  Tensor gQ = make_tensor(
+      make_gmem_ptr(reinterpret_cast<const cuteType*>(q_input) + q_bias_offset),
+      Shape<Int<kBlockM>, Int<kHeadDim>>{},
+      Stride<Int<kHeadDim>, _1>{});
+
+  const int32_t block_idx = partition_idx * kTileN;
+  const int* block_table =
+      params.block_table + bi * params.max_num_blocks_per_seq + block_idx;
+  const int32_t physical_block_number = block_table[0];
+
+  const int32_t cache_offset =
+      (physical_block_number * kv_head_num + kv_head_idx) * kBlockSize *
+      kHeadDimKV;
+
+  Tensor gK = make_tensor(
+      make_gmem_ptr(reinterpret_cast<const cuteType*>(params.cache_k) +
+                    cache_offset),
+      Shape<Int<kBlockSize>, Int<kHeadDimKV>>{},
+      Stride<Int<kHeadDimKV>, _1>{});
+
+  Tensor gV = make_tensor(
+      make_gmem_ptr(reinterpret_cast<const cuteType*>(params.cache_v) +
+                    cache_offset),
+      Shape<Int<kBlockSize>, Int<kHeadDimKV>>{},
+      Stride<Int<kHeadDimKV>, _1>{});
+
+  extern __shared__ char smem_[];
+  Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast<cuteType*>(smem_)),
+                          typename Kernel_traits::SmemLayoutQ{});
+  Tensor sQK =
+      make_tensor(sQ.data() + size(sQ), typename Kernel_traits::SmemLayoutQK{});
+
+  Tensor sK = make_tensor(sQK.data() + size(sQK),
+                          typename Kernel_traits::SmemLayoutKV{});
+  Tensor sV =
+      make_tensor(sK.data() + size(sK), typename Kernel_traits::SmemLayoutKV{});
+  Tensor sVt =
+      make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed{});
+  Tensor sVtNoSwizzle = make_tensor(
+      sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{});
+  __shared__ ElementAccum scores_warp[kNWarps][kMiLen * kBlockM];
+
+  auto gmem_tiled_copy_Q = typename Kernel_traits::GmemTiledCopyQ{};
+  auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx);
+
+  auto gmem_tiled_copy_KV = typename Kernel_traits::GmemTiledCopyKV{};
+  auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx);
+
+  Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ);
+  Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ);
+
+  Tensor tKgK = gmem_thr_copy_KV.partition_S(gK);
+  Tensor tKsK = gmem_thr_copy_KV.partition_D(sK);
+  Tensor tVgV = gmem_thr_copy_KV.partition_S(gV);
+  Tensor tVsV = gmem_thr_copy_KV.partition_D(sV);
+
+  Tensor cQ = make_identity_tensor(make_shape(kBlockM, kHeadDim));
+  Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ);
+
+  Tensor cKV = make_identity_tensor(make_shape(kBlockSize, kHeadDim));
+  Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV);
+
+  typename Kernel_traits::TiledMma tiled_mma;
+
+  auto thr_mma = tiled_mma.get_thread_slice(tidx);
+  using SmemCopyAtom = typename Kernel_traits::SmemCopyAtom;
+  auto smem_tiled_copy_Q = make_tiled_copy_A(SmemCopyAtom{}, tiled_mma);
+  auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx);
+  auto smem_tiled_copy_K = make_tiled_copy_B(SmemCopyAtom{}, tiled_mma);
+  auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx);
+  auto smem_tiled_copy_V = make_tiled_copy_B(
+      typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma);
+  auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx);
+  auto smem_tiled_copy_O =
+      make_tiled_copy_C(typename Kernel_traits::SmemCopyAtomO{}, tiled_mma);
+  auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(tidx);
+
+  Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ);
+  Tensor tSrQ = thr_mma.partition_fragment_A(sQ);
+
+  Tensor tSsQK = smem_thr_copy_Q.partition_S(sQK);
+  Tensor tSrQK = thr_mma.partition_fragment_A(sQK);
+
+  Tensor tSsK = smem_thr_copy_K.partition_S(sK);
+  Tensor tSrK = thr_mma.partition_fragment_B(sK);
+  Tensor tOsVt = smem_thr_copy_V.partition_S(sVt);
+  Tensor tOrVt = thr_mma.partition_fragment_B(sVtNoSwizzle);
+
+  copy<false>(gmem_tiled_copy_Q, tQgQ, tQsQ, tQcQ, kGqaGroupSize);
+
+  cute::cp_async_fence();
+  cp_async_wait<0>();
+
+  const int32_t remain_seq_len = seq_len - partition_idx * kTileN * kBlockSize;
+
+  copy(gmem_tiled_copy_KV, tKgK, tKsK, tKVcKV);
+
+  cute::cp_async_fence();
+
+  const int32_t warp_id = tidx / 32;
+  const int32_t lane_id = tidx % 32;
+  const int32_t row = lane_id / 4;
+  const int32_t col = lane_id % 4;
+  const int row_idx = tidx / 4;
+
+  using scale_k_vec = Vec<cuteType, 32>;
+  using scale_v_vec = Vec<cuteType, 4>;
+
+  scale_k_vec scale_k;
+  scale_k_vec zp_k;
+  scale_v_vec scale_v;
+  scale_v_vec zp_v;
+  if constexpr (kDataBits == 4) {
+    scale_k = *reinterpret_cast<const scale_k_vec*>(
+        params.cache_k_dequant_scale + kv_head_idx * kHeadDim + col * 32);
+    zp_k = *reinterpret_cast<const scale_k_vec*>(
+        params.cache_k_zp + kv_head_idx * kHeadDim + col * 32);
+    scale_v = *reinterpret_cast<const scale_v_vec*>(
+        params.cache_v_dequant_scale + kv_head_idx * kHeadDim + row_idx * 4);
+    zp_v = *reinterpret_cast<const scale_v_vec*>(
+        params.cache_v_zp + kv_head_idx * kHeadDim + row_idx * 4);
+  }
+
+  Tensor acc_o =
+      partition_fragment_C(tiled_mma, Shape<Int<kBlockM>, Int<kHeadDim>>{});
+  clear(acc_o);
+  Tensor acc_s =
+      partition_fragment_C(tiled_mma, Shape<Int<kBlockM>, Int<kBlockSize>>{});
+
+  ElementAccum scores_max[kMiLen];
+  ElementAccum scores_max_prev[kMiLen];
+  ElementAccum scores_sum[kMiLen];
+
+#pragma unroll
+  for (int mi = 0; mi < kMiLen; ++mi) {
+    scores_max[mi] = -INFINITY;
+    scores_sum[mi] = 0;
+  }
+
+  const int cache_offset_step = kv_head_num * kBlockSize * kHeadDimKV;
+
+#pragma unroll
+  for (int n = 0; n < kTileN; ++n) {
+    const int cur_remain_seq_len = remain_seq_len - n * kBlockSize;
+
+    if (cur_remain_seq_len <= 0) {
+      break;
     }
 
-    if (seq_len >= params.use_moba_seq_limit && params.qk_gate_topk_idx_ptr[(bi * kv_head_num + kv_head_idx) * Kernel_traits::kMaxN + partition_idx] == 0) {
-        return;
-    }
-
-
-    const int q_bias_offset = q_head_idx * kHeadDim;
-
-    cuteType * q_input = reinterpret_cast<cuteType *>(params.q_input) + params.cu_seq_q[bi] * head_num * kHeadDim;
-
-    Tensor gQ = make_tensor(
-        make_gmem_ptr(reinterpret_cast<const cuteType *>(q_input) + q_bias_offset),
-        Shape<Int<kBlockM>, Int<kHeadDim>>{},
-        Stride<Int<kHeadDim>, _1>{});
-
-    const int32_t block_idx = partition_idx * kTileN;
-    const int* block_table = params.block_table + bi * params.max_num_blocks_per_seq + block_idx;
-    const int32_t physical_block_number = block_table[0];
-
-    const int32_t cache_offset = (physical_block_number * kv_head_num + kv_head_idx) * kBlockSize * kHeadDimKV;
-
-    Tensor gK = make_tensor(
-        make_gmem_ptr(reinterpret_cast<const cuteType *>(params.cache_k) + cache_offset),
-        Shape<Int<kBlockSize>, Int<kHeadDimKV>>{},
-        Stride<Int<kHeadDimKV>, _1>{});
-
-    Tensor gV = make_tensor(
-        make_gmem_ptr(reinterpret_cast<const cuteType *>(params.cache_v) + cache_offset),
-        Shape<Int<kBlockSize>, Int<kHeadDimKV>>{},
-        Stride<Int<kHeadDimKV>, _1>{});
-
-    extern __shared__ char smem_[];
-    Tensor sQ = make_tensor(
-        make_smem_ptr(reinterpret_cast<cuteType *>(smem_)),
-        typename Kernel_traits::SmemLayoutQ{});
-    Tensor sQK = make_tensor(
-        sQ.data() + size(sQ),
-        typename Kernel_traits::SmemLayoutQK{});
-
-    Tensor sK = make_tensor(sQK.data() + size(sQK), typename Kernel_traits::SmemLayoutKV{});
-    Tensor sV = make_tensor(sK.data() + size(sK), typename Kernel_traits::SmemLayoutKV{});
-    Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed{});
-    Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{});
-    __shared__ ElementAccum scores_warp[kNWarps][kMiLen * kBlockM];
-
-    auto gmem_tiled_copy_Q = typename Kernel_traits::GmemTiledCopyQ{};
-    auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx);
-
-    auto gmem_tiled_copy_KV = typename Kernel_traits::GmemTiledCopyKV{};
-    auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx);
-
-    Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ);
-    Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ);
-
-    Tensor tKgK = gmem_thr_copy_KV.partition_S(gK);
-    Tensor tKsK = gmem_thr_copy_KV.partition_D(sK);
-    Tensor tVgV = gmem_thr_copy_KV.partition_S(gV);
-    Tensor tVsV = gmem_thr_copy_KV.partition_D(sV);
-
-    Tensor cQ = make_identity_tensor(make_shape(kBlockM, kHeadDim));
-    Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ);
-
-    Tensor cKV = make_identity_tensor(make_shape(kBlockSize, kHeadDim));
-    Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV);
-
-    typename Kernel_traits::TiledMma tiled_mma;
-
-    auto thr_mma = tiled_mma.get_thread_slice(tidx);
-    using SmemCopyAtom = typename Kernel_traits::SmemCopyAtom;
-    auto smem_tiled_copy_Q = make_tiled_copy_A(SmemCopyAtom{}, tiled_mma);
-    auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx);
-    auto smem_tiled_copy_K = make_tiled_copy_B(SmemCopyAtom{}, tiled_mma);
-    auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx);
-    auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma);
-    auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx);
-    auto smem_tiled_copy_O = make_tiled_copy_C(typename Kernel_traits::SmemCopyAtomO{}, tiled_mma);
-    auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(tidx);
-
-    Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ);
-    Tensor tSrQ  = thr_mma.partition_fragment_A(sQ);
-
-    Tensor tSsQK = smem_thr_copy_Q.partition_S(sQK);
-    Tensor tSrQK = thr_mma.partition_fragment_A(sQK);
-
-    Tensor tSsK = smem_thr_copy_K.partition_S(sK);
-    Tensor tSrK  = thr_mma.partition_fragment_B(sK);
-    Tensor tOsVt = smem_thr_copy_V.partition_S(sVt);
-    Tensor tOrVt  = thr_mma.partition_fragment_B(sVtNoSwizzle);
-
-    copy<false>(gmem_tiled_copy_Q, tQgQ, tQsQ, tQcQ, kGqaGroupSize);
-
-
-    cute::cp_async_fence();
+    clear(acc_s);
     cp_async_wait<0>();
+    __syncthreads();
 
-    const int32_t remain_seq_len = seq_len - partition_idx * kTileN * kBlockSize;
+    if (n > 0) {
+      tVgV.data() = tVgV.data() +
+                    (block_table[n] - block_table[n - 1]) * cache_offset_step;
+    }
 
-    copy(gmem_tiled_copy_KV, tKgK, tKsK, tKVcKV);
+    copy(gmem_tiled_copy_KV, tVgV, tVsV, tKVcKV);
 
     cute::cp_async_fence();
 
-    const int32_t warp_id = tidx / 32;
-    const int32_t lane_id = tidx % 32;
-    const int32_t row = lane_id / 4;
-    const int32_t col = lane_id % 4;
-    const int row_idx = tidx / 4;
+    if constexpr (kDataBits == 16) {
+      if (n == 0) {
+        gemm(acc_s,
+             tSrQ,
+             tSrK,
+             tSsQ,
+             tSsK,
+             tiled_mma,
+             smem_thr_copy_Q,
+             smem_thr_copy_K,
+             smem_tiled_copy_Q,
+             smem_tiled_copy_K);
+      } else {
+        gemm<true>(acc_s,
+                   tSrQ,
+                   tSrK,
+                   tSsQ,
+                   tSsK,
+                   tiled_mma,
+                   smem_thr_copy_Q,
+                   smem_thr_copy_K,
+                   smem_tiled_copy_Q,
+                   smem_tiled_copy_K);
+      }
+    } else {
+      Tensor tSrKQuant = make_tensor<cuteType>(
+          Layout<Shape<Shape<_2, _2>, Int<kBlockSize / 32>>,
+                 Stride<Shape<_1, _2>, _4>>{});
+      if (n == 0) {
+        gemm_qk_quant<CacheKV_traits, cuteType, kHeadDim, kDataBits>(
+            acc_s,
+            tSrQ,
+            tSsQ,
+            tSrKQuant,
+            sK,
+            tiled_mma,
+            smem_thr_copy_Q,
+            smem_tiled_copy_Q,
+            tidx,
+            scale_k.data.elt,
+            zp_k.data.elt);
+      } else {
+        gemm_qk_quant<CacheKV_traits, cuteType, kHeadDim, kDataBits, true>(
+            acc_s,
+            tSrQ,
+            tSsQ,
+            tSrKQuant,
+            sK,
+            tiled_mma,
+            smem_thr_copy_Q,
+            smem_tiled_copy_Q,
+            tidx,
+            scale_k.data.elt,
+            zp_k.data.elt);
+      }
+    }
+    Tensor scores =
+        make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
 
-    using scale_k_vec = Vec<cuteType, 32>;
-    using scale_v_vec = Vec<cuteType, 4>;
-
-    scale_k_vec scale_k;
-    scale_k_vec zp_k;
-    scale_v_vec scale_v;
-    scale_v_vec zp_v;
-    if constexpr (kDataBits == 4) {
-        scale_k = *reinterpret_cast<const scale_k_vec*>(params.cache_k_dequant_scale + kv_head_idx * kHeadDim + col * 32);
-        zp_k = *reinterpret_cast<const scale_k_vec*>(params.cache_k_zp + kv_head_idx * kHeadDim + col * 32);
-        scale_v = *reinterpret_cast<const scale_v_vec*>(params.cache_v_dequant_scale + kv_head_idx * kHeadDim + row_idx * 4);
-        zp_v = *reinterpret_cast<const scale_v_vec*>(params.cache_v_zp + kv_head_idx * kHeadDim + row_idx * 4);
+    if (partition_idx == partition_num - 1 && cur_remain_seq_len < kBlockSize) {
+      apply_mask<kMiLen>(scores, warp_id, col, cur_remain_seq_len);
     }
 
-    Tensor acc_o = partition_fragment_C(tiled_mma, Shape<Int<kBlockM>, Int<kHeadDim>>{});
-    clear(acc_o);
-    Tensor acc_s = partition_fragment_C(tiled_mma, Shape<Int<kBlockM>, Int<kBlockSize>>{});
-
-    ElementAccum scores_max[kMiLen];
-    ElementAccum scores_max_prev[kMiLen];
-    ElementAccum scores_sum[kMiLen];
-
-    #pragma unroll
+#pragma unroll
     for (int mi = 0; mi < kMiLen; ++mi) {
-        scores_max[mi] = -INFINITY;
-        scores_sum[mi] = 0;
+      scores_max_prev[mi] = scores_max[mi];
     }
 
-    const int cache_offset_step = kv_head_num * kBlockSize * kHeadDimKV;
-
-    #pragma unroll
-    for (int n = 0; n < kTileN; ++n) {
-        const int cur_remain_seq_len = remain_seq_len - n * kBlockSize;
-
-        if (cur_remain_seq_len <= 0) {
-            break;
-        }
-
-        clear(acc_s);
-        cp_async_wait<0>();
-        __syncthreads();
-
-        if (n > 0) {
-            tVgV.data() = tVgV.data() + (block_table[n] - block_table[n - 1]) * cache_offset_step;
-        }
-
-        copy(gmem_tiled_copy_KV, tVgV, tVsV, tKVcKV);
-
-        cute::cp_async_fence();
-
-        if constexpr (kDataBits == 16) {
-            if (n == 0) {
-                gemm(acc_s, tSrQ, tSrK, tSsQ, tSsK, tiled_mma, smem_thr_copy_Q, smem_thr_copy_K, smem_tiled_copy_Q, smem_tiled_copy_K);
-            } else {
-                gemm<true>(acc_s, tSrQ, tSrK, tSsQ, tSsK, tiled_mma, smem_thr_copy_Q, smem_thr_copy_K, smem_tiled_copy_Q, smem_tiled_copy_K);
-            }
-        } else {
-            Tensor tSrKQuant = make_tensor<cuteType>(
-                Layout<
-                    Shape<Shape<_2, _2>, Int<kBlockSize / 32>>,
-                    Stride<Shape<_1, _2>, _4>>{});
-            if (n == 0) {
-                gemm_qk_quant<CacheKV_traits, cuteType, kHeadDim, kDataBits>(acc_s, tSrQ, tSsQ, tSrKQuant, sK, tiled_mma, smem_thr_copy_Q, smem_tiled_copy_Q, tidx, scale_k.data.elt, zp_k.data.elt);
-            } else {
-                gemm_qk_quant<CacheKV_traits, cuteType, kHeadDim, kDataBits, true>(acc_s, tSrQ, tSsQ, tSrKQuant, sK, tiled_mma, smem_thr_copy_Q, smem_tiled_copy_Q, tidx, scale_k.data.elt, zp_k.data.elt);
-            }
-        }
-        Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
-
-        if (partition_idx == partition_num - 1 && cur_remain_seq_len < kBlockSize) {
-            apply_mask<kMiLen>(scores, warp_id, col, cur_remain_seq_len);
-        }
-
-        #pragma unroll
-        for (int mi = 0; mi < kMiLen; ++mi) {
-            scores_max_prev[mi] = scores_max[mi];
-        }
-
-        reduce_max<kMiLen>(scores, scores_max);
-
-        if (col == 0) {
-            scores_warp[warp_id][row] = scores_max[0];
-            if constexpr (kMiLen > 1) {
-                scores_warp[warp_id][row + 8] = scores_max[1];
-            }
-        }
-
-        __syncthreads();
-
-        MaxOp<ElementAccum> max_op;
-
-        if (tidx < kGqaGroupSize) {
-            float cur_max = scores_warp[0][tidx];
-            #pragma unroll
-            for (uint32_t i = 1; i < kNWarps; ++i) {
-                cur_max = max_op(scores_warp[i][tidx], cur_max);
-            }
-            scores_warp[0][tidx] = cur_max;
-        }
-
-        cp_async_wait<0>();
-        __syncthreads();
-
-        if (cur_remain_seq_len > kBlockSize && n < kTileN - 1) {
-            tKgK.data() = tKgK.data() + (block_table[n + 1] - block_table[n]) * cache_offset_step;
-            copy(gmem_tiled_copy_KV, tKgK, tKsK, tKVcKV);
-            cute::cp_async_fence();
-        }
-
-        #pragma unroll
-        for (int mi = 0; mi < kMiLen; ++mi) {
-            scores_max[mi] = scores_warp[0][row + mi * 8];
-        }
-
-        if (n == 0) {
-            softmax_rescale_o<true, kMiLen>(scores, acc_o, scores_max, scores_max_prev, scores_sum, params.inv_sqrt_dh);
-        } else {
-            softmax_rescale_o<false, kMiLen>(scores, acc_o, scores_max, scores_max_prev, scores_sum, params.inv_sqrt_dh);
-        }
-
-        Tensor rS = convert_type<cuteType>(acc_s);
-
-        Tensor trQK = smem_thr_copy_O.retile_S(rS);
-        Tensor tsQK = smem_thr_copy_O.partition_D(sQK);
-        cute::copy(smem_tiled_copy_O, trQK, tsQK);
-
-        __syncthreads();
-
-        if constexpr (kDataBits == 16) {
-            gemm(acc_o, tSrQK, tOrVt, tSsQK, tOsVt, tiled_mma, smem_thr_copy_Q, smem_thr_copy_V, smem_tiled_copy_Q, smem_tiled_copy_V);
-        } else {
-            Tensor tSrVQuant = make_tensor<cuteType>(
-                Layout<
-                    Shape<_4, Shape<_2, _2>>,
-                    Stride<_1, Shape<_4, _8>>>{});
-            gemm_value_quant<CacheKV_traits, cuteType, kHeadDim, kDataBits>(acc_o, tSrQK, tSsQK, tSrVQuant, sV, tiled_mma, smem_thr_copy_Q, smem_tiled_copy_Q, tidx, scale_v.data.elt, zp_v.data.elt);
-        }
-    }
-
-    const uint32_t pack_max_partition_num = (params.max_num_partitions + 3) / 4 * 4;
-    uint32_t max_sum_offset = bi * pack_max_partition_num * head_num + (tidx + q_head_idx) * pack_max_partition_num + partition_idx;
-
-    if (tidx < kGqaGroupSize) {
-        params.maxs[max_sum_offset] = scores_warp[0][tidx] * params.inv_sqrt_dh;
-    }
-
-    SumOp<ElementAccum> sum_op;
-    #pragma unroll
-    for (int mi = 0; mi < kMiLen; ++mi) {
-        scores_sum[mi] = Allreduce<4>::run(scores_sum[mi], sum_op);
-    }
-    __syncthreads();
+    reduce_max<kMiLen>(scores, scores_max);
 
     if (col == 0) {
-        scores_warp[warp_id][row] = scores_sum[0];
-        if constexpr (kMiLen > 1) {
-            scores_warp[warp_id][row + 8] = scores_sum[1];
-        }
+      scores_warp[warp_id][row] = scores_max[0];
+      if constexpr (kMiLen > 1) {
+        scores_warp[warp_id][row + 8] = scores_max[1];
+      }
     }
 
-
-    Tensor rO = convert_type<cuteType>(acc_o);
-    Tensor taccOrO = smem_thr_copy_O.retile_S(rO);
-    Tensor taccOsO = smem_thr_copy_O.partition_D(sQ);
-
-    cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);
-
     __syncthreads();
 
+    MaxOp<ElementAccum> max_op;
+
     if (tidx < kGqaGroupSize) {
-        float cur_sum = scores_warp[0][tidx];
-        #pragma unroll
-        for (uint32_t i = 1; i < kNWarps; ++i) {
-            cur_sum = sum_op(scores_warp[i][tidx], cur_sum);
-        }
-        scores_warp[0][tidx] = cur_sum;
+      float cur_max = scores_warp[0][tidx];
+#pragma unroll
+      for (uint32_t i = 1; i < kNWarps; ++i) {
+        cur_max = max_op(scores_warp[i][tidx], cur_max);
+      }
+      scores_warp[0][tidx] = cur_max;
     }
 
-    Tensor gO = make_tensor(
-        make_gmem_ptr(reinterpret_cast<cuteType *>(params.partition_attn_out) + ((bi * params.max_num_partitions + partition_idx) * head_num  + q_head_idx)* kHeadDim),
-                            Shape<Int<kBlockM>, Int<kHeadDim>>{},
-                            Stride<Int<kHeadDim>, _1>{});
-
-    auto gmem_tiled_copy_O = typename Kernel_traits::GmemTiledCopyO{};
-    auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(tidx);
-    Tensor tOsO = gmem_thr_copy_O.partition_S(sQ);
-    Tensor tOgO = gmem_thr_copy_O.partition_D(gO);
-    constexpr int32_t copy_size = kGqaGroupSize * 16;
+    cp_async_wait<0>();
     __syncthreads();
 
-    if (tidx < copy_size) {
-        cute::copy(gmem_tiled_copy_O, tOsO(_, 0, _), tOgO(_, 0, _));
+    if (cur_remain_seq_len > kBlockSize && n < kTileN - 1) {
+      tKgK.data() = tKgK.data() +
+                    (block_table[n + 1] - block_table[n]) * cache_offset_step;
+      copy(gmem_tiled_copy_KV, tKgK, tKsK, tKVcKV);
+      cute::cp_async_fence();
     }
 
+#pragma unroll
+    for (int mi = 0; mi < kMiLen; ++mi) {
+      scores_max[mi] = scores_warp[0][row + mi * 8];
+    }
+
+    if (n == 0) {
+      softmax_rescale_o<true, kMiLen>(scores,
+                                      acc_o,
+                                      scores_max,
+                                      scores_max_prev,
+                                      scores_sum,
+                                      params.inv_sqrt_dh);
+    } else {
+      softmax_rescale_o<false, kMiLen>(scores,
+                                       acc_o,
+                                       scores_max,
+                                       scores_max_prev,
+                                       scores_sum,
+                                       params.inv_sqrt_dh);
+    }
+
+    Tensor rS = convert_type<cuteType>(acc_s);
+
+    Tensor trQK = smem_thr_copy_O.retile_S(rS);
+    Tensor tsQK = smem_thr_copy_O.partition_D(sQK);
+    cute::copy(smem_tiled_copy_O, trQK, tsQK);
+
+    __syncthreads();
+
+    if constexpr (kDataBits == 16) {
+      gemm(acc_o,
+           tSrQK,
+           tOrVt,
+           tSsQK,
+           tOsVt,
+           tiled_mma,
+           smem_thr_copy_Q,
+           smem_thr_copy_V,
+           smem_tiled_copy_Q,
+           smem_tiled_copy_V);
+    } else {
+      Tensor tSrVQuant = make_tensor<cuteType>(
+          Layout<Shape<_4, Shape<_2, _2>>, Stride<_1, Shape<_4, _8>>>{});
+      gemm_value_quant<CacheKV_traits, cuteType, kHeadDim, kDataBits>(
+          acc_o,
+          tSrQK,
+          tSsQK,
+          tSrVQuant,
+          sV,
+          tiled_mma,
+          smem_thr_copy_Q,
+          smem_tiled_copy_Q,
+          tidx,
+          scale_v.data.elt,
+          zp_v.data.elt);
+    }
+  }
+
+  const uint32_t pack_max_partition_num =
+      (params.max_num_partitions + 3) / 4 * 4;
+  uint32_t max_sum_offset = bi * pack_max_partition_num * head_num +
+                            (tidx + q_head_idx) * pack_max_partition_num +
+                            partition_idx;
+
+  if (tidx < kGqaGroupSize) {
+    params.maxs[max_sum_offset] = scores_warp[0][tidx] * params.inv_sqrt_dh;
+  }
+
+  SumOp<ElementAccum> sum_op;
+#pragma unroll
+  for (int mi = 0; mi < kMiLen; ++mi) {
+    scores_sum[mi] = Allreduce<4>::run(scores_sum[mi], sum_op);
+  }
+  __syncthreads();
+
+  if (col == 0) {
+    scores_warp[warp_id][row] = scores_sum[0];
     if constexpr (kMiLen > 1) {
-        if (tidx < copy_size - 128) {
-            cute::copy(gmem_tiled_copy_O, tOsO(_, 1, _), tOgO(_, 1, _));
-        }
+      scores_warp[warp_id][row + 8] = scores_sum[1];
     }
+  }
 
-    if (tidx < kGqaGroupSize) {
-        params.sums[max_sum_offset] = scores_warp[0][tidx];
+  Tensor rO = convert_type<cuteType>(acc_o);
+  Tensor taccOrO = smem_thr_copy_O.retile_S(rO);
+  Tensor taccOsO = smem_thr_copy_O.partition_D(sQ);
+
+  cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);
+
+  __syncthreads();
+
+  if (tidx < kGqaGroupSize) {
+    float cur_sum = scores_warp[0][tidx];
+#pragma unroll
+    for (uint32_t i = 1; i < kNWarps; ++i) {
+      cur_sum = sum_op(scores_warp[i][tidx], cur_sum);
     }
+    scores_warp[0][tidx] = cur_sum;
+  }
+
+  Tensor gO = make_tensor(
+      make_gmem_ptr(
+          reinterpret_cast<cuteType*>(params.partition_attn_out) +
+          ((bi * params.max_num_partitions + partition_idx) * head_num +
+           q_head_idx) *
+              kHeadDim),
+      Shape<Int<kBlockM>, Int<kHeadDim>>{},
+      Stride<Int<kHeadDim>, _1>{});
+
+  auto gmem_tiled_copy_O = typename Kernel_traits::GmemTiledCopyO{};
+  auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(tidx);
+  Tensor tOsO = gmem_thr_copy_O.partition_S(sQ);
+  Tensor tOgO = gmem_thr_copy_O.partition_D(gO);
+  constexpr int32_t copy_size = kGqaGroupSize * 16;
+  __syncthreads();
+
+  if (tidx < copy_size) {
+    cute::copy(gmem_tiled_copy_O, tOsO(_, 0, _), tOgO(_, 0, _));
+  }
+
+  if constexpr (kMiLen > 1) {
+    if (tidx < copy_size - 128) {
+      cute::copy(gmem_tiled_copy_O, tOsO(_, 1, _), tOgO(_, 1, _));
+    }
+  }
+
+  if (tidx < kGqaGroupSize) {
+    params.sums[max_sum_offset] = scores_warp[0][tidx];
+  }
 }
 
+template <typename Kernel_traits, typename ParamType>
+inline __device__ float calculate_logit_scale(const int partition_num,
+                                              const int pack_max_partition_num,
+                                              ParamType& params,
+                                              char* shared_mem,
+                                              const int seq_len,
+                                              const int* qk_gate_topk_idx_ptr) {
+  constexpr int32_t kNFloatPacksize = 16 / sizeof(float);
+  constexpr int32_t kNReduceThreads = Kernel_traits::kNReduceThreads;
+  const int32_t bi = blockIdx.z;
+  const int32_t tidx = threadIdx.x;
+  const int32_t head_idx = blockIdx.y;
+  const int32_t head_num = params.head_num;
 
-template<typename Kernel_traits, typename ParamType>
-inline __device__ float calculate_logit_scale(const int partition_num, const int pack_max_partition_num, ParamType &params, char * shared_mem, const int seq_len, const int *qk_gate_topk_idx_ptr) {
-    constexpr int32_t kNFloatPacksize = 16 / sizeof(float);
-    constexpr int32_t kNReduceThreads = Kernel_traits::kNReduceThreads;
-    const int32_t bi = blockIdx.z;
-    const int32_t tidx = threadIdx.x;
-    const int32_t head_idx = blockIdx.y;
-    const int32_t head_num = params.head_num;
+  using float_vec = Vec<float, kNFloatPacksize>;
+  const int32_t offset = bi * head_num * pack_max_partition_num +
+                         head_idx * pack_max_partition_num;
 
-    using float_vec = Vec<float, kNFloatPacksize>;
-    const int32_t offset = bi * head_num * pack_max_partition_num + head_idx * pack_max_partition_num;
+  float* shared_max_logits = reinterpret_cast<float*>(shared_mem);
+  const float* max_logits_ptr = params.maxs + offset;
+  float global_max_logit = -FLT_MAX;
 
-    float* shared_max_logits = reinterpret_cast<float*>(shared_mem);
-    const float* max_logits_ptr = params.maxs + offset;
-    float global_max_logit = -FLT_MAX;
-
-    int32_t idx = tidx * kNFloatPacksize;
-    #pragma unroll
-    for (; idx <= partition_num - kNFloatPacksize; idx += kNReduceThreads * kNFloatPacksize) {
-        float_vec cur_max = *reinterpret_cast<const float_vec*>(max_logits_ptr + idx);
-        #pragma unroll
-        for (int32_t j = 0; j < kNFloatPacksize; ++j) {
-            if (seq_len >= params.use_moba_seq_limit) {
-                if (qk_gate_topk_idx_ptr[idx + j] != 0) {
-                    global_max_logit = fmaxf(global_max_logit, cur_max.data.elt[j]);
-                }
-            } else {
-                global_max_logit = fmaxf(global_max_logit, cur_max.data.elt[j]);
-            }
+  int32_t idx = tidx * kNFloatPacksize;
+#pragma unroll
+  for (; idx <= partition_num - kNFloatPacksize;
+       idx += kNReduceThreads * kNFloatPacksize) {
+    float_vec cur_max =
+        *reinterpret_cast<const float_vec*>(max_logits_ptr + idx);
+#pragma unroll
+    for (int32_t j = 0; j < kNFloatPacksize; ++j) {
+      if (seq_len >= params.use_moba_seq_limit) {
+        if (qk_gate_topk_idx_ptr[idx + j] != 0) {
+          global_max_logit = fmaxf(global_max_logit, cur_max.data.elt[j]);
         }
-        cur_max.store_to(shared_max_logits + idx);
+      } else {
+        global_max_logit = fmaxf(global_max_logit, cur_max.data.elt[j]);
+      }
     }
+    cur_max.store_to(shared_max_logits + idx);
+  }
 
-    const int32_t packed_data_num = partition_num / kNFloatPacksize * kNFloatPacksize;
+  const int32_t packed_data_num =
+      partition_num / kNFloatPacksize * kNFloatPacksize;
 
-    idx = packed_data_num + tidx;
-    #pragma unroll
-    for (; idx < partition_num; idx += kNReduceThreads) {
-        if (seq_len >= params.use_moba_seq_limit) {
-            if (qk_gate_topk_idx_ptr[idx] != 0) {
-                float cur_max = max_logits_ptr[idx];
-                global_max_logit = fmaxf(global_max_logit, cur_max);
-                shared_max_logits[idx] = cur_max;
-            }
-        } else {
-            float cur_max = max_logits_ptr[idx];
-            global_max_logit = fmaxf(global_max_logit, cur_max);
-            shared_max_logits[idx] = cur_max;
+  idx = packed_data_num + tidx;
+#pragma unroll
+  for (; idx < partition_num; idx += kNReduceThreads) {
+    if (seq_len >= params.use_moba_seq_limit) {
+      if (qk_gate_topk_idx_ptr[idx] != 0) {
+        float cur_max = max_logits_ptr[idx];
+        global_max_logit = fmaxf(global_max_logit, cur_max);
+        shared_max_logits[idx] = cur_max;
+      }
+    } else {
+      float cur_max = max_logits_ptr[idx];
+      global_max_logit = fmaxf(global_max_logit, cur_max);
+      shared_max_logits[idx] = cur_max;
+    }
+  }
+  __syncthreads();
+
+  global_max_logit =
+      BlockAllReduce<float, MaxOp<float>, kNReduceThreads>(global_max_logit);
+
+  float* share_sum_scale = reinterpret_cast<float*>(
+      shared_mem + sizeof(float) * pack_max_partition_num);
+  const float* exp_sums_ptr = params.sums + offset;
+  float global_exp_sum = 0.0f;
+
+  idx = tidx * kNFloatPacksize;
+#pragma unroll
+  for (; idx <= partition_num - kNFloatPacksize;
+       idx += kNReduceThreads * kNFloatPacksize) {
+    float_vec share_max =
+        *reinterpret_cast<const float_vec*>(shared_max_logits + idx);
+#pragma unroll
+    for (int32_t j = 0; j < kNFloatPacksize; ++j) {
+      if (seq_len >= params.use_moba_seq_limit) {
+        if (qk_gate_topk_idx_ptr[idx + j] != 0) {
+          float exp_sub_max = expf(share_max.data.elt[j] - global_max_logit);
+          float rescaled_exp_sum = exp_sums_ptr[idx + j] * exp_sub_max;
+          global_exp_sum += rescaled_exp_sum;
+          share_max.data.elt[j] = exp_sub_max;
         }
+      } else {
+        float exp_sub_max = expf(share_max.data.elt[j] - global_max_logit);
+        float rescaled_exp_sum = exp_sums_ptr[idx + j] * exp_sub_max;
+        global_exp_sum += rescaled_exp_sum;
+        share_max.data.elt[j] = exp_sub_max;
+      }
     }
-    __syncthreads();
+    share_max.store_to(share_sum_scale + idx);
+  }
 
-    global_max_logit = BlockAllReduce<float, MaxOp<float>, kNReduceThreads>(global_max_logit);
-
-    float* share_sum_scale = reinterpret_cast<float*>(shared_mem + sizeof(float) * pack_max_partition_num);
-    const float* exp_sums_ptr = params.sums + offset;
-    float global_exp_sum = 0.0f;
-
-    idx = tidx * kNFloatPacksize;
-    #pragma unroll
-    for (; idx <= partition_num - kNFloatPacksize; idx += kNReduceThreads * kNFloatPacksize) {
-        float_vec share_max = *reinterpret_cast<const float_vec*>(shared_max_logits + idx);
-        #pragma unroll
-        for (int32_t j = 0; j < kNFloatPacksize; ++j) {
-            if (seq_len >= params.use_moba_seq_limit) {
-                if (qk_gate_topk_idx_ptr[idx + j] != 0) {
-                    float exp_sub_max = expf(share_max.data.elt[j] - global_max_logit);
-                    float rescaled_exp_sum = exp_sums_ptr[idx + j] * exp_sub_max;
-                    global_exp_sum += rescaled_exp_sum;
-                    share_max.data.elt[j] = exp_sub_max;
-                }
-            } else {
-                float exp_sub_max = expf(share_max.data.elt[j] - global_max_logit);
-                float rescaled_exp_sum = exp_sums_ptr[idx + j] * exp_sub_max;
-                global_exp_sum += rescaled_exp_sum;
-                share_max.data.elt[j] = exp_sub_max;
-            }
-        }
-        share_max.store_to(share_sum_scale + idx);
+  idx = packed_data_num + tidx;
+#pragma unroll
+  for (; idx < partition_num; idx += kNReduceThreads) {
+    if (seq_len >= params.use_moba_seq_limit) {
+      if (qk_gate_topk_idx_ptr[idx] != 0) {
+        float share_max = shared_max_logits[idx];
+        float exp_sub_max = expf(share_max - global_max_logit);
+        float rescaled_exp_sum = exp_sums_ptr[idx] * exp_sub_max;
+        global_exp_sum += rescaled_exp_sum;
+        share_sum_scale[idx] = exp_sub_max;
+      }
+    } else {
+      float share_max = shared_max_logits[idx];
+      float exp_sub_max = expf(share_max - global_max_logit);
+      float rescaled_exp_sum = exp_sums_ptr[idx] * exp_sub_max;
+      global_exp_sum += rescaled_exp_sum;
+      share_sum_scale[idx] = exp_sub_max;
     }
+  }
+  __syncthreads();
 
-    idx = packed_data_num + tidx;
-    #pragma unroll
-    for (; idx < partition_num; idx += kNReduceThreads) {
-        if (seq_len >= params.use_moba_seq_limit) {
-            if (qk_gate_topk_idx_ptr[idx] != 0) {
-                float share_max = shared_max_logits[idx];
-                float exp_sub_max = expf(share_max - global_max_logit);
-                float rescaled_exp_sum = exp_sums_ptr[idx] * exp_sub_max;
-                global_exp_sum += rescaled_exp_sum;
-                share_sum_scale[idx] = exp_sub_max;
-            }
-        } else {
-            float share_max = shared_max_logits[idx];
-            float exp_sub_max = expf(share_max - global_max_logit);
-            float rescaled_exp_sum = exp_sums_ptr[idx] * exp_sub_max;
-            global_exp_sum += rescaled_exp_sum;
-            share_sum_scale[idx] = exp_sub_max;
-        }
-    }
-    __syncthreads();
+  global_exp_sum =
+      BlockAllReduce<float, SumOp<float>, kNReduceThreads>(global_exp_sum);
 
-    global_exp_sum = BlockAllReduce<float, SumOp<float>, kNReduceThreads>(global_exp_sum);
-
-    const float inv_global_exp_sum = fdividef(1.0f, global_exp_sum + 1e-6f);
-    return inv_global_exp_sum;
+  const float inv_global_exp_sum = fdividef(1.0f, global_exp_sum + 1e-6f);
+  return inv_global_exp_sum;
 }
 
-template<typename Kernel_traits, typename ParamType>
-__global__ void __launch_bounds__(Kernel_traits::kNReduceThreads) moba_decoder_attention_merge_kernel(ParamType params) {
-    using cuteType = typename Kernel_traits::cuteType;
-    constexpr int32_t kBlockN = Kernel_traits::kTileN * Kernel_traits::kBlockSize;
-    constexpr int32_t kNReducePacksize = 16 / sizeof(cuteType);
-    constexpr int32_t kNFloatPacksize = 16 / sizeof(float);
-    constexpr int32_t kNReduceWarps = Kernel_traits::kNReduceWarps;
-    constexpr int32_t kHeadDim = Kernel_traits::kHeadDim;
-    const int32_t bi = blockIdx.z;
-    const int32_t headdim_idx = kNReducePacksize * kNReduceWarps * blockIdx.x;
-    const int32_t tidx = threadIdx.x;
-    const int32_t head_idx = blockIdx.y;
-    const int32_t warp_id = tidx / 32;
-    const int32_t lane_id = tidx % 32;
-    const int32_t seq_len = params.seq_lens_decoder[bi] + 1;
-    const int32_t head_num = params.head_num;
-    using pack_half = typename PackedHalf<cuteType>::Type;
+template <typename Kernel_traits, typename ParamType>
+__global__ void __launch_bounds__(Kernel_traits::kNReduceThreads)
+    moba_decoder_attention_merge_kernel(ParamType params) {
+  using cuteType = typename Kernel_traits::cuteType;
+  constexpr int32_t kBlockN = Kernel_traits::kTileN * Kernel_traits::kBlockSize;
+  constexpr int32_t kNReducePacksize = 16 / sizeof(cuteType);
+  constexpr int32_t kNFloatPacksize = 16 / sizeof(float);
+  constexpr int32_t kNReduceWarps = Kernel_traits::kNReduceWarps;
+  constexpr int32_t kHeadDim = Kernel_traits::kHeadDim;
+  const int32_t bi = blockIdx.z;
+  const int32_t headdim_idx = kNReducePacksize * kNReduceWarps * blockIdx.x;
+  const int32_t tidx = threadIdx.x;
+  const int32_t head_idx = blockIdx.y;
+  const int32_t warp_id = tidx / 32;
+  const int32_t lane_id = tidx % 32;
+  const int32_t seq_len = params.seq_lens_decoder[bi] + 1;
+  const int32_t head_num = params.head_num;
+  using pack_half = typename PackedHalf<cuteType>::Type;
 
+  if (params.seq_lens_decoder[bi] == 0) {
+    return;
+  }
 
-    if (params.seq_lens_decoder[bi] == 0) {
-        return;
+  extern __shared__ char shared_mem[];
+
+  const int32_t partition_num = (seq_len + kBlockN - 1) / kBlockN;
+  const int32_t pack_max_partition_num =
+      (params.max_num_partitions + kNFloatPacksize - 1) / kNFloatPacksize *
+      kNFloatPacksize;
+
+  float* share_sum_scale = reinterpret_cast<float*>(
+      shared_mem + sizeof(float) * pack_max_partition_num);
+
+  constexpr int32_t kGqaGroupSize = Kernel_traits::kGqaGroupSize;
+  const int kv_head_idx = head_idx / Kernel_traits::kGqaGroupSize;
+  const int* qk_gate_topk_idx_ptr =
+      params.qk_gate_topk_idx_ptr +
+      (bi * params.kv_head_num + kv_head_idx) * Kernel_traits::kMaxN;
+
+  float inv_global_exp_sum =
+      calculate_logit_scale<Kernel_traits>(partition_num,
+                                           pack_max_partition_num,
+                                           params,
+                                           shared_mem,
+                                           seq_len,
+                                           qk_gate_topk_idx_ptr);
+
+  using T_vec = Vec<cuteType, kNReducePacksize>;
+
+  cuteType* partition_attn_out =
+      reinterpret_cast<cuteType*>(params.partition_attn_out) +
+      bi * head_num * params.max_num_partitions * kHeadDim +
+      head_idx * kHeadDim + headdim_idx;
+
+  Vec<float, kNReducePacksize> acc;
+  acc.set_zero();
+#pragma unroll
+  for (int idx = lane_id; idx < partition_num; idx += 32) {
+    if (seq_len >= params.use_moba_seq_limit &&
+        qk_gate_topk_idx_ptr[idx] == 0) {
+      continue;
     }
-
-    extern __shared__ char shared_mem[];
-
-    const int32_t partition_num = (seq_len + kBlockN - 1) / kBlockN;
-    const int32_t pack_max_partition_num = (params.max_num_partitions + kNFloatPacksize - 1) / kNFloatPacksize * kNFloatPacksize;
-
-    float* share_sum_scale = reinterpret_cast<float*>(shared_mem + sizeof(float) * pack_max_partition_num);
-
-    constexpr int32_t kGqaGroupSize = Kernel_traits::kGqaGroupSize;
-    const int kv_head_idx = head_idx / Kernel_traits::kGqaGroupSize;
-    const int * qk_gate_topk_idx_ptr = params.qk_gate_topk_idx_ptr + (bi * params.kv_head_num + kv_head_idx) * Kernel_traits::kMaxN;
-
-    float inv_global_exp_sum = calculate_logit_scale<Kernel_traits>(partition_num, pack_max_partition_num, params, shared_mem, seq_len, qk_gate_topk_idx_ptr);
-
-
-    using T_vec = Vec<cuteType, kNReducePacksize>;
-
-    cuteType* partition_attn_out = reinterpret_cast<cuteType*>(params.partition_attn_out) + bi * head_num * params.max_num_partitions * kHeadDim + head_idx * kHeadDim + headdim_idx;
-
-    Vec<float, kNReducePacksize> acc;
-    acc.set_zero();
-    #pragma unroll
-    for (int idx = lane_id; idx < partition_num; idx += 32) {
-        if (seq_len >= params.use_moba_seq_limit && qk_gate_topk_idx_ptr[idx] == 0) {
-            continue;
-        }
-        T_vec sub_logits = *reinterpret_cast<T_vec*>(&partition_attn_out[idx * head_num * kHeadDim + warp_id * kNReducePacksize]);
-        float scale = share_sum_scale[idx];
-        #pragma unroll
-        for (int k = 0; k < kNReducePacksize; ++k) {
-            acc.data.elt[k] += static_cast<float>(sub_logits.data.elt[k]) * scale;
-        }
-    }
-
-    __syncthreads();
-
-    T_vec out;
-    #pragma unroll
+    T_vec sub_logits = *reinterpret_cast<T_vec*>(
+        &partition_attn_out[idx * head_num * kHeadDim +
+                            warp_id * kNReducePacksize]);
+    float scale = share_sum_scale[idx];
+#pragma unroll
     for (int k = 0; k < kNReducePacksize; ++k) {
-        out.data.elt[k] = static_cast<cuteType>(WarpAllReduce<float, SumOp<float>>(acc.data.elt[k]) * inv_global_exp_sum);
+      acc.data.elt[k] += static_cast<float>(sub_logits.data.elt[k]) * scale;
     }
+  }
 
-    const int ori_token_idx = params.cu_seq_q[bi];
-    cuteType * attn_out = reinterpret_cast<cuteType *>(params.attn_out) + ori_token_idx * head_num * kHeadDim + head_idx * kHeadDim + headdim_idx + warp_id * kNReducePacksize;
+  __syncthreads();
 
-    if (lane_id == 0) {
-        out.store_to(attn_out);
-    }
+  T_vec out;
+#pragma unroll
+  for (int k = 0; k < kNReducePacksize; ++k) {
+    out.data.elt[k] = static_cast<cuteType>(
+        WarpAllReduce<float, SumOp<float>>(acc.data.elt[k]) *
+        inv_global_exp_sum);
+  }
+
+  const int ori_token_idx = params.cu_seq_q[bi];
+  cuteType* attn_out = reinterpret_cast<cuteType*>(params.attn_out) +
+                       ori_token_idx * head_num * kHeadDim +
+                       head_idx * kHeadDim + headdim_idx +
+                       warp_id * kNReducePacksize;
+
+  if (lane_id == 0) {
+    out.store_to(attn_out);
+  }
 }
 
+template <typename Kernel_traits, typename ParamType>
+void run_moba_decoder_attn(ParamType& params, cudaStream_t stream) {
+  dim3 grid;
+  grid.x = params.max_num_partitions;
+  grid.y = params.batch_size;
+  grid.z = params.kv_head_num;
+  constexpr int smem_size = Kernel_traits::kShareMemSize;
+  constexpr auto kernel =
+      &moba_decoder_attention_kernel<Kernel_traits, ParamType>;
+  if (smem_size >= 48 * 1024) {
+    cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  }
+  kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
 
-template<typename Kernel_traits, typename ParamType>
-void run_moba_decoder_attn(ParamType &params, cudaStream_t stream) {
-    dim3 grid;
-    grid.x = params.max_num_partitions;
-    grid.y = params.batch_size;
-    grid.z = params.kv_head_num;
-    constexpr int smem_size = Kernel_traits::kShareMemSize;
-    constexpr auto kernel = &moba_decoder_attention_kernel<Kernel_traits, ParamType>;
-    if (smem_size >= 48 * 1024) {
-        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-    }
-    kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
+  int32_t reduce_shared_mem_size =
+      2 * (params.max_num_partitions + 4) * sizeof(float);
+  constexpr int32_t pack_size = 16 / sizeof(typename Kernel_traits::cuteType);
+  static_assert(Kernel_traits::kHeadDim % pack_size == 0);
+  static_assert((Kernel_traits::kHeadDim / Kernel_traits::kNReduceWarps) %
+                    pack_size ==
+                0);
+  grid.x = Kernel_traits::kHeadDim / Kernel_traits::kNReduceWarps / pack_size;
+  grid.y = params.head_num;
+  grid.z = params.batch_size;
+  auto reduce_kernel =
+      &moba_decoder_attention_merge_kernel<Kernel_traits, ParamType>;
 
-    int32_t reduce_shared_mem_size = 2 * (params.max_num_partitions + 4) * sizeof(float);
-    constexpr int32_t pack_size = 16 / sizeof(typename Kernel_traits::cuteType);
-    static_assert(Kernel_traits::kHeadDim % pack_size == 0);
-    static_assert((Kernel_traits::kHeadDim / Kernel_traits::kNReduceWarps) % pack_size == 0);
-    grid.x = Kernel_traits::kHeadDim / Kernel_traits::kNReduceWarps / pack_size;
-    grid.y = params.head_num;
-    grid.z = params.batch_size;
-    auto reduce_kernel = &moba_decoder_attention_merge_kernel<Kernel_traits, ParamType>;
-
-    if (reduce_shared_mem_size >= 48 * 1024) {
-        cudaFuncSetAttribute(
-            reduce_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, reduce_shared_mem_size);
-    }
-    reduce_kernel<<<grid, Kernel_traits::kNReduceThreads, reduce_shared_mem_size, stream>>>(params);
+  if (reduce_shared_mem_size >= 48 * 1024) {
+    cudaFuncSetAttribute(reduce_kernel,
+                         cudaFuncAttributeMaxDynamicSharedMemorySize,
+                         reduce_shared_mem_size);
+  }
+  reduce_kernel<<<grid,
+                  Kernel_traits::kNReduceThreads,
+                  reduce_shared_mem_size,
+                  stream>>>(params);
 }
 
-
-template<typename cute_type, int kCacheBits, int kBlockN, int kMaxN, typename ParamType>
-void run_moba_decoder_attn_hdim128(ParamType &params, cudaStream_t stream) {
-    const int gqaGroupSize = params.head_num / params.kv_head_num;
-    using CacheKVTraits = CacheKV_quant_traits<cute_type, kCacheBits>;
-    constexpr int kTileN = kBlockN / CacheKVTraits::kBlockSize;
-    switch (gqaGroupSize) {
-        case 12: {
-            run_moba_decoder_attn<moba_decoder_attn_kernel_traits<12, kTileN, kMaxN,CacheKVTraits>>(params, stream);
-            break;
-        }
-        case 8: {
-            run_moba_decoder_attn<moba_decoder_attn_kernel_traits<8, kTileN, kMaxN,CacheKVTraits>>(params, stream);
-            break;
-        }
-        case 7: {
-            run_moba_decoder_attn<moba_decoder_attn_kernel_traits<7, kTileN, kMaxN,CacheKVTraits>>(params, stream);
-            break;
-        }
-        case 6: {
-            run_moba_decoder_attn<moba_decoder_attn_kernel_traits<6, kTileN, kMaxN,CacheKVTraits>>(params, stream);
-            break;
-        }
-        case 5: {
-            run_moba_decoder_attn<moba_decoder_attn_kernel_traits<5, kTileN, kMaxN,CacheKVTraits>>(params, stream);
-            break;
-        }
-        case 4: {
-            run_moba_decoder_attn<moba_decoder_attn_kernel_traits<4, kTileN, kMaxN,CacheKVTraits>>(params, stream);
-            break;
-        }
-        default: {
-            PADDLE_THROW(phi::errors::Unimplemented(
-            "DecoderBlockAttention not implemented for gqaGroupSize = %d", gqaGroupSize));
-        }
+template <typename cute_type,
+          int kCacheBits,
+          int kBlockN,
+          int kMaxN,
+          typename ParamType>
+void run_moba_decoder_attn_hdim128(ParamType& params, cudaStream_t stream) {
+  const int gqaGroupSize = params.head_num / params.kv_head_num;
+  using CacheKVTraits = CacheKV_quant_traits<cute_type, kCacheBits>;
+  constexpr int kTileN = kBlockN / CacheKVTraits::kBlockSize;
+  switch (gqaGroupSize) {
+    case 12: {
+      run_moba_decoder_attn<
+          moba_decoder_attn_kernel_traits<12, kTileN, kMaxN, CacheKVTraits>>(
+          params, stream);
+      break;
     }
+    case 8: {
+      run_moba_decoder_attn<
+          moba_decoder_attn_kernel_traits<8, kTileN, kMaxN, CacheKVTraits>>(
+          params, stream);
+      break;
+    }
+    case 7: {
+      run_moba_decoder_attn<
+          moba_decoder_attn_kernel_traits<7, kTileN, kMaxN, CacheKVTraits>>(
+          params, stream);
+      break;
+    }
+    case 6: {
+      run_moba_decoder_attn<
+          moba_decoder_attn_kernel_traits<6, kTileN, kMaxN, CacheKVTraits>>(
+          params, stream);
+      break;
+    }
+    case 5: {
+      run_moba_decoder_attn<
+          moba_decoder_attn_kernel_traits<5, kTileN, kMaxN, CacheKVTraits>>(
+          params, stream);
+      break;
+    }
+    case 4: {
+      run_moba_decoder_attn<
+          moba_decoder_attn_kernel_traits<4, kTileN, kMaxN, CacheKVTraits>>(
+          params, stream);
+      break;
+    }
+    default: {
+      PADDLE_THROW(phi::errors::Unimplemented(
+          "DecoderBlockAttention not implemented for gqaGroupSize = %d",
+          gqaGroupSize));
+    }
+  }
 }
 
-
 template <typename T>
 void DispatchMobaDecoderAttn(
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cache_k,
-        const paddle::Tensor& cache_v,
-        const paddle::Tensor& block_tables,
-        const paddle::Tensor& k_block_means,
-        const paddle::Tensor& out,
-        const paddle::Tensor& qk_gate_topk_idx,
-        const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_zero_points,
-        const paddle::optional<paddle::Tensor>& cache_v_zero_points,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int batch_size,
-        const int max_input_length,
-        const int use_moba_seq_limit,
-        const std::string &cache_quant_type_str) {
+    const paddle::Tensor& q_input,
+    const paddle::Tensor& seq_len_encoder,
+    const paddle::Tensor& seq_len_decoder,
+    const paddle::Tensor& cu_seq_q,
+    const paddle::Tensor& cache_k,
+    const paddle::Tensor& cache_v,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& k_block_means,
+    const paddle::Tensor& out,
+    const paddle::Tensor& qk_gate_topk_idx,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zero_points,
+    const paddle::optional<paddle::Tensor>& cache_v_zero_points,
+    const int head_num,
+    const int kv_head_num,
+    const int head_dim,
+    const int max_seq_q,
+    const int max_seq_k,
+    const int batch_size,
+    const int max_input_length,
+    const int use_moba_seq_limit,
+    const std::string& cache_quant_type_str) {
+  using cute_type = typename cuteType<T>::type;
+  const int kMobaBlockSize = 128;
+  const int kMaxN = 1024;
 
-    using cute_type = typename cuteType<T>::type;
-    const int kMobaBlockSize = 128;
-    const int kMaxN = 1024;
+  constexpr int max_seq_per_block = kMobaBlockSize;
+  moba_decoder_attn_params<cute_type> params;
+  memset(&params, 0, sizeof(params));
+  const uint32_t max_num_partitions =
+      (max_seq_k + max_seq_per_block) / max_seq_per_block;
+  assert(head_dim == 128);
 
-    constexpr int max_seq_per_block = kMobaBlockSize;
-    moba_decoder_attn_params<cute_type> params;
-    memset(&params, 0, sizeof(params));
-    const uint32_t max_num_partitions = (max_seq_k + max_seq_per_block) / max_seq_per_block;
-    assert(head_dim == 128);
+  paddle::Tensor maxs =
+      paddle::empty({batch_size, head_num, (max_num_partitions + 3) / 4 * 4},
+                    paddle::DataType::FLOAT32,
+                    q_input.place());
+  paddle::Tensor sums =
+      paddle::empty({batch_size, head_num, (max_num_partitions + 3) / 4 * 4},
+                    paddle::DataType::FLOAT32,
+                    q_input.place());
+  paddle::Tensor partition_attn_out =
+      paddle::empty({batch_size, max_num_partitions, head_num, head_dim},
+                    q_input.dtype(),
+                    q_input.place());
 
-    paddle::Tensor maxs = paddle::empty({batch_size, head_num, (max_num_partitions + 3) / 4 * 4}, paddle::DataType::FLOAT32, q_input.place());
-    paddle::Tensor sums = paddle::empty({batch_size, head_num, (max_num_partitions + 3) / 4 * 4}, paddle::DataType::FLOAT32, q_input.place());
-    paddle::Tensor partition_attn_out = paddle::empty({batch_size, max_num_partitions, head_num, head_dim}, q_input.dtype(), q_input.place());
+  params.q_input =
+      reinterpret_cast<cute_type*>(const_cast<T*>(q_input.data<T>()));
+  params.attn_out = reinterpret_cast<cute_type*>(const_cast<T*>(out.data<T>()));
+  params.seq_lens_encoder = const_cast<int*>(seq_len_encoder.data<int>());
+  params.seq_lens_decoder = const_cast<int*>(seq_len_decoder.data<int>());
+  params.block_table = const_cast<int*>(block_tables.data<int>());
+  params.max_input_length = max_input_length;
+  params.head_num = head_num;
+  params.kv_head_num = kv_head_num;
+  params.max_num_blocks_per_seq = block_tables.dims()[1];
+  params.batch_size = batch_size;
+  params.inv_sqrt_dh = 1.0f / std::sqrt(head_dim);
+  params.max_num_partitions = max_num_partitions;
+  params.maxs = reinterpret_cast<float*>(maxs.data<float>());
+  params.sums = reinterpret_cast<float*>(sums.data<float>());
+  params.partition_attn_out =
+      reinterpret_cast<cute_type*>(partition_attn_out.data<T>());
+  params.qk_gate_topk_idx_ptr = const_cast<int*>(qk_gate_topk_idx.data<int>());
+  params.use_moba_seq_limit = use_moba_seq_limit;
+  params.cu_seq_q = const_cast<int*>(cu_seq_q.data<int>());
 
-    params.q_input = reinterpret_cast<cute_type *>(const_cast<T*>(q_input.data<T>()));
-    params.attn_out = reinterpret_cast<cute_type *>(const_cast<T*>(out.data<T>()));
-    params.seq_lens_encoder = const_cast<int*>(seq_len_encoder.data<int>());
-    params.seq_lens_decoder = const_cast<int*>(seq_len_decoder.data<int>());
-    params.block_table = const_cast<int*>(block_tables.data<int>());
-    params.max_input_length = max_input_length;
-    params.head_num = head_num;
-    params.kv_head_num = kv_head_num;
-    params.max_num_blocks_per_seq = block_tables.dims()[1];
-    params.batch_size = batch_size;
-    params.inv_sqrt_dh = 1.0f / std::sqrt(head_dim);
-    params.max_num_partitions = max_num_partitions;
-    params.maxs = reinterpret_cast<float*>(maxs.data<float>());
-    params.sums = reinterpret_cast<float*>(sums.data<float>());
-    params.partition_attn_out = reinterpret_cast<cute_type *>(partition_attn_out.data<T>());
-    params.qk_gate_topk_idx_ptr = const_cast<int*>(qk_gate_topk_idx.data<int>());
-    params.use_moba_seq_limit = use_moba_seq_limit;
-    params.cu_seq_q = const_cast<int*>(cu_seq_q.data<int>());
-
-
-    if (cache_quant_type_str == "none") {
-        params.cache_k = reinterpret_cast<cute_type *>(const_cast<T*>(cache_k.data<T>()));
-        params.cache_v = reinterpret_cast<cute_type *>(const_cast<T*>(cache_v.data<T>()));
-        run_moba_decoder_attn_hdim128<cute_type, 16, max_seq_per_block, kMaxN>(params, q_input.stream());
+  if (cache_quant_type_str == "none") {
+    params.cache_k =
+        reinterpret_cast<cute_type*>(const_cast<T*>(cache_k.data<T>()));
+    params.cache_v =
+        reinterpret_cast<cute_type*>(const_cast<T*>(cache_v.data<T>()));
+    run_moba_decoder_attn_hdim128<cute_type, 16, max_seq_per_block, kMaxN>(
+        params, q_input.stream());
+  } else {
+    params.cache_k = const_cast<uint8_t*>(cache_k.data<uint8_t>());
+    params.cache_v = const_cast<uint8_t*>(cache_v.data<uint8_t>());
+    params.cache_k_quant_scale = reinterpret_cast<cute_type*>(
+        const_cast<T*>(cache_k_quant_scale.get().data<T>()));
+    params.cache_v_quant_scale = reinterpret_cast<cute_type*>(
+        const_cast<T*>(cache_v_quant_scale.get().data<T>()));
+    params.cache_k_dequant_scale = reinterpret_cast<cute_type*>(
+        const_cast<T*>(cache_k_dequant_scale.get().data<T>()));
+    params.cache_v_dequant_scale = reinterpret_cast<cute_type*>(
+        const_cast<T*>(cache_v_dequant_scale.get().data<T>()));
+    params.cache_k_zp = reinterpret_cast<cute_type*>(
+        const_cast<T*>(cache_k_zero_points.get().data<T>()));
+    params.cache_v_zp = reinterpret_cast<cute_type*>(
+        const_cast<T*>(cache_v_zero_points.get().data<T>()));
+    if (cache_quant_type_str == "cache_int8_zp") {
+      run_moba_decoder_attn_hdim128<cute_type, 8, max_seq_per_block, kMaxN>(
+          params, q_input.stream());
+    } else if (cache_quant_type_str == "cache_int4_zp") {
+      run_moba_decoder_attn_hdim128<cute_type, 4, max_seq_per_block, kMaxN>(
+          params, q_input.stream());
     } else {
-        params.cache_k = const_cast<uint8_t*>(cache_k.data<uint8_t>());
-        params.cache_v = const_cast<uint8_t*>(cache_v.data<uint8_t>());
-        params.cache_k_quant_scale = reinterpret_cast<cute_type *>(const_cast<T*>(cache_k_quant_scale.get().data<T>()));
-        params.cache_v_quant_scale = reinterpret_cast<cute_type *>(const_cast<T*>(cache_v_quant_scale.get().data<T>()));
-        params.cache_k_dequant_scale = reinterpret_cast<cute_type *>(const_cast<T*>(cache_k_dequant_scale.get().data<T>()));
-        params.cache_v_dequant_scale = reinterpret_cast<cute_type *>(const_cast<T*>(cache_v_dequant_scale.get().data<T>()));
-        params.cache_k_zp = reinterpret_cast<cute_type *>(const_cast<T*>(cache_k_zero_points.get().data<T>()));
-        params.cache_v_zp = reinterpret_cast<cute_type *>(const_cast<T*>(cache_v_zero_points.get().data<T>()));
-        if (cache_quant_type_str == "cache_int8_zp") {
-            run_moba_decoder_attn_hdim128<cute_type, 8, max_seq_per_block, kMaxN>(params, q_input.stream());
-        } else if (cache_quant_type_str == "cache_int4_zp") {
-            run_moba_decoder_attn_hdim128<cute_type, 4, max_seq_per_block, kMaxN>(params, q_input.stream());
-        } else {
-            PADDLE_THROW(phi::errors::Unimplemented(
-            "GQA Attention not implemented for cache_quant_type_str = %s", cache_quant_type_str.c_str()));
-        }
+      PADDLE_THROW(phi::errors::Unimplemented(
+          "GQA Attention not implemented for cache_quant_type_str = %s",
+          cache_quant_type_str.c_str()));
     }
+  }
 }
 
 void MobaDecoderAttn(
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cache_k,
-        const paddle::Tensor& cache_v,
-        const paddle::Tensor& block_tables,
-        const paddle::Tensor& k_block_means,
-        const paddle::Tensor& out,
-        const paddle::Tensor& qk_gate_topk_idx,
-        const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_zero_points,
-        const paddle::optional<paddle::Tensor>& cache_v_zero_points,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_input_length,
-        const int use_moba_seq_limit,
-        const int max_seq_q,
-        const int max_seq_k,
-        const std::string &cache_quant_type_str) {
-
-    const int batch_size = block_tables.dims()[0];
-    if (q_input.dtype() == paddle::DataType::FLOAT16) {
-        return DispatchMobaDecoderAttn<phi::dtype::float16>(
-            q_input,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            cache_k,
-            cache_v,
-            block_tables,
-            k_block_means,
-            out,
-            qk_gate_topk_idx,
-            cache_k_quant_scale,
-            cache_v_quant_scale,
-            cache_k_dequant_scale,
-            cache_v_dequant_scale,
-            cache_k_zero_points,
-            cache_v_zero_points,
-            head_num,
-            kv_head_num,
-            head_dim,
-            max_seq_q,
-            max_seq_k,
-            batch_size,
-            max_input_length,
-            use_moba_seq_limit,
-            cache_quant_type_str);
-    } else if (q_input.dtype() == paddle::DataType::BFLOAT16) {
-        return DispatchMobaDecoderAttn<phi::dtype::bfloat16>(
-            q_input,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            cache_k,
-            cache_v,
-            block_tables,
-            k_block_means,
-            out,
-            qk_gate_topk_idx,
-            cache_k_quant_scale,
-            cache_v_quant_scale,
-            cache_k_dequant_scale,
-            cache_v_dequant_scale,
-            cache_k_zero_points,
-            cache_v_zero_points,
-            head_num,
-            kv_head_num,
-            head_dim,
-            max_seq_q,
-            max_seq_k,
-            batch_size,
-            max_input_length,
-            use_moba_seq_limit,
-            cache_quant_type_str);
-    }
+    const paddle::Tensor& q_input,
+    const paddle::Tensor& seq_len_encoder,
+    const paddle::Tensor& seq_len_decoder,
+    const paddle::Tensor& cu_seq_q,
+    const paddle::Tensor& cache_k,
+    const paddle::Tensor& cache_v,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& k_block_means,
+    const paddle::Tensor& out,
+    const paddle::Tensor& qk_gate_topk_idx,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zero_points,
+    const paddle::optional<paddle::Tensor>& cache_v_zero_points,
+    const int head_num,
+    const int kv_head_num,
+    const int head_dim,
+    const int max_input_length,
+    const int use_moba_seq_limit,
+    const int max_seq_q,
+    const int max_seq_k,
+    const std::string& cache_quant_type_str) {
+  const int batch_size = block_tables.dims()[0];
+  if (q_input.dtype() == paddle::DataType::FLOAT16) {
+    return DispatchMobaDecoderAttn<phi::dtype::float16>(q_input,
+                                                        seq_len_encoder,
+                                                        seq_len_decoder,
+                                                        cu_seq_q,
+                                                        cache_k,
+                                                        cache_v,
+                                                        block_tables,
+                                                        k_block_means,
+                                                        out,
+                                                        qk_gate_topk_idx,
+                                                        cache_k_quant_scale,
+                                                        cache_v_quant_scale,
+                                                        cache_k_dequant_scale,
+                                                        cache_v_dequant_scale,
+                                                        cache_k_zero_points,
+                                                        cache_v_zero_points,
+                                                        head_num,
+                                                        kv_head_num,
+                                                        head_dim,
+                                                        max_seq_q,
+                                                        max_seq_k,
+                                                        batch_size,
+                                                        max_input_length,
+                                                        use_moba_seq_limit,
+                                                        cache_quant_type_str);
+  } else if (q_input.dtype() == paddle::DataType::BFLOAT16) {
+    return DispatchMobaDecoderAttn<phi::dtype::bfloat16>(q_input,
+                                                         seq_len_encoder,
+                                                         seq_len_decoder,
+                                                         cu_seq_q,
+                                                         cache_k,
+                                                         cache_v,
+                                                         block_tables,
+                                                         k_block_means,
+                                                         out,
+                                                         qk_gate_topk_idx,
+                                                         cache_k_quant_scale,
+                                                         cache_v_quant_scale,
+                                                         cache_k_dequant_scale,
+                                                         cache_v_dequant_scale,
+                                                         cache_k_zero_points,
+                                                         cache_v_zero_points,
+                                                         head_num,
+                                                         kv_head_num,
+                                                         head_dim,
+                                                         max_seq_q,
+                                                         max_seq_k,
+                                                         batch_size,
+                                                         max_input_length,
+                                                         use_moba_seq_limit,
+                                                         cache_quant_type_str);
+  }
 }
diff --git a/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_attn_kernel.h b/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_attn_kernel.h
index 1b21f9b32a..64ce5a7887 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_attn_kernel.h
+++ b/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_attn_kernel.h
@@ -22,204 +22,190 @@
 using namespace cute;
 template <typename T>
 struct moba_decoder_attn_params {
-    T *__restrict__ q_input;
-    void *__restrict__ cache_k;
-    void *__restrict__ cache_v;
+  T *__restrict__ q_input;
+  void *__restrict__ cache_k;
+  void *__restrict__ cache_v;
 
-    T *__restrict__ attn_out;
-    T *__restrict__ partition_attn_out;
-    T *__restrict__ cache_k_dequant_scale;
-    T *__restrict__ cache_v_dequant_scale;
-    T *__restrict__ cache_k_quant_scale;
-    T *__restrict__ cache_v_quant_scale;
-    T *__restrict__ cache_k_zp;
-    T *__restrict__ cache_v_zp;
-    int * __restrict__ cu_seq_q;
-    float * sums;
-    float * maxs;
-    int * seq_lens_encoder;
-    int * seq_lens_decoder;
-    int * block_table;
-    int max_input_length;
-    int max_seq_len;
-    int head_num;
-    int kv_head_num;
-    int max_num_blocks_per_seq;
-    float scale_softmax;
-    int batch_size;
-    int max_num_partitions;
-    float inv_sqrt_dh;
-    int *qk_gate_topk_idx_ptr;
-    int use_moba_seq_limit;
+  T *__restrict__ attn_out;
+  T *__restrict__ partition_attn_out;
+  T *__restrict__ cache_k_dequant_scale;
+  T *__restrict__ cache_v_dequant_scale;
+  T *__restrict__ cache_k_quant_scale;
+  T *__restrict__ cache_v_quant_scale;
+  T *__restrict__ cache_k_zp;
+  T *__restrict__ cache_v_zp;
+  int *__restrict__ cu_seq_q;
+  float *sums;
+  float *maxs;
+  int *seq_lens_encoder;
+  int *seq_lens_decoder;
+  int *block_table;
+  int max_input_length;
+  int max_seq_len;
+  int head_num;
+  int kv_head_num;
+  int max_num_blocks_per_seq;
+  float scale_softmax;
+  int batch_size;
+  int max_num_partitions;
+  float inv_sqrt_dh;
+  int *qk_gate_topk_idx_ptr;
+  int use_moba_seq_limit;
 };
 
 template <typename cute_type_, int DataBits_>
 struct CacheKV_quant_traits {
-    using cuteType = cute_type_;
-    static constexpr int kDataBits = DataBits_;
-    static constexpr int kBlockSize = 64;
-    static constexpr int kHeadDim = 128;
-    static constexpr int kBlockKSmem = 64;
-    using SmemLayoutAtomQ = decltype(
-        composition(Swizzle<3, 3, 3>{},
-                    Layout<
-                        Shape<Int<8>, Int<kBlockKSmem>>,
-                        Stride<Int<kBlockKSmem>, _1>>{}));
+  using cuteType = cute_type_;
+  static constexpr int kDataBits = DataBits_;
+  static constexpr int kBlockSize = 64;
+  static constexpr int kHeadDim = 128;
+  static constexpr int kBlockKSmem = 64;
+  using SmemLayoutAtomQ = decltype(composition(
+      Swizzle<3, 3, 3>{},
+      Layout<Shape<Int<8>, Int<kBlockKSmem>>, Stride<Int<kBlockKSmem>, _1>>{}));
 
-    using SmemLayoutKV = decltype(tile_to_shape(
-        SmemLayoutAtomQ{},
-        Shape<Int<kBlockSize>, Int<kHeadDim>>{}));
+  using SmemLayoutKV = decltype(tile_to_shape(
+      SmemLayoutAtomQ{}, Shape<Int<kBlockSize>, Int<kHeadDim>>{}));
 
-    static constexpr int kNWarps = 4;
-    static constexpr int kNThreads = kNWarps * 32;
+  static constexpr int kNWarps = 4;
+  static constexpr int kNThreads = kNWarps * 32;
 
+  static constexpr int kThreadPerValue = 16 / sizeof(cuteType);
+  static constexpr int kThreadsPerRow = kHeadDim / kThreadPerValue;
 
-    static constexpr int kThreadPerValue = 16 / sizeof(cuteType);
-    static constexpr int kThreadsPerRow = kHeadDim / kThreadPerValue;
+  using GmemLayoutAtom =
+      Layout<Shape<Int<kNThreads / kThreadsPerRow>, Int<kThreadsPerRow>>,
+             Stride<Int<kThreadsPerRow>, _1>>;
 
-    using GmemLayoutAtom = Layout<
-        Shape <Int<kNThreads / kThreadsPerRow>, Int<kThreadsPerRow>>,
-        Stride<Int<kThreadsPerRow>, _1>>;
+  using GmemTiledCopyQ = decltype(make_tiled_copy(
+      Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, cuteType>{},
+      GmemLayoutAtom{},
+      Layout<Shape<_1, Int<kThreadPerValue>>>{}));
 
-    using GmemTiledCopyQ = decltype(
-        make_tiled_copy(Copy_Atom<
-            SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, cuteType>{},
-            GmemLayoutAtom{},
-            Layout<Shape<_1, Int<kThreadPerValue>>>{}));
+  using MMA_Atom_Arch =
+      std::conditional_t<std::is_same_v<cuteType, cutlass::half_t>,
+                         MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
+                         MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>>;
 
-    using MMA_Atom_Arch = std::conditional_t<
-        std::is_same_v<cuteType, cutlass::half_t>,
-        MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
-        MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>
-    >;
+  using ValLayoutMNK = Layout<Shape<_1, _4, _1>>;
 
-    using ValLayoutMNK = Layout<Shape<_1,_4,_1>>;
+  using PermutationMNK = Tile<_16, Int<16 * kNWarps>, _16>;
 
-    using PermutationMNK = Tile<_16, Int<16 * kNWarps>, _16>;
+  using TiledMma = TiledMMA<MMA_Atom_Arch, ValLayoutMNK, PermutationMNK>;
 
-    using TiledMma = TiledMMA<
-        MMA_Atom_Arch,
-        ValLayoutMNK,
-        PermutationMNK>;
+  using SmemCopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, cuteType>;
 
-    using SmemCopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, cuteType>;
+  using SmemLayoutAtomVtransposed =
+      decltype(composition(Swizzle<3, 3, 3>{},
+                           Layout<Shape<Int<kBlockKSmem>, Int<kBlockSize>>,
+                                  Stride<_1, Int<kBlockKSmem>>>{}));
 
-    using SmemLayoutAtomVtransposed = decltype(
-        composition(Swizzle<3, 3, 3>{},
-                    Layout<Shape<Int<kBlockKSmem>, Int<kBlockSize>>,
-                           Stride<_1, Int<kBlockKSmem>>>{}));
+  using SmemLayoutVtransposed = decltype(tile_to_shape(
+      SmemLayoutAtomVtransposed{}, Shape<Int<kHeadDim>, Int<kBlockSize>>{}));
 
-    using SmemLayoutVtransposed = decltype(tile_to_shape(
-        SmemLayoutAtomVtransposed{},
-        Shape<Int<kHeadDim>, Int<kBlockSize>>{}));
+  using SmemLayoutVtransposedNoSwizzle =
+      decltype(get_nonswizzle_portion(SmemLayoutVtransposed{}));
 
-    using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{}));
+  using SmemCopyAtomTransposed = Copy_Atom<SM75_U16x8_LDSM_T, cuteType>;
 
-    using SmemCopyAtomTransposed = Copy_Atom<SM75_U16x8_LDSM_T, cuteType>;
-
-    static constexpr int kShareMemSize = size(SmemLayoutKV{}) * 2 * sizeof(cuteType);
+  static constexpr int kShareMemSize =
+      size(SmemLayoutKV{}) * 2 * sizeof(cuteType);
 };
 
 template <int kGqaGroupSize_, int kTileN_, int kMaxN_, typename CacheKV_traits_>
 struct moba_decoder_attn_kernel_traits {
-    using ElementAccum = float;
-    using CacheKV_traits = CacheKV_traits_;
-    using cuteType = typename CacheKV_traits::cuteType;
-    static constexpr int kDataBits = CacheKV_traits::kDataBits;
-    static constexpr int kTileN = kTileN_;
-    static constexpr int kMaxN = kMaxN_;
-    static constexpr int kGqaGroupSize = kGqaGroupSize_;
-    static constexpr int kHeadDim = CacheKV_traits::kHeadDim;
-    static constexpr int kHeadDimKV = kHeadDim / (16 / kDataBits);
-    static constexpr int kMinGemmM = 16;
-    static constexpr int kBlockM = (kGqaGroupSize + kMinGemmM - 1) / kMinGemmM * kMinGemmM;
-    static constexpr int kBlockSize = CacheKV_traits::kBlockSize;
-    static_assert(kGqaGroupSize <= 16);
-    static constexpr int32_t kNWarps = CacheKV_traits::kNWarps;
+  using ElementAccum = float;
+  using CacheKV_traits = CacheKV_traits_;
+  using cuteType = typename CacheKV_traits::cuteType;
+  static constexpr int kDataBits = CacheKV_traits::kDataBits;
+  static constexpr int kTileN = kTileN_;
+  static constexpr int kMaxN = kMaxN_;
+  static constexpr int kGqaGroupSize = kGqaGroupSize_;
+  static constexpr int kHeadDim = CacheKV_traits::kHeadDim;
+  static constexpr int kHeadDimKV = kHeadDim / (16 / kDataBits);
+  static constexpr int kMinGemmM = 16;
+  static constexpr int kBlockM =
+      (kGqaGroupSize + kMinGemmM - 1) / kMinGemmM * kMinGemmM;
+  static constexpr int kBlockSize = CacheKV_traits::kBlockSize;
+  static_assert(kGqaGroupSize <= 16);
+  static constexpr int32_t kNWarps = CacheKV_traits::kNWarps;
 
-    static constexpr int kBlockKSmem = CacheKV_traits::kBlockKSmem;
-    static constexpr int kBlockKVSmem = kHeadDimKV <= 64 ? kHeadDimKV : 64;
-    static_assert(kHeadDim % kBlockKSmem == 0);
-    static constexpr int kNReduceWarps = 4;
-    static constexpr int kNReduceThreads = kNReduceWarps * 32;
+  static constexpr int kBlockKSmem = CacheKV_traits::kBlockKSmem;
+  static constexpr int kBlockKVSmem = kHeadDimKV <= 64 ? kHeadDimKV : 64;
+  static_assert(kHeadDim % kBlockKSmem == 0);
+  static constexpr int kNReduceWarps = 4;
+  static constexpr int kNReduceThreads = kNReduceWarps * 32;
 
+  using SmemLayoutAtomQ = typename CacheKV_traits::SmemLayoutAtomQ;
 
-    using SmemLayoutAtomQ = typename CacheKV_traits::SmemLayoutAtomQ;
+  using SmemLayoutQ = decltype(tile_to_shape(
+      SmemLayoutAtomQ{}, Shape<Int<kBlockM>, Int<kHeadDim>>{}));
 
-    using SmemLayoutQ = decltype(tile_to_shape(
-        SmemLayoutAtomQ{},
-        Shape<Int<kBlockM>, Int<kHeadDim>>{}));
+  using SmemLayoutQK = decltype(tile_to_shape(
+      SmemLayoutAtomQ{}, Shape<Int<kBlockM>, Int<kBlockSize>>{}));
 
-    using SmemLayoutQK = decltype(tile_to_shape(
-        SmemLayoutAtomQ{},
-        Shape<Int<kBlockM>, Int<kBlockSize>>{}));
+  using SmemLayoutAtomKV =
+      decltype(composition(Swizzle<3, 3, 3>{},
+                           Layout<Shape<Int<8>, Int<kBlockKVSmem>>,
+                                  Stride<Int<kBlockKVSmem>, _1>>{}));
 
-    using SmemLayoutAtomKV = decltype(
-        composition(Swizzle<3, 3, 3>{},
-                    Layout<
-                        Shape<Int<8>, Int<kBlockKVSmem>>,
-                        Stride<Int<kBlockKVSmem>, _1>>{}));
+  using SmemLayoutKV_ = decltype(tile_to_shape(
+      SmemLayoutAtomKV{}, Shape<Int<kBlockSize>, Int<kHeadDimKV>>{}));
 
-    using SmemLayoutKV_ = decltype(tile_to_shape(
-        SmemLayoutAtomKV{},
-        Shape<Int<kBlockSize>, Int<kHeadDimKV>>{}));
+  using SmemLayoutKV =
+      std::conditional_t<kDataBits == 16,
+                         SmemLayoutKV_,
+                         decltype(get_nonswizzle_portion(SmemLayoutKV_{}))>;
 
-    using SmemLayoutKV = std::conditional_t<
-        kDataBits == 16,
-        SmemLayoutKV_,
-        decltype(get_nonswizzle_portion(SmemLayoutKV_{}))
-    >;
+  constexpr static int kBlockKVSize = kDataBits == 4 ? 32 : kBlockSize;
+  using SmemLayoutAtomVtransposed =
+      decltype(composition(Swizzle<3, 3, 3>{},
+                           Layout<Shape<Int<kBlockKSmem>, Int<kBlockKVSize>>,
+                                  Stride<_1, Int<kBlockKSmem>>>{}));
 
-    constexpr static int kBlockKVSize = kDataBits == 4 ? 32 : kBlockSize;
-    using SmemLayoutAtomVtransposed = decltype(
-        composition(Swizzle<3, 3, 3>{},
-                    Layout<Shape<Int<kBlockKSmem>, Int<kBlockKVSize>>,
-                           Stride<_1, Int<kBlockKSmem>>>{}));
+  using SmemLayoutVtransposed = decltype(tile_to_shape(
+      SmemLayoutAtomVtransposed{}, Shape<Int<kHeadDim>, Int<kBlockKVSize>>{}));
 
-    using SmemLayoutVtransposed = decltype(tile_to_shape(
-        SmemLayoutAtomVtransposed{},
-        Shape<Int<kHeadDim>, Int<kBlockKVSize>>{}));
+  using SmemLayoutVtransposedNoSwizzle =
+      decltype(get_nonswizzle_portion(SmemLayoutVtransposed{}));
 
-    using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{}));
+  static constexpr int kThreadsPerRow = CacheKV_traits::kThreadsPerRow;
+  static constexpr int kThreadsKVPerRow = kThreadsPerRow / (16 / kDataBits);
+  static constexpr int kNThreads = CacheKV_traits::kNThreads;
 
-    static constexpr int kThreadsPerRow = CacheKV_traits::kThreadsPerRow;
-    static constexpr int kThreadsKVPerRow = kThreadsPerRow / (16 / kDataBits);
-    static constexpr int kNThreads = CacheKV_traits::kNThreads;
+  using GmemKVLayoutAtom =
+      Layout<Shape<Int<kNThreads / kThreadsKVPerRow>, Int<kThreadsKVPerRow>>,
+             Stride<Int<kThreadsKVPerRow>, _1>>;
 
-    using GmemKVLayoutAtom = Layout<
-        Shape<Int<kNThreads / kThreadsKVPerRow>, Int<kThreadsKVPerRow>>,
-        Stride<Int<kThreadsKVPerRow>, _1>>;
+  using SmemCopyAtom = typename CacheKV_traits::SmemCopyAtom;
+  using TiledMma = typename CacheKV_traits::TiledMma;
 
-    using SmemCopyAtom = typename CacheKV_traits::SmemCopyAtom;
-    using TiledMma = typename CacheKV_traits::TiledMma;
+  static constexpr int kThreadPerValue = CacheKV_traits::kThreadPerValue;
 
-    static constexpr int kThreadPerValue = CacheKV_traits::kThreadPerValue;
+  using GmemTiledCopyQ = typename CacheKV_traits::GmemTiledCopyQ;
+  using GmemLayoutAtom = typename CacheKV_traits::GmemLayoutAtom;
+  using GmemTiledCopyKV = decltype(make_tiled_copy(
+      Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, cuteType>{},
+      GmemKVLayoutAtom{},
+      Layout<Shape<_1, Int<kThreadPerValue>>>{}));
 
-    using GmemTiledCopyQ = typename CacheKV_traits::GmemTiledCopyQ;
-    using GmemLayoutAtom = typename CacheKV_traits::GmemLayoutAtom;
-    using GmemTiledCopyKV = decltype(
-        make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, cuteType>{},
-                        GmemKVLayoutAtom{},
-                        Layout<Shape<_1, Int<kThreadPerValue>>>{}));
+  using SmemCopyAtomTransposed =
+      typename CacheKV_traits::SmemCopyAtomTransposed;
 
+  using GmemTiledCopyO =
+      decltype(make_tiled_copy(Copy_Atom<DefaultCopy, cuteType>{},
+                               GmemLayoutAtom{},
+                               Layout<Shape<_1, Int<kThreadPerValue>>>{}));
+  using SmemCopyAtomO = Copy_Atom<DefaultCopy, cuteType>;
 
-    using SmemCopyAtomTransposed = typename CacheKV_traits::SmemCopyAtomTransposed;
+  using SmemLayoutAtomO = decltype(composition(
+      Swizzle<3, 3, 3>{},
+      Layout<Shape<Int<8>, Int<kBlockKSmem>>, Stride<Int<kBlockKSmem>, _1>>{}));
 
-    using GmemTiledCopyO = decltype(
-        make_tiled_copy(Copy_Atom<DefaultCopy, cuteType>{},
-                        GmemLayoutAtom{},
-                        Layout<Shape<_1, Int<kThreadPerValue>>>{}));
-    using SmemCopyAtomO = Copy_Atom<DefaultCopy, cuteType>;
+  using SmemLayoutO = decltype(tile_to_shape(
+      SmemLayoutAtomO{}, Shape<Int<kBlockM>, Int<kHeadDim>>{}));
 
-    using SmemLayoutAtomO = decltype(
-        composition(Swizzle<3, 3, 3>{},
-                    Layout<
-                        Shape<Int<8>, Int<kBlockKSmem>>,
-                        Stride<Int<kBlockKSmem>, _1>>{}));
-
-    using SmemLayoutO = decltype(tile_to_shape(
-        SmemLayoutAtomO{},
-        Shape<Int<kBlockM>, Int<kHeadDim>>{}));
-
-    static constexpr int kShareMemSize = (size(SmemLayoutQ{}) + size(SmemLayoutQK{}) + size(SmemLayoutKV{}) * 2) * sizeof(cuteType);
+  static constexpr int kShareMemSize =
+      (size(SmemLayoutQ{}) + size(SmemLayoutQK{}) + size(SmemLayoutKV{}) * 2) *
+      sizeof(cuteType);
 };
diff --git a/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_write_cache.cu b/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_write_cache.cu
index b1cc5b1ef0..8ee6c8c1d4 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_write_cache.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_decoder_write_cache.cu
@@ -16,174 +16,197 @@
 #include "../moba_attn_utils.hpp"
 #include "moba_attn/moba_attn.h"
 
-template <typename T, int kBlockSize, int kHeadDim, int moba_block_size, int kMaxN>
-__global__ void moba_decoder_attn_write_c16(
-        const T * qkv_out,
-        const T * qkv_bias,
-        T * q_input,
-        const int * cu_seq_q,
-        const int * cu_seq_k,
-        const int * seq_len_encoder,
-        const int * seq_len_decoder,
-        T * cache_k,
-        T * cache_v,
-        const int * block_tables,
-        const float * rope_sin_cos,
-        T *k_block_means,
-        const int head_num,
-        const int kv_head_num,
-        const int max_blocks_per_seq,
-        const int max_input_length) {
+template <typename T,
+          int kBlockSize,
+          int kHeadDim,
+          int moba_block_size,
+          int kMaxN>
+__global__ void moba_decoder_attn_write_c16(const T* qkv_out,
+                                            const T* qkv_bias,
+                                            T* q_input,
+                                            const int* cu_seq_q,
+                                            const int* cu_seq_k,
+                                            const int* seq_len_encoder,
+                                            const int* seq_len_decoder,
+                                            T* cache_k,
+                                            T* cache_v,
+                                            const int* block_tables,
+                                            const float* rope_sin_cos,
+                                            T* k_block_means,
+                                            const int head_num,
+                                            const int kv_head_num,
+                                            const int max_blocks_per_seq,
+                                            const int max_input_length) {
+  int bidh = blockIdx.x;
+  const int bidb = blockIdx.y;
+  const int tidx = threadIdx.x;
+  const int seq_len = seq_len_decoder[bidb];
 
-    int bidh = blockIdx.x;
-    const int bidb = blockIdx.y;
-    const int tidx = threadIdx.x;
-    const int seq_len = seq_len_decoder[bidb];
+  if (seq_len == 0) {
+    return;
+  }
 
-    if (seq_len == 0) {
-        return;
+  constexpr int kPackSize = 4;
+  using SrcType = Vec<T, kPackSize>;
+  using rope_type = Vec<float, kPackSize / 2>;
+  SrcType src, bias, k_prev;
+  rope_type sin, cos;
+  const int bias_idx = bidh * kHeadDim + tidx * kPackSize;
+  const int ori_token_idx = cu_seq_q[bidb];
+  src.load_from(qkv_out +
+                ori_token_idx * (head_num + 2 * kv_head_num) * kHeadDim +
+                bias_idx);
+  if (qkv_bias != nullptr) {
+    bias.load_from(qkv_bias + bias_idx);
+    src.add(bias);
+  }
+
+  const int32_t* block_table_now = block_tables + bidb * max_blocks_per_seq;
+  const int32_t physical_block_number = block_table_now[seq_len / kBlockSize];
+
+  if (bidh < head_num) {
+    const float* cos_rope =
+        rope_sin_cos + seq_len * (kHeadDim / 2) + tidx * (kPackSize / 2);
+    const float* sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
+    sin.load_from(sin_rope);
+    cos.load_from(cos_rope);
+    apply_rotary_embedding<T, kPackSize>(src, cos, sin);
+
+    src.store_to(q_input + cu_seq_q[bidb] * head_num * kHeadDim + bias_idx);
+  } else if (bidh < head_num + kv_head_num) {
+    bidh -= head_num;
+    const int token_in_blocks = seq_len % kBlockSize;
+    const float* cos_rope =
+        rope_sin_cos + seq_len * (kHeadDim / 2) + tidx * (kPackSize / 2);
+    const float* sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
+    sin.load_from(sin_rope);
+    cos.load_from(cos_rope);
+    apply_rotary_embedding<T, kPackSize>(src, cos, sin);
+
+    T* cache = cache_k +
+               physical_block_number * kv_head_num * kBlockSize * kHeadDim +
+               bidh * kBlockSize * kHeadDim + tidx * kPackSize +
+               token_in_blocks * kHeadDim;
+    src.store_to(cache);
+
+    const int seq_len_block = seq_len / moba_block_size;
+
+    const int store_mean_idx =
+        (bidb * kMaxN + seq_len_block) * kv_head_num * kHeadDim +
+        bidh * kHeadDim + tidx * kPackSize;
+
+    if (seq_len % moba_block_size != 0) {
+      const int token_num_prev = seq_len % moba_block_size;
+      const float inv_tokens_sum = fdividef(1.0f, token_num_prev + 1);
+      k_prev.load_from(k_block_means + store_mean_idx);
+
+#pragma unroll
+      for (int i = 0; i < kPackSize; i++) {
+        src.data.elt[i] =
+            T(inv_tokens_sum * (float(src.data.elt[i]) +
+                                float(k_prev.data.elt[i]) * token_num_prev));
+      }
     }
 
-    constexpr int kPackSize = 4;
-    using SrcType = Vec<T, kPackSize>;
-    using rope_type = Vec<float, kPackSize / 2>;
-    SrcType src, bias, k_prev;
-    rope_type sin, cos;
-    const int bias_idx = bidh * kHeadDim + tidx * kPackSize;
-    const int ori_token_idx = cu_seq_q[bidb];
-    src.load_from(qkv_out + ori_token_idx * (head_num + 2 * kv_head_num) * kHeadDim + bias_idx);
-    if (qkv_bias != nullptr) {
-        bias.load_from(qkv_bias + bias_idx);
-        src.add(bias);
-    }
-
-    const int32_t *block_table_now = block_tables + bidb * max_blocks_per_seq;
-    const int32_t physical_block_number = block_table_now[seq_len / kBlockSize];
-
-
-    if (bidh < head_num) {
-        const float * cos_rope = rope_sin_cos + seq_len * (kHeadDim / 2) + tidx * (kPackSize / 2);
-        const float * sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
-        sin.load_from(sin_rope);
-        cos.load_from(cos_rope);
-        apply_rotary_embedding<T, kPackSize>(src, cos, sin);
-
-        src.store_to(q_input + cu_seq_q[bidb] * head_num * kHeadDim + bias_idx);
-    } else if (bidh < head_num + kv_head_num) {
-        bidh -= head_num;
-        const int token_in_blocks = seq_len % kBlockSize;
-        const float * cos_rope = rope_sin_cos + seq_len * (kHeadDim / 2) + tidx * (kPackSize / 2);
-        const float * sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
-        sin.load_from(sin_rope);
-        cos.load_from(cos_rope);
-        apply_rotary_embedding<T, kPackSize>(src, cos, sin);
-
-        T * cache = cache_k + physical_block_number * kv_head_num * kBlockSize * kHeadDim + bidh * kBlockSize * kHeadDim + tidx * kPackSize + token_in_blocks * kHeadDim;
-        src.store_to(cache);
-
-        const int seq_len_block = seq_len / moba_block_size;
-
-        const int store_mean_idx = (bidb * kMaxN + seq_len_block) * kv_head_num * kHeadDim + bidh * kHeadDim + tidx * kPackSize;
-
-        if (seq_len % moba_block_size != 0) {
-            const int token_num_prev = seq_len % moba_block_size;
-            const float inv_tokens_sum = fdividef(1.0f, token_num_prev + 1);
-            k_prev.load_from(k_block_means + store_mean_idx);
-
-            #pragma unroll
-            for (int i = 0; i < kPackSize; i++) {
-                src.data.elt[i] = T(inv_tokens_sum * (float(src.data.elt[i]) + float(k_prev.data.elt[i]) * token_num_prev));
-            }
-        }
-
-        src.store_to(k_block_means + store_mean_idx);
-
-    } else {
-        bidh -= head_num + kv_head_num;
-        const int token_in_blocks = seq_len % kBlockSize;
-        T * cache = cache_v + physical_block_number * kv_head_num * kBlockSize * kHeadDim + bidh * kBlockSize * kHeadDim + tidx * kPackSize + token_in_blocks * kHeadDim;
-        src.store_to(cache);
-    }
+    src.store_to(k_block_means + store_mean_idx);
 
+  } else {
+    bidh -= head_num + kv_head_num;
+    const int token_in_blocks = seq_len % kBlockSize;
+    T* cache = cache_v +
+               physical_block_number * kv_head_num * kBlockSize * kHeadDim +
+               bidh * kBlockSize * kHeadDim + tidx * kPackSize +
+               token_in_blocks * kHeadDim;
+    src.store_to(cache);
+  }
 }
 
 void MobaDecoderAttnWriteCacheKv(
-        const paddle::Tensor& qkv_out,
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cache_k,
-        const paddle::Tensor& cache_v,
-        const paddle::Tensor& block_tables,
-        const paddle::Tensor& rope_sin_cos,
-        const paddle::Tensor& k_block_means,
-        const paddle::optional<paddle::Tensor>& qkv_bias,
-        const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_zero_points,
-        const paddle::optional<paddle::Tensor>& cache_v_zero_points,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_input_length,
-        const std::string &cache_quant_type_str) {
-
-    constexpr int kThreads = 32;
-    constexpr int kHeadDim = 128;
-    constexpr int kMobaBlockSize = 128;
-    constexpr int kMaxN = 1024;
-    assert(kHeadDim == head_dim);
-    constexpr int kBlockSize = 64;
-    const int max_blocks_per_seq = block_tables.dims()[1];
-    const int batch_size = block_tables.dims()[0];
-    if (cache_quant_type_str == "none") {
-        dim3 grid_dims;
-        grid_dims.x = head_num + kv_head_num * 2;
-        grid_dims.y = batch_size;
-        if (qkv_out.dtype() == paddle::DataType::FLOAT16) {
-            using T = phi::dtype::float16;
-            moba_decoder_attn_write_c16<T, kBlockSize, kHeadDim, kMobaBlockSize, kMaxN><<<grid_dims, kThreads, 0, qkv_out.stream()>>>(
-                qkv_out.data<T>(),
-                qkv_bias ? qkv_bias.get().data<T>() : nullptr,
-                const_cast<T*>(q_input.data<T>()),
-                cu_seq_q.data<int>(),
-                cu_seq_k.data<int>(),
-                seq_len_encoder.data<int>(),
-                seq_len_decoder.data<int>(),
-                const_cast<T *>(cache_k.data<T>()),
-                const_cast<T *>(cache_v.data<T>()),
-                block_tables.data<int>(),
-                rope_sin_cos.data<float>(),
-                const_cast<T*>(k_block_means.data<T>()),
-                head_num,
-                kv_head_num,
-                max_blocks_per_seq,
-                max_input_length);
-        } else if (qkv_out.dtype() == paddle::DataType::BFLOAT16) {
-            using T = phi::dtype::bfloat16;
-            moba_decoder_attn_write_c16<T, kBlockSize, kHeadDim, kMobaBlockSize, kMaxN><<<grid_dims, kThreads, 0, qkv_out.stream()>>>(
-                qkv_out.data<T>(),
-                qkv_bias ? qkv_bias.get().data<T>() : nullptr,
-                const_cast<T*>(q_input.data<T>()),
-                cu_seq_q.data<int>(),
-                cu_seq_k.data<int>(),
-                seq_len_encoder.data<int>(),
-                seq_len_decoder.data<int>(),
-                const_cast<T *>(cache_k.data<T>()),
-                const_cast<T *>(cache_v.data<T>()),
-                block_tables.data<int>(),
-                rope_sin_cos.data<float>(),
-                const_cast<T*>(k_block_means.data<T>()),
-                head_num,
-                kv_head_num,
-                max_blocks_per_seq,
-                max_input_length);
-        }
-    } else {
-        PD_THROW("Only supported cache_quant_type_str in ['none'].");
+    const paddle::Tensor& qkv_out,
+    const paddle::Tensor& q_input,
+    const paddle::Tensor& cu_seq_q,
+    const paddle::Tensor& cu_seq_k,
+    const paddle::Tensor& seq_len_encoder,
+    const paddle::Tensor& seq_len_decoder,
+    const paddle::Tensor& cache_k,
+    const paddle::Tensor& cache_v,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& rope_sin_cos,
+    const paddle::Tensor& k_block_means,
+    const paddle::optional<paddle::Tensor>& qkv_bias,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zero_points,
+    const paddle::optional<paddle::Tensor>& cache_v_zero_points,
+    const int head_num,
+    const int kv_head_num,
+    const int head_dim,
+    const int max_input_length,
+    const std::string& cache_quant_type_str) {
+  constexpr int kThreads = 32;
+  constexpr int kHeadDim = 128;
+  constexpr int kMobaBlockSize = 128;
+  constexpr int kMaxN = 1024;
+  assert(kHeadDim == head_dim);
+  constexpr int kBlockSize = 64;
+  const int max_blocks_per_seq = block_tables.dims()[1];
+  const int batch_size = block_tables.dims()[0];
+  if (cache_quant_type_str == "none") {
+    dim3 grid_dims;
+    grid_dims.x = head_num + kv_head_num * 2;
+    grid_dims.y = batch_size;
+    if (qkv_out.dtype() == paddle::DataType::FLOAT16) {
+      using T = phi::dtype::float16;
+      moba_decoder_attn_write_c16<T,
+                                  kBlockSize,
+                                  kHeadDim,
+                                  kMobaBlockSize,
+                                  kMaxN>
+          <<<grid_dims, kThreads, 0, qkv_out.stream()>>>(
+              qkv_out.data<T>(),
+              qkv_bias ? qkv_bias.get().data<T>() : nullptr,
+              const_cast<T*>(q_input.data<T>()),
+              cu_seq_q.data<int>(),
+              cu_seq_k.data<int>(),
+              seq_len_encoder.data<int>(),
+              seq_len_decoder.data<int>(),
+              const_cast<T*>(cache_k.data<T>()),
+              const_cast<T*>(cache_v.data<T>()),
+              block_tables.data<int>(),
+              rope_sin_cos.data<float>(),
+              const_cast<T*>(k_block_means.data<T>()),
+              head_num,
+              kv_head_num,
+              max_blocks_per_seq,
+              max_input_length);
+    } else if (qkv_out.dtype() == paddle::DataType::BFLOAT16) {
+      using T = phi::dtype::bfloat16;
+      moba_decoder_attn_write_c16<T,
+                                  kBlockSize,
+                                  kHeadDim,
+                                  kMobaBlockSize,
+                                  kMaxN>
+          <<<grid_dims, kThreads, 0, qkv_out.stream()>>>(
+              qkv_out.data<T>(),
+              qkv_bias ? qkv_bias.get().data<T>() : nullptr,
+              const_cast<T*>(q_input.data<T>()),
+              cu_seq_q.data<int>(),
+              cu_seq_k.data<int>(),
+              seq_len_encoder.data<int>(),
+              seq_len_decoder.data<int>(),
+              const_cast<T*>(cache_k.data<T>()),
+              const_cast<T*>(cache_v.data<T>()),
+              block_tables.data<int>(),
+              rope_sin_cos.data<float>(),
+              const_cast<T*>(k_block_means.data<T>()),
+              head_num,
+              kv_head_num,
+              max_blocks_per_seq,
+              max_input_length);
     }
+  } else {
+    PD_THROW("Only supported cache_quant_type_str in ['none'].");
+  }
 }
diff --git a/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_qk_sort_decoder.cu b/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_qk_sort_decoder.cu
index 8808533f88..4cabbd4f31 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_qk_sort_decoder.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_decoder_attn/moba_qk_sort_decoder.cu
@@ -20,220 +20,214 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-template <typename T, int knthreads, int moba_block_size, int kBlockMaxN, int searchtimes>
-__global__ void qk_gate_sort_decoder_kernel(
-        const T* qk_gate_weight,
-        int * qk_gate_topk_idx,
-        const int *decoder_seq_lens,
-        const int head_num,
-        const int kv_head_num,
-        const int kGqaGroupSize,
-        const int top_k_left,
-        const int top_k_right,
-        const int use_moba_seq_limit) {
+template <typename T,
+          int knthreads,
+          int moba_block_size,
+          int kBlockMaxN,
+          int searchtimes>
+__global__ void qk_gate_sort_decoder_kernel(const T* qk_gate_weight,
+                                            int* qk_gate_topk_idx,
+                                            const int* decoder_seq_lens,
+                                            const int head_num,
+                                            const int kv_head_num,
+                                            const int kGqaGroupSize,
+                                            const int top_k_left,
+                                            const int top_k_right,
+                                            const int use_moba_seq_limit) {
+  const int bidb = blockIdx.x;
+  const int bidh = blockIdx.y;
+  const int tidx = threadIdx.x;
+  const int bidh_kv = bidh / kGqaGroupSize;
 
-    const int bidb = blockIdx.x;
-    const int bidh = blockIdx.y;
-    const int tidx = threadIdx.x;
-    const int bidh_kv = bidh / kGqaGroupSize;
+  if (decoder_seq_lens[bidb] == 0 ||
+      decoder_seq_lens[bidb] < use_moba_seq_limit) {
+    return;
+  }
+  const int seq_len =
+      (decoder_seq_lens[bidb] + moba_block_size - 1) / moba_block_size;
 
-    if (decoder_seq_lens[bidb] == 0 || decoder_seq_lens[bidb] < use_moba_seq_limit) {
-        return;
+  constexpr int kPackSize = kBlockMaxN / knthreads;
+
+  static_assert(kBlockMaxN % knthreads == 0);
+
+  T token_mean[kPackSize];
+
+  using SrcType = Vec<T, kPackSize>;
+  using SrcType_f = Vec<float, kPackSize>;
+  using SrcType_i = Vec<int, kPackSize>;
+
+  SrcType src;
+  SrcType_f src_f;
+  SrcType_i select_idx;
+
+  select_idx.set_zero();
+
+  const int load_offset =
+      bidb * head_num * kBlockMaxN + bidh * kBlockMaxN + tidx * kPackSize;
+
+  src.load_from(qk_gate_weight + load_offset);
+
+  float max_global = -FLT_MAX;
+  float min_global = FLT_MAX;
+
+  const int data_len = seq_len - tidx * kPackSize;
+
+#pragma unroll
+  for (int i = 0; i < kPackSize; i++) {
+    if (i < data_len) {
+      src_f.data.elt[i] = float(src.data.elt[i]);
+      min_global = min(min_global, src_f.data.elt[i]);
+    } else {
+      src_f.data.elt[i] = -FLT_MAX;
     }
-    const int seq_len = (decoder_seq_lens[bidb] + moba_block_size - 1) / moba_block_size;
+    max_global = max(max_global, src_f.data.elt[i]);
+  }
 
-    constexpr int kPackSize = kBlockMaxN / knthreads;
+  max_global = BlockAllReduce<float, MaxOp<float>, knthreads>(max_global);
+  min_global = BlockAllReduce<float, MinOp<float>, knthreads>(min_global);
 
-    static_assert(kBlockMaxN % knthreads == 0);
+  float right_limit = max_global;
+  float left_limit = min_global;
 
-    T token_mean[kPackSize];
+  float mid_limit;
+  int count;
 
-    using SrcType = Vec<T, kPackSize>;
-    using SrcType_f = Vec<float, kPackSize>;
-    using SrcType_i = Vec<int, kPackSize>;
-
-    SrcType src;
-    SrcType_f src_f;
-    SrcType_i select_idx;
-
-    select_idx.set_zero();
-
-    const int load_offset = bidb * head_num * kBlockMaxN + bidh * kBlockMaxN + tidx * kPackSize;
-
-    src.load_from(qk_gate_weight + load_offset);
-
-    float max_global = -FLT_MAX;
-    float min_global = FLT_MAX;
-
-    const int data_len = seq_len - tidx * kPackSize;
-
-    #pragma unroll
-    for (int i = 0; i < kPackSize; i++) {
-        if (i < data_len) {
-            src_f.data.elt[i] = float(src.data.elt[i]);
-            min_global = min(min_global, src_f.data.elt[i]);
-        } else {
-            src_f.data.elt[i] = -FLT_MAX;
-        }
-        max_global = max(max_global, src_f.data.elt[i]);
+#pragma unroll
+  for (int i = 0; i < searchtimes; i++) {
+    mid_limit = (left_limit + right_limit) * 0.5f;
+    count = get_data_count<kPackSize, knthreads>(src_f.data.elt, mid_limit);
+    if (count < top_k_left) {
+      right_limit = mid_limit;
+    } else if (count > top_k_right) {
+      left_limit = mid_limit;
+    } else {
+      break;
     }
+  }
 
+  const int store_idx =
+      bidb * kv_head_num * kBlockMaxN + bidh_kv * kBlockMaxN + tidx * kPackSize;
 
-    max_global = BlockAllReduce<float, MaxOp<float>, knthreads>(max_global);
-    min_global = BlockAllReduce<float, MinOp<float>, knthreads>(min_global);
-
-
-    float right_limit = max_global;
-    float left_limit = min_global;
-
-    float mid_limit;
-    int count;
-
-    #pragma unroll
-    for (int i = 0; i < searchtimes; i++) {
-        mid_limit = (left_limit + right_limit) * 0.5f;
-        count = get_data_count<kPackSize, knthreads>(src_f.data.elt, mid_limit);
-        if (count < top_k_left) {
-            right_limit = mid_limit;
-        } else if (count > top_k_right) {
-            left_limit = mid_limit;
-        } else {
-            break;
-        }
+#pragma unroll
+  for (int i = 0; i < kPackSize; i++) {
+    if (src_f.data.elt[i] >= mid_limit) {
+      qk_gate_topk_idx[store_idx + i] = 1;
     }
+  }
 
-    const int store_idx = bidb * kv_head_num * kBlockMaxN + bidh_kv * kBlockMaxN + tidx * kPackSize;
-
-    #pragma unroll
-    for (int i = 0; i < kPackSize; i++) {
-        if (src_f.data.elt[i] >= mid_limit) {
-            qk_gate_topk_idx[store_idx + i] = 1;
-        }
-    }
-
-    if (tidx == 0) {
-        qk_gate_topk_idx[store_idx] = 1;
-        qk_gate_topk_idx[store_idx + seq_len - 1] = 1;
-        qk_gate_topk_idx[store_idx + seq_len - 2] = 1;
-    }
+  if (tidx == 0) {
+    qk_gate_topk_idx[store_idx] = 1;
+    qk_gate_topk_idx[store_idx + seq_len - 1] = 1;
+    qk_gate_topk_idx[store_idx + seq_len - 2] = 1;
+  }
 }
 
 template <int kBlockMaxN, int moba_block_size, typename T>
-void qk_gate_sort_decoder(
-        const T* qk_gate_weight,
-        int * qk_gate_topk_idx,
-        const int *decoder_seq_lens,
-        const int head_num,
-        const int kv_head_num,
-        const int batch_size,
-        const int top_k_left,
-        const int top_k_right,
-        const int use_moba_seq_limit,
-        cudaStream_t stream) {
+void qk_gate_sort_decoder(const T* qk_gate_weight,
+                          int* qk_gate_topk_idx,
+                          const int* decoder_seq_lens,
+                          const int head_num,
+                          const int kv_head_num,
+                          const int batch_size,
+                          const int top_k_left,
+                          const int top_k_right,
+                          const int use_moba_seq_limit,
+                          cudaStream_t stream) {
+  const int gqa_group_size = head_num / kv_head_num;
+  constexpr int kPackSize = 16 / sizeof(T);
+  const int knthreads = kBlockMaxN / kPackSize;
+  dim3 grid_dims;
+  grid_dims.x = batch_size;
+  grid_dims.y = head_num;
+  const int searchtimes = 6;
 
-    const int gqa_group_size = head_num / kv_head_num;
-    constexpr int kPackSize = 16 / sizeof(T);
-    const int knthreads = kBlockMaxN / kPackSize;
-    dim3 grid_dims;
-    grid_dims.x = batch_size;
-    grid_dims.y = head_num;
-    const int searchtimes = 6;
+  constexpr auto kernel = qk_gate_sort_decoder_kernel<T,
+                                                      knthreads,
+                                                      moba_block_size,
+                                                      kBlockMaxN,
+                                                      searchtimes>;
 
-    constexpr auto kernel = qk_gate_sort_decoder_kernel<T, knthreads, moba_block_size, kBlockMaxN, searchtimes>;
-
-    kernel<<<grid_dims, knthreads, 0, 0>>>(
-        qk_gate_weight,
-        qk_gate_topk_idx,
-        decoder_seq_lens,
-        head_num,
-        kv_head_num,
-        gqa_group_size,
-        top_k_left,
-        top_k_right,
-        use_moba_seq_limit);
+  kernel<<<grid_dims, knthreads, 0, 0>>>(qk_gate_weight,
+                                         qk_gate_topk_idx,
+                                         decoder_seq_lens,
+                                         head_num,
+                                         kv_head_num,
+                                         gqa_group_size,
+                                         top_k_left,
+                                         top_k_right,
+                                         use_moba_seq_limit);
 }
 
-
 template <typename T>
 std::vector<paddle::Tensor> DispatchQkSortDecoder(
-        const paddle::Tensor& qk_gate_weight,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const int head_num,
-        const int kv_head_num,
-        const int top_k_left,
-        const int top_k_right,
-        const int use_moba_seq_limit) {
+    const paddle::Tensor& qk_gate_weight,
+    const paddle::Tensor& seq_len_encoder,
+    const paddle::Tensor& seq_len_decoder,
+    const int head_num,
+    const int kv_head_num,
+    const int top_k_left,
+    const int top_k_right,
+    const int use_moba_seq_limit) {
+  constexpr int kMobaBlockSize = 128;
+  constexpr int kMaxN = 1024;
 
-    constexpr int kMobaBlockSize = 128;
-    constexpr int kMaxN = 1024;
+  const int batch_size = seq_len_decoder.dims()[0];
+  paddle::Tensor qk_gate_topk_idx =
+      paddle::empty({batch_size, kv_head_num, kMaxN},
+                    paddle::DataType::INT32,
+                    qk_gate_weight.place());
 
-    const int batch_size = seq_len_decoder.dims()[0];
-    paddle::Tensor qk_gate_topk_idx = paddle::empty({batch_size, kv_head_num, kMaxN}, paddle::DataType::INT32, qk_gate_weight.place());
+  qk_gate_sort_decoder<kMaxN, kMobaBlockSize, T>(qk_gate_weight.data<T>(),
+                                                 qk_gate_topk_idx.data<int>(),
+                                                 seq_len_decoder.data<int>(),
+                                                 head_num,
+                                                 kv_head_num,
+                                                 batch_size,
+                                                 top_k_left,
+                                                 top_k_right,
+                                                 use_moba_seq_limit,
+                                                 qk_gate_weight.stream());
 
-    qk_gate_sort_decoder<kMaxN, kMobaBlockSize, T>(
-        qk_gate_weight.data<T>(),
-        qk_gate_topk_idx.data<int>(),
-        seq_len_decoder.data<int>(),
-        head_num,
-        kv_head_num,
-        batch_size,
-        top_k_left,
-        top_k_right,
-        use_moba_seq_limit,
-        qk_gate_weight.stream()
-    );
-
-    return {qk_gate_topk_idx};
+  return {qk_gate_topk_idx};
 }
 
-std::vector<paddle::Tensor> QkSortDecoder(
-        const paddle::Tensor& qk_gate_weight,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const int head_num,
-        const int kv_head_num,
-        const int top_k_left,
-        const int top_k_right,
-        const int use_moba_seq_limit) {
-
-    if (qk_gate_weight.dtype() == paddle::DataType::FLOAT16) {
-        return std::move(
-            DispatchQkSortDecoder<phi::dtype::float16>(
-                qk_gate_weight,
-                seq_len_encoder,
-                seq_len_decoder,
-                head_num,
-                kv_head_num,
-                top_k_left,
-                top_k_right,
-                use_moba_seq_limit)
-        );
-    } else if (qk_gate_weight.dtype() == paddle::DataType::BFLOAT16) {
-        return std::move(
-            DispatchQkSortDecoder<phi::dtype::bfloat16>(
-                qk_gate_weight,
-                seq_len_encoder,
-                seq_len_decoder,
-                head_num,
-                kv_head_num,
-                top_k_left,
-                top_k_right,
-                use_moba_seq_limit)
-        );
-    }
+std::vector<paddle::Tensor> QkSortDecoder(const paddle::Tensor& qk_gate_weight,
+                                          const paddle::Tensor& seq_len_encoder,
+                                          const paddle::Tensor& seq_len_decoder,
+                                          const int head_num,
+                                          const int kv_head_num,
+                                          const int top_k_left,
+                                          const int top_k_right,
+                                          const int use_moba_seq_limit) {
+  if (qk_gate_weight.dtype() == paddle::DataType::FLOAT16) {
+    return std::move(
+        DispatchQkSortDecoder<phi::dtype::float16>(qk_gate_weight,
+                                                   seq_len_encoder,
+                                                   seq_len_decoder,
+                                                   head_num,
+                                                   kv_head_num,
+                                                   top_k_left,
+                                                   top_k_right,
+                                                   use_moba_seq_limit));
+  } else if (qk_gate_weight.dtype() == paddle::DataType::BFLOAT16) {
+    return std::move(
+        DispatchQkSortDecoder<phi::dtype::bfloat16>(qk_gate_weight,
+                                                    seq_len_encoder,
+                                                    seq_len_decoder,
+                                                    head_num,
+                                                    kv_head_num,
+                                                    top_k_left,
+                                                    top_k_right,
+                                                    use_moba_seq_limit));
+  }
 }
 
 PD_BUILD_STATIC_OP(moba_qk_sort_decoder)
-    .Inputs({
-        "qk_gate_weight",
-        "seq_len_encoder",
-        "seq_len_decoder"})
-    .Attrs({
-        "head_num: int",
-        "kv_head_num: int",
-        "top_k_left: int",
-        "top_k_right: int",
-        "use_moba_seq_limit: int"})
+    .Inputs({"qk_gate_weight", "seq_len_encoder", "seq_len_decoder"})
+    .Attrs({"head_num: int",
+            "kv_head_num: int",
+            "top_k_left: int",
+            "top_k_right: int",
+            "use_moba_seq_limit: int"})
     .Outputs({"qk_gate_topk_idx"})
     .SetKernelFn(PD_KERNEL(QkSortDecoder));
diff --git a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/kernel_traits.h b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/kernel_traits.h
index 2673708e67..cd09cf9f78 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/kernel_traits.h
+++ b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/kernel_traits.h
@@ -13,7 +13,8 @@
 // limitations under the License.
 
 /******************************************************************************
- * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar,
+ *Pradeep Ramani, Tri Dao.
  ******************************************************************************/
 
 #pragma once
@@ -30,114 +31,161 @@
 using namespace cute;
 
 struct moba_encoder_attn_params {
-    void *__restrict__ q_ptr;
-    void *__restrict__ k_ptr;
-    void *__restrict__ v_ptr;
-    void * __restrict__ o_ptr;
-    int * __restrict__ cu_seq_q;
-    int * __restrict__ cu_seq_k;
-    int * __restrict__ qk_gate_topk_idx;
-    int * __restrict__ seq_len_encoder;
-    int * __restrict__ cu_seq_q_pack;
-    int head_num;
-    int kv_head_num;
-    int max_seq_q;
-    int max_seq_k;
-    int batch_size;
-    int gqa_group_size;
-    float scale_softmax_log2;
-    int use_moba_seq_limit;
+  void *__restrict__ q_ptr;
+  void *__restrict__ k_ptr;
+  void *__restrict__ v_ptr;
+  void *__restrict__ o_ptr;
+  int *__restrict__ cu_seq_q;
+  int *__restrict__ cu_seq_k;
+  int *__restrict__ qk_gate_topk_idx;
+  int *__restrict__ seq_len_encoder;
+  int *__restrict__ cu_seq_q_pack;
+  int head_num;
+  int kv_head_num;
+  int max_seq_q;
+  int max_seq_k;
+  int batch_size;
+  int gqa_group_size;
+  float scale_softmax_log2;
+  int use_moba_seq_limit;
 };
 
-template <int kStages, class Gemm1Type, class Gemm2Type, class OutputType, class SmemLayoutQ,
-          class SmemLayoutK, class SmemLayoutV, class SmemLayoutO>
+template <int kStages,
+          class Gemm1Type,
+          class Gemm2Type,
+          class OutputType,
+          class SmemLayoutQ,
+          class SmemLayoutK,
+          class SmemLayoutV,
+          class SmemLayoutO>
 struct SharedStorageQKVO {
-    cute::array_aligned<Gemm1Type, cute::cosize_v<SmemLayoutQ>> smem_q;
-    cute::array_aligned<Gemm1Type, cute::cosize_v<SmemLayoutK>> smem_k;
-    union {
-        cute::array_aligned<Gemm2Type, cute::cosize_v<SmemLayoutV>> smem_v;
-        cute::array_aligned<OutputType, cute::cosize_v<SmemLayoutO>> smem_o;
-    };
-    struct {
-        cutlass::arch::ClusterTransactionBarrier barrier_Q;
-        typename cutlass::PipelineTmaAsync<kStages>::SharedStorage pipeline_k;
-        typename cutlass::PipelineTmaAsync<kStages>::SharedStorage pipeline_v;
-    };
+  cute::array_aligned<Gemm1Type, cute::cosize_v<SmemLayoutQ>> smem_q;
+  cute::array_aligned<Gemm1Type, cute::cosize_v<SmemLayoutK>> smem_k;
+  union {
+    cute::array_aligned<Gemm2Type, cute::cosize_v<SmemLayoutV>> smem_v;
+    cute::array_aligned<OutputType, cute::cosize_v<SmemLayoutO>> smem_o;
+  };
+  struct {
+    cutlass::arch::ClusterTransactionBarrier barrier_Q;
+    typename cutlass::PipelineTmaAsync<kStages>::SharedStorage pipeline_k;
+    typename cutlass::PipelineTmaAsync<kStages>::SharedStorage pipeline_v;
+  };
 };
 
-template<int kHeadDim_, int kBlockM_, int kBlockN_, int kNWarps_, int kStages_, int kMaxN_, bool UseMoba_, typename elem_type=cutlass::half_t>
+template <int kHeadDim_,
+          int kBlockM_,
+          int kBlockN_,
+          int kNWarps_,
+          int kStages_,
+          int kMaxN_,
+          bool UseMoba_,
+          typename elem_type = cutlass::half_t>
 struct moba_encoder_attn_kernel_traits {
-    using Element = elem_type;
-    using ElementAccum = float;
-    using index_t = int32_t;
+  using Element = elem_type;
+  using ElementAccum = float;
+  using index_t = int32_t;
 
-    static constexpr int kNWarps = kNWarps_;
-    static constexpr int kNThreads = kNWarps * cutlass::NumThreadsPerWarp;
+  static constexpr int kNWarps = kNWarps_;
+  static constexpr int kNThreads = kNWarps * cutlass::NumThreadsPerWarp;
 
-    static constexpr int UseMoba = UseMoba_;
+  static constexpr int UseMoba = UseMoba_;
 
-    static constexpr int kBlockM = kBlockM_;
-    static constexpr int kBlockN = kBlockN_;
-    static constexpr int kHeadDim = kHeadDim_;
-    static constexpr int kMaxN = kMaxN_;
-    static_assert(kHeadDim % 32 == 0);
-    using TileShape_MNK = Shape<Int<kBlockM>, Int<kBlockN>, Int<kHeadDim>>;
-    using ClusterShape_MNK = Shape<Int<1>, Int<1>, Int<1>>;
-    static constexpr int kStages = kStages_;
+  static constexpr int kBlockM = kBlockM_;
+  static constexpr int kBlockN = kBlockN_;
+  static constexpr int kHeadDim = kHeadDim_;
+  static constexpr int kMaxN = kMaxN_;
+  static_assert(kHeadDim % 32 == 0);
+  using TileShape_MNK = Shape<Int<kBlockM>, Int<kBlockN>, Int<kHeadDim>>;
+  using ClusterShape_MNK = Shape<Int<1>, Int<1>, Int<1>>;
+  static constexpr int kStages = kStages_;
 
-    using AtomLayoutMNK = Layout<Shape<Int<kBlockM / 64>, _1, _1>>;
-    using TiledMma0 = decltype(cute::make_tiled_mma(
-        cute::GMMA::ss_op_selector<Element, Element, ElementAccum, TileShape_MNK>(),
-        AtomLayoutMNK{}));
-    using TiledMma1 = decltype(cute::make_tiled_mma(
-        cute::GMMA::rs_op_selector<Element, Element, ElementAccum, decltype(select<0, 2, 1>(TileShape_MNK{})),
-            GMMA::Major::K, GMMA::Major::MN>(),
-        AtomLayoutMNK{}));
+  using AtomLayoutMNK = Layout<Shape<Int<kBlockM / 64>, _1, _1>>;
+  using TiledMma0 = decltype(cute::make_tiled_mma(
+      cute::GMMA::
+          ss_op_selector<Element, Element, ElementAccum, TileShape_MNK>(),
+      AtomLayoutMNK{}));
+  using TiledMma1 = decltype(cute::make_tiled_mma(
+      cute::GMMA::rs_op_selector<Element,
+                                 Element,
+                                 ElementAccum,
+                                 decltype(select<0, 2, 1>(TileShape_MNK{})),
+                                 GMMA::Major::K,
+                                 GMMA::Major::MN>(),
+      AtomLayoutMNK{}));
 
-    using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::ss_smem_selector<GMMA::Major::K, Element,
-        decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
-    using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_MNK{})));
+  using SmemLayoutAtomQ =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               Element,
+               decltype(cute::get<0>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutQ =
+      decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_MNK{})));
 
-    using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::ss_smem_selector<GMMA::Major::K, Element,
-        decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
-    using SmemLayoutK =
-        decltype(tile_to_shape(SmemLayoutAtomK{},
-                 make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{}), Int<kStages>{})));
+  using SmemLayoutAtomK =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               Element,
+               decltype(cute::get<1>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutK =
+      decltype(tile_to_shape(SmemLayoutAtomK{},
+                             make_shape(shape<1>(TileShape_MNK{}),
+                                        shape<2>(TileShape_MNK{}),
+                                        Int<kStages>{})));
 
-    using SmemLayoutAtomV = decltype(cutlass::gemm::collective::detail::ss_smem_selector<GMMA::Major::K, Element,
-        decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
-    using SmemLayoutV =
-        decltype(tile_to_shape(SmemLayoutAtomV{},
-                 make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{}), Int<kStages>{})));
+  using SmemLayoutAtomV =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               Element,
+               decltype(cute::get<1>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutV =
+      decltype(tile_to_shape(SmemLayoutAtomV{},
+                             make_shape(shape<1>(TileShape_MNK{}),
+                                        shape<2>(TileShape_MNK{}),
+                                        Int<kStages>{})));
 
-    using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector<GMMA::Major::K, Element,
-        decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
-    using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 2>(TileShape_MNK{})));
+  using SmemLayoutAtomO =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               Element,
+               decltype(cute::get<0>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutO =
+      decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 2>(TileShape_MNK{})));
 
-    using SmemCopyAtomQ = Copy_Atom<cute::SM75_U32x4_LDSM_N, Element>;
-    using SmemCopyAtomO = Copy_Atom<cute::SM90_U32x4_STSM_N, Element>;
+  using SmemCopyAtomQ = Copy_Atom<cute::SM75_U32x4_LDSM_N, Element>;
+  using SmemCopyAtomO = Copy_Atom<cute::SM90_U32x4_STSM_N, Element>;
 
-    using SharedStorage = SharedStorageQKVO<kStages, Element, Element, Element, SmemLayoutQ, SmemLayoutK, SmemLayoutV, SmemLayoutO>;
+  using SharedStorage = SharedStorageQKVO<kStages,
+                                          Element,
+                                          Element,
+                                          Element,
+                                          SmemLayoutQ,
+                                          SmemLayoutK,
+                                          SmemLayoutV,
+                                          SmemLayoutO>;
 
-    static constexpr int NumProducerThreads = cutlass::NumThreadsPerWarpGroup;
-    static constexpr int NumMmaThreads = kNThreads - NumProducerThreads;
-    static constexpr int kNumVecElem = ceil_div(128, sizeof_bits_v<Element>);
-    static constexpr int kNumThreadsPerRow = kHeadDim / kNumVecElem;
-    static_assert(NumMmaThreads % kNumThreadsPerRow == 0);
-    static constexpr int kNumRows = NumMmaThreads / kNumThreadsPerRow;
-    using TiledCopyOAtom = cute::Copy_Atom<cute::UniversalCopy<cutlass::uint128_t>, Element>;
-    using TiledCopyOThrLayout = decltype(cute::make_layout(
-        cute::make_shape(Int<kNumRows>{}, Int<kNumThreadsPerRow>{}),
-        LayoutRight{}));
-    using TiledCopyOValLayout = decltype(cute::make_layout(
-        cute::make_shape(_1{}, Int<kNumVecElem>{}),
-        LayoutRight{}));
-    using GmemTiledCopyO = decltype(make_tiled_copy(
-        TiledCopyOAtom{},
-        TiledCopyOThrLayout{}, // Thr layout
-        TiledCopyOValLayout{} // Val layout
-    ));
+  static constexpr int NumProducerThreads = cutlass::NumThreadsPerWarpGroup;
+  static constexpr int NumMmaThreads = kNThreads - NumProducerThreads;
+  static constexpr int kNumVecElem = ceil_div(128, sizeof_bits_v<Element>);
+  static constexpr int kNumThreadsPerRow = kHeadDim / kNumVecElem;
+  static_assert(NumMmaThreads % kNumThreadsPerRow == 0);
+  static constexpr int kNumRows = NumMmaThreads / kNumThreadsPerRow;
+  using TiledCopyOAtom =
+      cute::Copy_Atom<cute::UniversalCopy<cutlass::uint128_t>, Element>;
+  using TiledCopyOThrLayout = decltype(cute::make_layout(
+      cute::make_shape(Int<kNumRows>{}, Int<kNumThreadsPerRow>{}),
+      LayoutRight{}));
+  using TiledCopyOValLayout = decltype(cute::make_layout(
+      cute::make_shape(_1{}, Int<kNumVecElem>{}), LayoutRight{}));
+  using GmemTiledCopyO =
+      decltype(make_tiled_copy(TiledCopyOAtom{},
+                               TiledCopyOThrLayout{},  // Thr layout
+                               TiledCopyOValLayout{}   // Val layout
+                               ));
 
-    using MainloopPipeline = typename cutlass::PipelineTmaAsync<kStages>;
-    using PipelineState = typename cutlass::PipelineState<kStages>;
+  using MainloopPipeline = typename cutlass::PipelineTmaAsync<kStages>;
+  using PipelineState = typename cutlass::PipelineState<kStages>;
 };
diff --git a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/mainloop_attn.hpp b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/mainloop_attn.hpp
index 84d41af838..f2ae99c776 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/mainloop_attn.hpp
+++ b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/mainloop_attn.hpp
@@ -13,7 +13,8 @@
 // limitations under the License.
 
 /******************************************************************************
- * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar,
+ *Pradeep Ramani, Tri Dao.
  ******************************************************************************/
 
 #pragma once
@@ -31,443 +32,550 @@
 using namespace cute;
 
 enum class AttnNamedBarriers {
-    QueryEmpty = 0,
-    ValueEmpty = 1,
-    TileCountSmemEmpty = 2,
-    TileCountSmemFull = 3,
-    WarpSchedulerWG1 = 4,
-    WarpSchedulerWG2 = 5,
-    WarpSchedulerWG3 = 6,
+  QueryEmpty = 0,
+  ValueEmpty = 1,
+  TileCountSmemEmpty = 2,
+  TileCountSmemFull = 3,
+  WarpSchedulerWG1 = 4,
+  WarpSchedulerWG2 = 5,
+  WarpSchedulerWG3 = 6,
 };
 
-
-
 template <typename Ktraits>
 struct CollectiveMainloopAttn {
+  using Element = typename Ktraits::Element;
+  using TileShape_MNK = typename Ktraits::TileShape_MNK;
+  using ClusterShape = typename Ktraits::ClusterShape_MNK;
 
-    using Element = typename Ktraits::Element;
-    using TileShape_MNK = typename Ktraits::TileShape_MNK;
-    using ClusterShape = typename Ktraits::ClusterShape_MNK;
+  static constexpr int kStages = Ktraits::kStages;
+  static constexpr int kHeadDim = Ktraits::kHeadDim;
+  static constexpr int kBlockM = Ktraits::kBlockM;
+  static constexpr int kBlockN = Ktraits::kBlockN;
 
-    static constexpr int kStages = Ktraits::kStages;
-    static constexpr int kHeadDim = Ktraits::kHeadDim;
-    static constexpr int kBlockM = Ktraits::kBlockM;
-    static constexpr int kBlockN = Ktraits::kBlockN;
+  using ShapeT = cute::Shape<int32_t, int32_t, int32_t>;
+  using StrideT = cute::Shape<int32_t, _1, int32_t>;
+  using LayoutT = cute::Layout<ShapeT, StrideT>;
 
-    using ShapeT = cute::Shape<int32_t, int32_t, int32_t>;
-    using StrideT = cute::Shape<int32_t, _1, int32_t>;
-    using LayoutT = cute::Layout<ShapeT, StrideT>;
+  using GmemTiledCopyQ = cute::SM90_TMA_LOAD;
+  using GmemTiledCopyKV =
+      decltype(cutlass::gemm::collective::detail::
+                   sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape{})));
+  using GmemTiledCopyO = typename Ktraits::GmemTiledCopyO;
 
+  using SmemLayoutAtomQ =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               Element,
+               decltype(cute::get<0>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutQ =
+      decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_MNK{})));
 
-    using GmemTiledCopyQ = cute::SM90_TMA_LOAD;
-    using GmemTiledCopyKV = decltype(cutlass::gemm::collective::detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape{})));
-    using GmemTiledCopyO = typename Ktraits::GmemTiledCopyO;
+  using SmemLayoutAtomK =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               Element,
+               decltype(cute::get<1>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutK =
+      decltype(tile_to_shape(SmemLayoutAtomK{},
+                             make_shape(shape<1>(TileShape_MNK{}),
+                                        shape<2>(TileShape_MNK{}),
+                                        Int<kStages>{})));
+  using SmemLayoutV = SmemLayoutK;
+  // Note this is the transpose in terms of the view, not in terms of memory.
+  using SmemLayoutVt = decltype(cute::composition(
+      SmemLayoutV{},
+      make_layout(
+          make_shape(
+              get<2>(TileShape_MNK{}), get<1>(TileShape_MNK{}), Int<kStages>{}),
+          make_stride(get<1>(TileShape_MNK{}),
+                      _1{},
+                      Int<size(SmemLayoutV{}(_, _, _0{}))>{}))));
+  using SmemLayoutO = typename Ktraits::SmemLayoutO;
+  using SmemCopyAtomO = typename Ktraits::SmemCopyAtomO;
 
-    using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::ss_smem_selector<GMMA::Major::K, Element,
-        decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
-    using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_MNK{})));
+  using TMA_Q = decltype(make_tma_copy(
+      GmemTiledCopyQ{},
+      make_tensor(make_gmem_ptr(static_cast<Element const*>(nullptr)),
+                  repeat_like(StrideT{}, int32_t(0)),
+                  StrideT{}),
+      SmemLayoutQ{},
+      select<0, 2>(TileShape_MNK{}),
+      _1{}));  // no mcast for Q
 
-    using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::ss_smem_selector<GMMA::Major::K, Element,
-        decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
-    using SmemLayoutK =
-        decltype(tile_to_shape(SmemLayoutAtomK{},
-                 make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{}), Int<kStages>{})));
-    using SmemLayoutV = SmemLayoutK;
-    // Note this is the transpose in terms of the view, not in terms of memory.
-    using SmemLayoutVt =
-        decltype(cute::composition(SmemLayoutV{},
-                                   make_layout(make_shape(get<2>(TileShape_MNK{}), get<1>(TileShape_MNK{}), Int<kStages>{}),
-                                               make_stride(get<1>(TileShape_MNK{}), _1{}, Int<size(SmemLayoutV{}(_, _, _0{}))>{}))));
-    using SmemLayoutO = typename Ktraits::SmemLayoutO;
-    using SmemCopyAtomO = typename Ktraits::SmemCopyAtomO;
+  using TMA_KV = decltype(make_tma_copy(
+      GmemTiledCopyKV{},
+      make_tensor(make_gmem_ptr(static_cast<Element const*>(nullptr)),
+                  repeat_like(StrideT{}, int32_t(0)),
+                  StrideT{}),
+      take<0, 2>(SmemLayoutK{}),
+      select<1, 2>(TileShape_MNK{}),
+      size<0>(ClusterShape{})));  // mcast along M mode for this N load, if any
 
-    using TMA_Q = decltype(make_tma_copy(
-        GmemTiledCopyQ{},
-        make_tensor(
-            make_gmem_ptr(static_cast<Element const*>(nullptr)),
-            repeat_like(StrideT{}, int32_t(0)),
-            StrideT{}
-        ),
-        SmemLayoutQ{},
-        select<0, 2>(TileShape_MNK{}),
-        _1{}));  // no mcast for Q
+  static constexpr int NumMmaThreads = size(typename Ktraits::TiledMma0{});
+  using MainloopPipeline = typename Ktraits::MainloopPipeline;
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState = typename MainloopPipeline::PipelineState;
 
-    using TMA_KV = decltype(make_tma_copy(
+  // Set the bytes transferred in this TMA transaction (may involve multiple
+  // issues)
+  static constexpr uint32_t TmaTransactionBytesQ = static_cast<uint32_t>(
+      size(SmemLayoutQ{}) * cutlass::sizeof_bits_v<Element> / 8);
+  static constexpr uint32_t TmaTransactionBytesK = static_cast<uint32_t>(
+      size(take<0, 2>(SmemLayoutK{})) * cutlass::sizeof_bits_v<Element> / 8);
+
+  static constexpr bool UseSchedulerBarrier = kHeadDim <= 128;
+
+  // Host side kernel arguments
+  struct Arguments {
+    Element const* ptr_Q;
+    LayoutT layout_Q;
+    Element const* ptr_K;
+    LayoutT layout_K;
+    Element const* ptr_V;
+    LayoutT layout_V;
+    float const softmax_scale_log2;
+  };
+
+  // Device side kernel params
+  struct Params {
+    LayoutT layout_Q;
+    LayoutT layout_K;
+    LayoutT layout_V;
+    cutlass::FastDivmod qhead_per_khead_divmod;
+    TMA_Q tma_load_Q;
+    TMA_KV tma_load_K, tma_load_V;
+    float const softmax_scale_log2;
+  };
+
+  static Params to_underlying_arguments(Arguments const& args) {
+    Tensor mQ = make_tensor(make_gmem_ptr(args.ptr_Q), args.layout_Q);
+    TMA_Q tma_load_Q = make_tma_copy(GmemTiledCopyQ{},
+                                     mQ,
+                                     SmemLayoutQ{},
+                                     select<0, 2>(TileShape_MNK{}),
+                                     _1{});  // no mcast for Q
+    Tensor mK = make_tensor(make_gmem_ptr(args.ptr_K), args.layout_K);
+    TMA_KV tma_load_K = make_tma_copy(
         GmemTiledCopyKV{},
-        make_tensor(
-            make_gmem_ptr(static_cast<Element const*>(nullptr)),
-            repeat_like(StrideT{}, int32_t(0)),
-            StrideT{}
-        ),
-        take<0, 2>(SmemLayoutK{}),
+        mK,
+        SmemLayoutK{}(_, _, _0{}),
         select<1, 2>(TileShape_MNK{}),
-        size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any
+        size<0>(ClusterShape{}));  // mcast along M mode for this N load, if any
+    Tensor mV = make_tensor(make_gmem_ptr(args.ptr_V), args.layout_V);
+    TMA_KV tma_load_V = make_tma_copy(
+        GmemTiledCopyKV{},
+        mV,
+        SmemLayoutV{}(_, _, _0{}),
+        select<1, 2>(TileShape_MNK{}),
+        size<0>(ClusterShape{}));  // mcast along M mode for this N load, if any
+    return {args.layout_Q,
+            args.layout_K,
+            args.layout_V,
+            cutlass::FastDivmod(cute::ceil_div(get<2>(args.layout_Q.shape()),
+                                               get<2>(args.layout_K.shape()))),
+            tma_load_Q,
+            tma_load_K,
+            tma_load_V,
+            args.softmax_scale_log2};
+  }
 
-    static constexpr int NumMmaThreads = size(typename Ktraits::TiledMma0{});
-    using MainloopPipeline = typename Ktraits::MainloopPipeline;
-    using PipelineParams = typename MainloopPipeline::Params;
-    using PipelineState = typename MainloopPipeline::PipelineState;
+  /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best
+  /// performance
+  CUTLASS_DEVICE
+  static void prefetch_tma_descriptors(Params const& mainloop_params) {
+    cute::prefetch_tma_descriptor(
+        mainloop_params.tma_load_Q.get_tma_descriptor());
+    cute::prefetch_tma_descriptor(
+        mainloop_params.tma_load_K.get_tma_descriptor());
+    cute::prefetch_tma_descriptor(
+        mainloop_params.tma_load_V.get_tma_descriptor());
+  }
 
-    // Set the bytes transferred in this TMA transaction (may involve multiple issues)
-    static constexpr uint32_t TmaTransactionBytesQ = static_cast<uint32_t>(size(SmemLayoutQ{}) * cutlass::sizeof_bits_v<Element> / 8);
-    static constexpr uint32_t TmaTransactionBytesK = static_cast<uint32_t>(size(take<0, 2>(SmemLayoutK{})) * cutlass::sizeof_bits_v<Element> / 8);
+  template <typename MTensor, typename Shape>
+  CUTLASS_DEVICE auto get_local_tile_tensor(const MTensor& m_tensor,
+                                            const Shape& tile_shape,
+                                            const int* cu_seq_len,
+                                            const int bidh,
+                                            const int bidb,
+                                            const int actual_seq_len) const {
+    auto g_offset = local_tile(m_tensor(_, _, bidh),
+                               cute::make_shape(1, get<1>(tile_shape)),
+                               make_coord(cu_seq_len[bidb], _0{}));
+    auto g_sequence = make_tensor(
+        g_offset.data(),
+        make_layout(cute::make_shape(actual_seq_len, get<1>(tile_shape)),
+                    g_offset.stride()));
+    auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_, _0{}));
+    return g_tensor;
+  }
 
-    static constexpr bool UseSchedulerBarrier = kHeadDim <= 128;
+  template <bool UseMoba, typename SharedStorage>
+  CUTLASS_DEVICE void load(Params const& mainloop_params,
+                           MainloopPipeline pipeline_k,
+                           MainloopPipeline pipeline_v,
+                           PipelineState& smem_pipe_write_k,
+                           PipelineState& smem_pipe_write_v,
+                           SharedStorage& shared_storage,
+                           const int* qk_gate_topk_idx,
+                           const int n_block_max,
+                           const int m_block,
+                           const int bidh,
+                           const int bidb,
+                           const int* cu_seq_q,
+                           const int* cu_seq_k,
+                           const int seq_len_q,
+                           const int seq_len_k) {
+    Tensor sQ =
+        make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+    Tensor sK =
+        make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{});
+    Tensor sV =
+        make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutV{});
 
-    // Host side kernel arguments
-    struct Arguments {
-        Element const* ptr_Q;
-        LayoutT layout_Q;
-        Element const* ptr_K;
-        LayoutT layout_K;
-        Element const* ptr_V;
-        LayoutT layout_V;
-        float const softmax_scale_log2;
+    Tensor mQ = mainloop_params.tma_load_Q.get_tma_tensor(
+        mainloop_params.layout_Q.shape());
+    Tensor mK = mainloop_params.tma_load_K.get_tma_tensor(
+        mainloop_params.layout_K.shape());
+    Tensor mV = mainloop_params.tma_load_V.get_tma_tensor(
+        mainloop_params.layout_V.shape());
+    int bidh_kv = mainloop_params.qhead_per_khead_divmod.divide(bidh);
+
+    Tensor gQ = get_local_tile_tensor(
+        mQ, select<0, 2>(TileShape_MNK{}), cu_seq_q, bidh, bidb, seq_len_q)(
+        _, _, m_block);
+    Tensor gK = get_local_tile_tensor(
+        mK, select<1, 2>(TileShape_MNK{}), cu_seq_k, bidh_kv, bidb, seq_len_k);
+    Tensor gV = get_local_tile_tensor(
+        mV, select<1, 2>(TileShape_MNK{}), cu_seq_k, bidh_kv, bidb, seq_len_k);
+
+    Tensor sQ_x =
+        make_tensor(sQ.data(), make_layout(sQ.layout(), Layout<_1>{}));
+    Tensor gQ_x =
+        make_tensor(gQ.data(), make_layout(gQ.layout(), Layout<_1>{}));
+    auto [tQgQ, tQsQ] = tma_partition(mainloop_params.tma_load_Q,
+                                      _0{},
+                                      Layout<_1>{},
+                                      group_modes<0, 2>(sQ_x),
+                                      group_modes<0, 2>(gQ_x));
+    auto [tKgK, tKsK] = tma_partition(mainloop_params.tma_load_K,
+                                      _0{},
+                                      Layout<_1>{},
+                                      group_modes<0, 2>(sK),
+                                      group_modes<0, 2>(gK));
+    auto [tVgV, tVsV] = tma_partition(mainloop_params.tma_load_V,
+                                      _0{},
+                                      Layout<_1>{},
+                                      group_modes<0, 2>(sV),
+                                      group_modes<0, 2>(gV));
+
+    uint16_t mcast_mask_kv = 0;
+
+    int n_block = n_block_max - 1;
+
+    int lane_predicate = cute::elect_one_sync();
+
+    if (lane_predicate) {
+      shared_storage.barrier_Q.arrive_and_expect_tx(TmaTransactionBytesQ);
+      copy(mainloop_params.tma_load_Q.with(
+               reinterpret_cast<
+                   cutlass::arch::ClusterTransactionBarrier::ValueType&>(
+                   shared_storage.barrier_Q),
+               0 /*mcast_mask*/),
+           tQgQ,
+           tQsQ);
+    }
+
+    if (lane_predicate) {
+      pipeline_k.producer_acquire(smem_pipe_write_k);
+      copy(mainloop_params.tma_load_K.with(
+               *pipeline_k.producer_get_barrier(smem_pipe_write_k),
+               mcast_mask_kv),
+           tKgK(_, n_block),
+           tKsK(_, smem_pipe_write_k.index()));
+      ++smem_pipe_write_k;
+    }
+
+    if (lane_predicate) {
+      int idx = 0;
+#pragma unroll 2
+      for (; n_block > 0;) {
+        pipeline_k.producer_acquire(smem_pipe_write_k);
+        int pre_idx = 1;
+        if constexpr (UseMoba) {
+          pre_idx = qk_gate_topk_idx[idx];
+        }
+        copy(mainloop_params.tma_load_K.with(
+                 *pipeline_k.producer_get_barrier(smem_pipe_write_k),
+                 mcast_mask_kv),
+             tKgK(_, n_block - pre_idx),
+             tKsK(_, smem_pipe_write_k.index()));
+
+        ++smem_pipe_write_k;
+        pipeline_v.producer_acquire(smem_pipe_write_v);
+        copy(mainloop_params.tma_load_V.with(
+                 *pipeline_v.producer_get_barrier(smem_pipe_write_v),
+                 mcast_mask_kv),
+             tVgV(_, n_block),
+             tVsV(_, smem_pipe_write_v.index()));
+        ++smem_pipe_write_v;
+        n_block -= pre_idx;
+        idx += 1;
+      }
+    }
+    if (lane_predicate) {
+      pipeline_v.producer_acquire(smem_pipe_write_v);
+      copy(mainloop_params.tma_load_V.with(
+               *pipeline_v.producer_get_barrier(smem_pipe_write_v),
+               mcast_mask_kv),
+           tVgV(_, n_block),
+           tVsV(_, smem_pipe_write_v.index()));
+      ++smem_pipe_write_v;
+    }
+  }
+
+  CUTLASS_DEVICE void warp_scheduler_barrier_sync() {
+    if constexpr (UseSchedulerBarrier) {
+      cutlass::arch::NamedBarrier::sync(
+          NumMmaThreads,
+          static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 +
+              cutlass::canonical_warp_group_idx() /*id*/);
+    }
+  }
+
+  CUTLASS_DEVICE void mma_init() {
+    if constexpr (!UseSchedulerBarrier) {
+      return;
+    }
+    static_assert(NumMmaThreads == 2 * cutlass::NumThreadsPerWarpGroup ||
+                  NumMmaThreads == 3 * cutlass::NumThreadsPerWarpGroup);
+    if (cutlass::canonical_warp_group_idx() > 1) {
+      cutlass::arch::NamedBarrier::arrive(
+          NumMmaThreads,
+          static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 + 1 /*id*/);
+    }
+    if constexpr (NumMmaThreads == 3 * cutlass::NumThreadsPerWarpGroup) {
+      if (cutlass::canonical_warp_group_idx() > 2) {
+        cutlass::arch::NamedBarrier::arrive(
+            NumMmaThreads,
+            static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 +
+                2 /*id*/);
+      }
+    }
+  }
+
+  CUTLASS_DEVICE void warp_scheduler_barrier_arrive() {
+    if constexpr (!UseSchedulerBarrier) {
+      return;
+    }
+    static_assert(NumMmaThreads == 2 * cutlass::NumThreadsPerWarpGroup ||
+                  NumMmaThreads == 3 * cutlass::NumThreadsPerWarpGroup);
+    if constexpr (NumMmaThreads == 2 * cutlass::NumThreadsPerWarpGroup) {
+      cutlass::arch::NamedBarrier::arrive(
+          NumMmaThreads,
+          static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 +
+              (3 - cutlass::canonical_warp_group_idx()) /*id*/);
+    } else {
+      cutlass::arch::NamedBarrier::arrive(
+          NumMmaThreads,
+          static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 +
+              (cutlass::canonical_warp_group_idx() <= 2
+                   ? cutlass::canonical_warp_group_idx() + 1
+                   : cutlass::canonical_warp_group_idx() + 1 - 3) /*id*/);
+      cutlass::arch::NamedBarrier::arrive(
+          NumMmaThreads,
+          static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 +
+              (cutlass::canonical_warp_group_idx() <= 1
+                   ? cutlass::canonical_warp_group_idx() + 2
+                   : cutlass::canonical_warp_group_idx() + 2 - 3) /*id*/);
+    }
+  }
+
+  template <bool UseMoba,
+            typename SharedStorage,
+            typename FrgTensorO,
+            typename Softmax>
+  CUTLASS_DEVICE void mma(Params const& mainloop_params,
+                          MainloopPipeline pipeline_k,
+                          MainloopPipeline pipeline_v,
+                          PipelineState& smem_pipe_read_k,
+                          PipelineState& smem_pipe_read_v,
+                          FrgTensorO& tOrO,
+                          Softmax& softmax,
+                          const int* qk_gate_topk_idx,
+                          const int n_block_max,
+                          const int thread_idx,
+                          const int m_block,
+                          const int seq_len_q,
+                          const int seq_len_k,
+                          SharedStorage& shared_storage) {
+    Tensor sQ =
+        make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+    Tensor sK =
+        make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{});
+    Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.data()),
+                             SmemLayoutVt{});
+
+    typename Ktraits::TiledMma0 tiled_mma0;
+    typename Ktraits::TiledMma1 tiled_mma1;
+    auto threadMma0 = tiled_mma0.get_thread_slice(thread_idx);
+    auto threadMma1 = tiled_mma1.get_thread_slice(thread_idx);
+
+    Tensor tSrQ = threadMma0.partition_fragment_A(sQ);
+    Tensor tSrK = threadMma0.partition_fragment_B(sK);
+    Tensor tOrV = threadMma1.partition_fragment_B(sVt);
+
+    auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
+      auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+      pipeline.consumer_wait(smem_pipe_read, barrier_token);
     };
 
-    // Device side kernel params
-    struct Params {
-        LayoutT layout_Q;
-        LayoutT layout_K;
-        LayoutT layout_V;
-        cutlass::FastDivmod qhead_per_khead_divmod;
-        TMA_Q tma_load_Q;
-        TMA_KV tma_load_K, tma_load_V;
-        float const softmax_scale_log2;
+    tiled_mma1.accumulate_ = GMMA::ScaleOut::Zero;
+
+    int n_block = n_block_max - 1;
+
+    cutlass::ConsumerToken barrier_token = static_cast<cutlass::BarrierStatus>(
+        shared_storage.barrier_Q.try_wait(0));
+    if (barrier_token == cutlass::BarrierStatus::WaitAgain) {
+      shared_storage.barrier_Q.wait(0);
+    }
+
+    Tensor tSrS =
+        partition_fragment_C(tiled_mma0, select<0, 1>(TileShape_MNK{}));
+    consumer_wait(pipeline_k, smem_pipe_read_k);
+    warp_scheduler_barrier_sync();
+    gemm</*zero_init=*/true, /*wg_wait=*/-1>(
+        tiled_mma0, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), tSrS);
+    warp_scheduler_barrier_arrive();
+    warpgroup_wait<0>();
+    pipeline_k.consumer_release(smem_pipe_read_k);
+    ++smem_pipe_read_k;
+
+    auto col_limit_causal = [&](int row, int n_block) {
+      return row + 1 + seq_len_k - n_block * kBlockN - seq_len_q +
+             m_block * kBlockM;
     };
-
-
-    static Params
-    to_underlying_arguments(Arguments const& args) {
-        Tensor mQ = make_tensor(make_gmem_ptr(args.ptr_Q), args.layout_Q);
-        TMA_Q tma_load_Q = make_tma_copy(
-            GmemTiledCopyQ{},
-            mQ,
-            SmemLayoutQ{},
-            select<0, 2>(TileShape_MNK{}),
-            _1{}); // no mcast for Q
-        Tensor mK = make_tensor(make_gmem_ptr(args.ptr_K), args.layout_K);
-        TMA_KV tma_load_K = make_tma_copy(
-            GmemTiledCopyKV{},
-            mK,
-            SmemLayoutK{}(_, _, _0{}),
-            select<1, 2>(TileShape_MNK{}),
-            size<0>(ClusterShape{})); // mcast along M mode for this N load, if any
-        Tensor mV = make_tensor(make_gmem_ptr(args.ptr_V), args.layout_V);
-        TMA_KV tma_load_V = make_tma_copy(
-            GmemTiledCopyKV{},
-            mV,
-            SmemLayoutV{}(_, _, _0{}),
-            select<1, 2>(TileShape_MNK{}),
-            size<0>(ClusterShape{})); // mcast along M mode for this N load, if any
-        return {args.layout_Q, args.layout_K, args.layout_V,
-                cutlass::FastDivmod(cute::ceil_div(get<2>(args.layout_Q.shape()), get<2>(args.layout_K.shape()))),
-                tma_load_Q, tma_load_K, tma_load_V,
-                args.softmax_scale_log2};
+    Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_MNK{}));
+    Tensor tScS = threadMma0.partition_C(cS);
+#pragma unroll
+    for (int i = 0; i < size(tSrS); ++i) {
+      if (int(get<1>(tScS(i))) >=
+          std::min(seq_len_k - n_block * kBlockN,
+                   col_limit_causal(int(get<0>(tScS(i))), n_block))) {
+        tSrS(i) = -INFINITY;
+      }
     }
 
-    /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
-    CUTLASS_DEVICE
-    static void prefetch_tma_descriptors(Params const& mainloop_params) {
-        cute::prefetch_tma_descriptor(mainloop_params.tma_load_Q.get_tma_descriptor());
-        cute::prefetch_tma_descriptor(mainloop_params.tma_load_K.get_tma_descriptor());
-        cute::prefetch_tma_descriptor(mainloop_params.tma_load_V.get_tma_descriptor());
+    softmax.template online_softmax</*Is_first=*/true>(
+        tSrS, mainloop_params.softmax_scale_log2);
+
+    Tensor tOrP = make_tensor(
+        convert_type<Element>(tSrS).data(),
+        convert_layout_acc_Aregs<typename Ktraits::TiledMma1>(tSrS.layout()));
+    Tensor scores_scale = make_fragment_like(softmax.row_max);
+    clear(scores_scale);
+
+    int idx = 0;
+#pragma unroll 2
+    for (; n_block > 0;) {
+      Tensor tSrS =
+          partition_fragment_C(tiled_mma0, select<0, 1>(TileShape_MNK{}));
+      consumer_wait(pipeline_k, smem_pipe_read_k);
+      warp_scheduler_barrier_sync();
+      gemm</*zero_init=*/true, /*wg_wait=*/-1>(
+          tiled_mma0, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), tSrS);
+      softmax.rescale_o(tOrO, scores_scale);
+      consumer_wait(pipeline_v, smem_pipe_read_v);
+      gemm</*zero_init=*/false, /*wg_wait=*/-1>(
+          tiled_mma1, tOrP, tOrV(_, _, _, smem_pipe_read_v.index()), tOrO);
+      warp_scheduler_barrier_arrive();
+      warpgroup_wait<1>();
+      pipeline_k.consumer_release(smem_pipe_read_k);  // release K
+      cute::copy(softmax.template max</*Is_first=*/false>(
+                     tSrS, mainloop_params.softmax_scale_log2),
+                 scores_scale);
+      softmax.template online_softmax</*Is_first=*/false>(
+          tSrS, mainloop_params.softmax_scale_log2);
+      warpgroup_wait<0>();
+      pipeline_v.consumer_release(smem_pipe_read_v);  // release V
+      ++smem_pipe_read_k;
+      ++smem_pipe_read_v;
+      cute::copy(
+          make_tensor(convert_type<Element>(tSrS).data(),
+                      convert_layout_acc_Aregs<typename Ktraits::TiledMma1>(
+                          tSrS.layout())),
+          tOrP);
+      if constexpr (UseMoba) {
+        n_block -= qk_gate_topk_idx[idx];
+        idx += 1;
+      } else {
+        n_block -= 1;
+      }
     }
 
-    template <typename MTensor, typename Shape>
-    CUTLASS_DEVICE auto get_local_tile_tensor(
-        const MTensor &m_tensor,
-        const Shape &tile_shape,
-        const int *cu_seq_len,
-        const int bidh,
-        const int bidb,
-        const int actual_seq_len) const {
-        auto g_offset = local_tile(
-            m_tensor(_, _, bidh),
-            cute::make_shape(1, get<1>(tile_shape)),
-            make_coord(cu_seq_len[bidb], _0{}));
-        auto g_sequence = make_tensor(
-            g_offset.data(),
-            make_layout(
-                cute::make_shape(actual_seq_len, get<1>(tile_shape)),
-                g_offset.stride()
-            ));
-        auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_, _0{}));
-        return g_tensor;
+    softmax.rescale_o(tOrO, scores_scale);
+    consumer_wait(pipeline_v, smem_pipe_read_v);
+    gemm</*zero_init=*/false, /*wg_wait=*/-1>(
+        tiled_mma1, tOrP, tOrV(_, _, _, smem_pipe_read_v.index()), tOrO);
+    cute::copy(softmax.finalize(mainloop_params.softmax_scale_log2),
+               scores_scale);
+    warpgroup_wait<0>();
+    pipeline_v.consumer_release(smem_pipe_read_v);
+    ++smem_pipe_read_v;
+
+    softmax.rescale_o(tOrO, scores_scale);
+  }
+
+  template <int NumMmaThreads,
+            typename SharedStorage,
+            typename FrgTensorO,
+            typename TiledMma,
+            typename T>
+  CUTLASS_DEVICE void store(Params const& mainloop_params,
+                            FrgTensorO const& tOrO,
+                            SharedStorage& shared_storage,
+                            TiledMma tiled_mma,
+                            int thread_idx,
+                            const int o_head_stride,
+                            const int real_seq,
+                            T* out_ptr) {
+    Tensor sO =
+        make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{});
+    auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma);
+    auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx);
+
+    Tensor tOrO_out = convert_type<Element>(tOrO);
+    Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out);
+    Tensor taccOsO = smem_thr_copy_O.partition_D(sO);
+
+    cutlass::arch::NamedBarrier::sync(
+        NumMmaThreads, static_cast<int>(AttnNamedBarriers::ValueEmpty) /*id*/);
+    cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);
+    cutlass::arch::fence_view_async_shared();  // ensure smem writes are visible
+                                               // to TMA
+    cutlass::arch::NamedBarrier::arrive(
+        NumMmaThreads + cutlass::NumThreadsPerWarp,
+        cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+
+    Tensor gO = make_tensor(make_gmem_ptr(out_ptr),
+                            Shape<Int<kBlockM>, Int<kHeadDim>>{},
+                            make_stride(o_head_stride, _1{}));
+
+    GmemTiledCopyO gmem_tiled_copy_O;
+    auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx);
+
+    Tensor tOsO = gmem_thr_copy_O.partition_S(sO);
+    Tensor tOgO = gmem_thr_copy_O.partition_D(gO);
+
+    Tensor cO = make_identity_tensor(Shape<Int<kBlockM>, Int<kHeadDim>>{});
+
+    Tensor tOcO = gmem_thr_copy_O.partition_S(cO);
+
+    if (real_seq >= kBlockM) {
+      copy<true>(gmem_tiled_copy_O, tOsO, tOgO, tOcO);
+    } else {
+      copy<false>(gmem_tiled_copy_O, tOsO, tOgO, tOcO, real_seq);
     }
-
-
-    template <bool UseMoba, typename SharedStorage>
-    CUTLASS_DEVICE void
-    load(Params const& mainloop_params,
-         MainloopPipeline pipeline_k,
-         MainloopPipeline pipeline_v,
-         PipelineState& smem_pipe_write_k,
-         PipelineState& smem_pipe_write_v,
-         SharedStorage &shared_storage,
-         const int *qk_gate_topk_idx,
-         const int n_block_max,
-         const int m_block,
-         const int bidh,
-         const int bidb,
-         const int *cu_seq_q,
-         const int *cu_seq_k,
-         const int seq_len_q,
-         const int seq_len_k) {
-
-        Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
-        Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{});
-        Tensor sV = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutV{});
-
-        Tensor mQ = mainloop_params.tma_load_Q.get_tma_tensor(mainloop_params.layout_Q.shape());
-        Tensor mK = mainloop_params.tma_load_K.get_tma_tensor(mainloop_params.layout_K.shape());
-        Tensor mV = mainloop_params.tma_load_V.get_tma_tensor(mainloop_params.layout_V.shape());
-        int bidh_kv = mainloop_params.qhead_per_khead_divmod.divide(bidh);
-
-        Tensor gQ = get_local_tile_tensor(
-            mQ, select<0, 2>(TileShape_MNK{}), cu_seq_q, bidh, bidb, seq_len_q)(_, _, m_block);
-        Tensor gK = get_local_tile_tensor(
-            mK, select<1, 2>(TileShape_MNK{}), cu_seq_k, bidh_kv, bidb, seq_len_k);
-        Tensor gV = get_local_tile_tensor(
-            mV, select<1, 2>(TileShape_MNK{}), cu_seq_k, bidh_kv, bidb, seq_len_k);
-
-        Tensor sQ_x = make_tensor(sQ.data(), make_layout(sQ.layout(), Layout<_1>{}));
-        Tensor gQ_x = make_tensor(gQ.data(), make_layout(gQ.layout(), Layout<_1>{}));
-        auto [tQgQ, tQsQ] = tma_partition(mainloop_params.tma_load_Q, _0{}, Layout<_1>{},group_modes<0, 2>(sQ_x), group_modes<0, 2>(gQ_x));
-        auto [tKgK, tKsK] = tma_partition(mainloop_params.tma_load_K, _0{}, Layout<_1>{},group_modes<0, 2>(sK), group_modes<0, 2>(gK));
-        auto [tVgV, tVsV] = tma_partition(mainloop_params.tma_load_V, _0{}, Layout<_1>{},group_modes<0, 2>(sV), group_modes<0, 2>(gV));
-
-        uint16_t mcast_mask_kv = 0;
-
-        int n_block = n_block_max - 1;
-
-        int lane_predicate = cute::elect_one_sync();
-
-        if (lane_predicate) {
-            shared_storage.barrier_Q.arrive_and_expect_tx(TmaTransactionBytesQ);
-            copy(mainloop_params.tma_load_Q.with(reinterpret_cast<cutlass::arch::ClusterTransactionBarrier::ValueType&>(shared_storage.barrier_Q), 0 /*mcast_mask*/), tQgQ, tQsQ);
-        }
-
-
-        if (lane_predicate) {
-            pipeline_k.producer_acquire(smem_pipe_write_k);
-            copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), mcast_mask_kv), tKgK(_, n_block), tKsK(_, smem_pipe_write_k.index()));
-            ++smem_pipe_write_k;
-        }
-
-        if (lane_predicate) {
-            int idx = 0;
-            #pragma unroll 2
-            for (; n_block > 0; ) {
-                pipeline_k.producer_acquire(smem_pipe_write_k);
-                int pre_idx = 1;
-                if constexpr (UseMoba) {
-                    pre_idx = qk_gate_topk_idx[idx];
-                }
-                copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), mcast_mask_kv), tKgK(_, n_block - pre_idx), tKsK(_, smem_pipe_write_k.index()));
-
-                ++smem_pipe_write_k;
-                pipeline_v.producer_acquire(smem_pipe_write_v);
-                copy(mainloop_params.tma_load_V.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), mcast_mask_kv), tVgV(_, n_block), tVsV(_, smem_pipe_write_v.index()));
-                ++smem_pipe_write_v;
-                n_block -= pre_idx;
-                idx += 1;
-            }
-        }
-        if (lane_predicate) {
-            pipeline_v.producer_acquire(smem_pipe_write_v);
-            copy(mainloop_params.tma_load_V.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), mcast_mask_kv), tVgV(_, n_block), tVsV(_, smem_pipe_write_v.index()));
-            ++smem_pipe_write_v;
-        }
-    }
-
-    CUTLASS_DEVICE void
-    warp_scheduler_barrier_sync() {
-        if constexpr (UseSchedulerBarrier) {
-            cutlass::arch::NamedBarrier::sync(NumMmaThreads, static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 + cutlass::canonical_warp_group_idx() /*id*/);
-        }
-    }
-
-    CUTLASS_DEVICE void
-    mma_init() {
-        if constexpr (!UseSchedulerBarrier) { return; }
-        static_assert(NumMmaThreads == 2 * cutlass::NumThreadsPerWarpGroup || NumMmaThreads == 3 * cutlass::NumThreadsPerWarpGroup);
-        if (cutlass::canonical_warp_group_idx() > 1) {
-            cutlass::arch::NamedBarrier::arrive(NumMmaThreads, static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 + 1 /*id*/);
-        }
-        if constexpr (NumMmaThreads == 3 * cutlass::NumThreadsPerWarpGroup) {
-            if (cutlass::canonical_warp_group_idx() > 2) {
-                cutlass::arch::NamedBarrier::arrive(NumMmaThreads, static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 + 2 /*id*/);
-            }
-        }
-
-    }
-
-    CUTLASS_DEVICE void
-    warp_scheduler_barrier_arrive() {
-        if constexpr (!UseSchedulerBarrier) { return; }
-        static_assert(NumMmaThreads == 2 * cutlass::NumThreadsPerWarpGroup || NumMmaThreads == 3 * cutlass::NumThreadsPerWarpGroup);
-        if constexpr (NumMmaThreads == 2 * cutlass::NumThreadsPerWarpGroup) {
-            cutlass::arch::NamedBarrier::arrive(NumMmaThreads, static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 + (3 - cutlass::canonical_warp_group_idx()) /*id*/);
-        } else {
-            cutlass::arch::NamedBarrier::arrive(NumMmaThreads, static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 + (cutlass::canonical_warp_group_idx() <= 2 ? cutlass::canonical_warp_group_idx() + 1 : cutlass::canonical_warp_group_idx() + 1 - 3)  /*id*/);
-            cutlass::arch::NamedBarrier::arrive(NumMmaThreads, static_cast<int>(AttnNamedBarriers::WarpSchedulerWG1) - 1 + (cutlass::canonical_warp_group_idx() <= 1 ? cutlass::canonical_warp_group_idx() + 2 : cutlass::canonical_warp_group_idx() + 2 - 3)  /*id*/);
-        }
-    }
-
-
-    template <bool UseMoba, typename SharedStorage, typename FrgTensorO, typename Softmax>
-    CUTLASS_DEVICE void
-    mma(Params const& mainloop_params,
-        MainloopPipeline pipeline_k,
-        MainloopPipeline pipeline_v,
-        PipelineState& smem_pipe_read_k,
-        PipelineState& smem_pipe_read_v,
-        FrgTensorO& tOrO,
-        Softmax& softmax,
-        const int *qk_gate_topk_idx,
-        const int n_block_max,
-        const int thread_idx,
-        const int m_block,
-        const int seq_len_q,
-        const int seq_len_k,
-        SharedStorage& shared_storage) {
-
-        Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
-        Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{});
-        Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutVt{});
-
-        typename Ktraits::TiledMma0 tiled_mma0;
-        typename Ktraits::TiledMma1 tiled_mma1;
-        auto threadMma0 = tiled_mma0.get_thread_slice(thread_idx);
-        auto threadMma1 = tiled_mma1.get_thread_slice(thread_idx);
-
-        Tensor tSrQ = threadMma0.partition_fragment_A(sQ);
-        Tensor tSrK = threadMma0.partition_fragment_B(sK);
-        Tensor tOrV = threadMma1.partition_fragment_B(sVt);
-
-        auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
-            auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
-            pipeline.consumer_wait(smem_pipe_read, barrier_token);
-        };
-
-        tiled_mma1.accumulate_ = GMMA::ScaleOut::Zero;
-
-        int n_block = n_block_max - 1;
-
-        cutlass::ConsumerToken barrier_token = static_cast<cutlass::BarrierStatus>(shared_storage.barrier_Q.try_wait(0));
-        if (barrier_token == cutlass::BarrierStatus::WaitAgain) { shared_storage.barrier_Q.wait(0); }
-
-        Tensor tSrS = partition_fragment_C(tiled_mma0, select<0, 1>(TileShape_MNK{}));
-        consumer_wait(pipeline_k, smem_pipe_read_k);
-        warp_scheduler_barrier_sync();
-        gemm</*zero_init=*/true, /*wg_wait=*/-1>(tiled_mma0, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), tSrS);
-        warp_scheduler_barrier_arrive();
-        warpgroup_wait<0>();
-        pipeline_k.consumer_release(smem_pipe_read_k);
-        ++smem_pipe_read_k;
-
-        auto col_limit_causal = [&](int row, int n_block) {
-            return row + 1 + seq_len_k - n_block * kBlockN - seq_len_q + m_block * kBlockM;
-        };
-        Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_MNK{}));
-        Tensor tScS = threadMma0.partition_C(cS);
-        #pragma unroll
-        for (int i = 0; i < size(tSrS); ++i) {
-            if (int(get<1>(tScS(i))) >=
-                std::min(seq_len_k - n_block * kBlockN, col_limit_causal(int(get<0>(tScS(i))), n_block))) {
-                tSrS(i) = -INFINITY;
-            }
-        }
-
-        softmax.template online_softmax</*Is_first=*/true>(tSrS, mainloop_params.softmax_scale_log2);
-
-        Tensor tOrP = make_tensor(convert_type<Element>(tSrS).data(), convert_layout_acc_Aregs<typename Ktraits::TiledMma1>(tSrS.layout()));
-        Tensor scores_scale = make_fragment_like(softmax.row_max);
-        clear(scores_scale);
-
-        int idx = 0;
-        #pragma unroll 2
-        for (; n_block > 0; ) {
-            Tensor tSrS = partition_fragment_C(tiled_mma0, select<0, 1>(TileShape_MNK{}));
-            consumer_wait(pipeline_k, smem_pipe_read_k);
-            warp_scheduler_barrier_sync();
-            gemm</*zero_init=*/true, /*wg_wait=*/-1>(tiled_mma0, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), tSrS);
-            softmax.rescale_o(tOrO, scores_scale);
-            consumer_wait(pipeline_v, smem_pipe_read_v);
-            gemm</*zero_init=*/false, /*wg_wait=*/-1>(tiled_mma1, tOrP, tOrV(_, _, _, smem_pipe_read_v.index()), tOrO);
-            warp_scheduler_barrier_arrive();
-            warpgroup_wait<1>();
-            pipeline_k.consumer_release(smem_pipe_read_k);  // release K
-            cute::copy(softmax.template max</*Is_first=*/false>(tSrS, mainloop_params.softmax_scale_log2), scores_scale);
-            softmax.template online_softmax</*Is_first=*/false>(tSrS, mainloop_params.softmax_scale_log2);
-            warpgroup_wait<0>();
-            pipeline_v.consumer_release(smem_pipe_read_v);  // release V
-            ++smem_pipe_read_k;
-            ++smem_pipe_read_v;
-            cute::copy(make_tensor(convert_type<Element>(tSrS).data(), convert_layout_acc_Aregs<typename Ktraits::TiledMma1>(tSrS.layout())), tOrP);
-            if constexpr (UseMoba) {
-                n_block -= qk_gate_topk_idx[idx];
-                idx += 1;
-            } else {
-                n_block -= 1;
-            }
-        }
-
-        softmax.rescale_o(tOrO, scores_scale);
-        consumer_wait(pipeline_v, smem_pipe_read_v);
-        gemm</*zero_init=*/false, /*wg_wait=*/-1>(tiled_mma1, tOrP, tOrV(_, _, _, smem_pipe_read_v.index()), tOrO);
-        cute::copy(softmax.finalize(mainloop_params.softmax_scale_log2), scores_scale);
-        warpgroup_wait<0>();
-        pipeline_v.consumer_release(smem_pipe_read_v);
-        ++smem_pipe_read_v;
-
-        softmax.rescale_o(tOrO, scores_scale);
-    }
-
-    template <int NumMmaThreads, typename SharedStorage, typename FrgTensorO, typename TiledMma, typename T>
-    CUTLASS_DEVICE void
-    store(Params const& mainloop_params,
-        FrgTensorO const& tOrO,
-        SharedStorage& shared_storage,
-        TiledMma tiled_mma,
-        int thread_idx,
-        const int o_head_stride,
-        const int real_seq,
-        T * out_ptr) {
-
-        Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{});
-        auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma);
-        auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx);
-
-        Tensor tOrO_out = convert_type<Element>(tOrO);
-        Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out);
-        Tensor taccOsO = smem_thr_copy_O.partition_D(sO);
-
-        cutlass::arch::NamedBarrier::sync(NumMmaThreads, static_cast<int>(AttnNamedBarriers::ValueEmpty) /*id*/);
-        cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);
-        cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA
-        cutlass::arch::NamedBarrier::arrive(NumMmaThreads + cutlass::NumThreadsPerWarp,cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
-
-        Tensor gO = make_tensor(make_gmem_ptr(out_ptr),
-            Shape<Int<kBlockM>, Int<kHeadDim>>{},
-            make_stride(o_head_stride, _1{}));
-
-        GmemTiledCopyO gmem_tiled_copy_O;
-        auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx);
-
-        Tensor tOsO = gmem_thr_copy_O.partition_S(sO);
-        Tensor tOgO = gmem_thr_copy_O.partition_D(gO);
-
-        Tensor cO = make_identity_tensor(Shape<Int<kBlockM>, Int<kHeadDim>>{});
-
-        Tensor tOcO = gmem_thr_copy_O.partition_S(cO);
-
-        if (real_seq >= kBlockM) {
-            copy<true>(gmem_tiled_copy_O, tOsO, tOgO, tOcO);
-        } else {
-            copy<false>(gmem_tiled_copy_O, tOsO, tOgO, tOcO, real_seq);
-        }
-    }
-
+  }
 };
diff --git a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_encoder_attn.cu b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_encoder_attn.cu
index d4df8b27e2..c138cd1c36 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_encoder_attn.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_encoder_attn.cu
@@ -18,7 +18,7 @@
 #include "cutlass/util/print_error.hpp"
 #include "cutlass/util/GPU_Clock.hpp"
 #if defined(CUTLASS_ENABLE_CUBLAS) && CUTLASS_ENABLE_CUBLAS != 0
-#  include "cutlass/util/cublas_wrappers.hpp"
+#include "cutlass/util/cublas_wrappers.hpp"
 #endif
 #include "moba_attn/moba_attn_utils.hpp"
 #include "moba_attn/moba_attn.h"
@@ -33,356 +33,377 @@
 
 template <int kHeadDim>
 auto get_gmem_layout(int token_num, int head_num) {
-    return make_layout(
-        make_shape(token_num, kHeadDim, head_num),
-        make_stride(head_num * kHeadDim, _1{}, kHeadDim));
+  return make_layout(make_shape(token_num, kHeadDim, head_num),
+                     make_stride(head_num * kHeadDim, _1{}, kHeadDim));
 }
 
 template <typename Ktraits>
-__global__ void __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp, 1)
+__global__ void __launch_bounds__(Ktraits::kNWarps* cutlass::NumThreadsPerWarp,
+                                  1)
     moba_encoder_attention_kernel(
-        CUTE_GRID_CONSTANT typename CollectiveMainloopAttn<Ktraits>::Params const mainloop_params,
+        CUTE_GRID_CONSTANT
+        typename CollectiveMainloopAttn<Ktraits>::Params const mainloop_params,
         CUTE_GRID_CONSTANT moba_encoder_attn_params const data_params) {
+  using Element = typename Ktraits::Element;
+  using ElementAccum = typename Ktraits::ElementAccum;
+  using SoftType = ElementAccum;
+  using TileShape_MNK = typename Ktraits::TileShape_MNK;
+  using ClusterShape = typename Ktraits::ClusterShape_MNK;
 
-    using Element = typename Ktraits::Element;
-    using ElementAccum = typename Ktraits::ElementAccum;
-    using SoftType = ElementAccum;
-    using TileShape_MNK = typename Ktraits::TileShape_MNK;
-    using ClusterShape = typename Ktraits::ClusterShape_MNK;
+  static constexpr int NumMmaThreads = size(typename Ktraits::TiledMma0{});
+  static constexpr int NumCopyThreads = cutlass::NumThreadsPerWarpGroup;
+  static constexpr int kBlockM = Ktraits::kBlockM;
+  static constexpr int kBlockN = Ktraits::kBlockN;
+  constexpr int kHeadDim = Ktraits::kHeadDim;
+  constexpr int kMaxN = Ktraits::kMaxN;
 
-    static constexpr int NumMmaThreads = size(typename Ktraits::TiledMma0{});
-    static constexpr int NumCopyThreads = cutlass::NumThreadsPerWarpGroup;
-    static constexpr int kBlockM = Ktraits::kBlockM;
-    static constexpr int kBlockN = Ktraits::kBlockN;
-    constexpr int kHeadDim = Ktraits::kHeadDim;
-    constexpr int kMaxN = Ktraits::kMaxN;
+  using CollectiveMainloop = CollectiveMainloopAttn<Ktraits>;
 
-    using CollectiveMainloop = CollectiveMainloopAttn<Ktraits>;
+  using MainloopPipeline = typename Ktraits::MainloopPipeline;
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState = typename MainloopPipeline::PipelineState;
 
-    using MainloopPipeline = typename Ktraits::MainloopPipeline;
-    using PipelineParams = typename MainloopPipeline::Params;
-    using PipelineState = typename MainloopPipeline::PipelineState;
+  extern __shared__ char shared_memory[];
+  auto& shared_storage =
+      *reinterpret_cast<typename Ktraits::SharedStorage*>(shared_memory);
 
-    extern __shared__ char shared_memory[];
-    auto &shared_storage = *reinterpret_cast<typename Ktraits::SharedStorage*>(shared_memory);
+  const int m_block = blockIdx.x;
+  const int bidh = blockIdx.y;
+  const int bidb = blockIdx.z;
 
-    const int m_block = blockIdx.x;
-    const int bidh = blockIdx.y;
-    const int bidb = blockIdx.z;
+  const int seq_len_q = data_params.seq_len_encoder[bidb];
+  const int seq_len_k =
+      data_params.cu_seq_k[bidb + 1] - data_params.cu_seq_k[bidb];
 
-    const int seq_len_q = data_params.seq_len_encoder[bidb];
-    const int seq_len_k = data_params.cu_seq_k[bidb + 1] - data_params.cu_seq_k[bidb];
+  if (seq_len_q == 0) {
+    return;
+  }
 
+  __align__(16) __shared__ int qk_gate_topk_idx[kMaxN];
+  const int* qk_gate_idx_cur_offset =
+      data_params.qk_gate_topk_idx +
+      data_params.cu_seq_q_pack[bidb] / kBlockM * data_params.head_num * kMaxN +
+      (m_block * data_params.head_num + bidh) * kMaxN;
 
-    if (seq_len_q == 0) {
-        return;
+#pragma unroll
+  for (int i = threadIdx.x; i < kMaxN / 4;
+       i += Ktraits::kNWarps * cutlass::NumThreadsPerWarp) {
+    reinterpret_cast<int4*>(qk_gate_topk_idx)[i] =
+        reinterpret_cast<const int4*>(qk_gate_idx_cur_offset)[i];
+  }
+
+  const int n_block_max = min(
+      cute::ceil_div((m_block + 1) * kBlockM + seq_len_k - seq_len_q, kBlockN),
+      cute::ceil_div(seq_len_k, kBlockN));
+
+  if (m_block * kBlockM >= seq_len_q) {
+    return;
+  }
+
+  int const lane_predicate = cute::elect_one_sync();
+  int const warp_idx = cutlass::canonical_warp_idx_sync();
+
+  if (warp_idx == 0 && lane_predicate) {
+    CollectiveMainloop::prefetch_tma_descriptors(mainloop_params);
+  }
+
+  // Obtain warp index
+  int const warp_group_thread_idx =
+      threadIdx.x % cutlass::NumThreadsPerWarpGroup;
+
+  PipelineParams pipeline_params;
+  pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesK;
+  int warp_group_idx = cutlass::canonical_warp_group_idx();
+  pipeline_params.role = warp_group_idx == 0
+                             ? MainloopPipeline::ThreadCategory::Producer
+                             : MainloopPipeline::ThreadCategory::Consumer;
+  pipeline_params.is_leader = warp_group_thread_idx == 0;
+  pipeline_params.num_consumers = NumMmaThreads;
+
+  if (warp_idx == 0 && lane_predicate) {
+    shared_storage.barrier_Q.init(1);
+  }
+
+  MainloopPipeline pipeline_k(
+      shared_storage.pipeline_k, pipeline_params, ClusterShape{});
+  MainloopPipeline pipeline_v(
+      shared_storage.pipeline_v, pipeline_params, ClusterShape{});
+
+  __syncthreads();
+
+  CollectiveMainloop collective_mainloop;
+
+  if (warp_group_idx == 0) {
+    cutlass::arch::warpgroup_reg_dealloc<Ktraits::kNWarps == 8 ? 56 : 24>();
+
+    int warp_idx_in_warpgroup =
+        __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);
+    if (warp_idx_in_warpgroup == 0) {
+      PipelineState smem_pipe_write_k =
+          cutlass::make_producer_start_state<MainloopPipeline>();
+      PipelineState smem_pipe_write_v =
+          cutlass::make_producer_start_state<MainloopPipeline>();
+
+      collective_mainloop.load<Ktraits::UseMoba>(mainloop_params,
+                                                 pipeline_k,
+                                                 pipeline_v,
+                                                 smem_pipe_write_k,
+                                                 smem_pipe_write_v,
+                                                 shared_storage,
+                                                 qk_gate_topk_idx,
+                                                 n_block_max,
+                                                 m_block,
+                                                 bidh,
+                                                 bidb,
+                                                 data_params.cu_seq_q,
+                                                 data_params.cu_seq_k,
+                                                 seq_len_q,
+                                                 seq_len_k);
     }
+  } else {
+    cutlass::arch::warpgroup_reg_alloc<Ktraits::kNWarps == 8 ? 256 : 240>();
+    typename Ktraits::TiledMma1 tiled_mma1;
 
-    __align__(16) __shared__ int qk_gate_topk_idx[kMaxN];
-    const int *qk_gate_idx_cur_offset = data_params.qk_gate_topk_idx + data_params.cu_seq_q_pack[bidb] / kBlockM * data_params.head_num * kMaxN + (m_block * data_params.head_num + bidh) * kMaxN;
+    collective_mainloop.mma_init();
 
-    #pragma unroll
-    for (int i = threadIdx.x; i < kMaxN / 4; i += Ktraits::kNWarps * cutlass::NumThreadsPerWarp) {
-        reinterpret_cast<int4*>(qk_gate_topk_idx)[i] = reinterpret_cast<const int4*>(qk_gate_idx_cur_offset)[i];
-    }
+    PipelineState smem_pipe_read_k, smem_pipe_read_v;
 
+    Tensor tOrO =
+        partition_fragment_C(tiled_mma1, select<0, 2>(TileShape_MNK{}));
+    Softmax<2 * (2 * kBlockM / NumMmaThreads)> softmax;
 
-    const int n_block_max = min(cute::ceil_div((m_block + 1) * kBlockM + seq_len_k - seq_len_q, kBlockN), cute::ceil_div(seq_len_k, kBlockN));
+    collective_mainloop.mma<Ktraits::UseMoba>(mainloop_params,
+                                              pipeline_k,
+                                              pipeline_v,
+                                              smem_pipe_read_k,
+                                              smem_pipe_read_v,
+                                              tOrO,
+                                              softmax,
+                                              qk_gate_topk_idx,
+                                              n_block_max,
+                                              threadIdx.x - NumCopyThreads,
+                                              m_block,
+                                              seq_len_q,
+                                              seq_len_k,
+                                              shared_storage);
 
-    if (m_block * kBlockM >= seq_len_q) {
-        return;
-    }
+    const int o_head_stride = data_params.head_num * kHeadDim;
+    const int store_offset =
+        (data_params.cu_seq_q[bidb] + m_block * kBlockM) * o_head_stride +
+        bidh * kHeadDim;
 
-    int const lane_predicate = cute::elect_one_sync();
-    int const warp_idx = cutlass::canonical_warp_idx_sync();
-
-    if (warp_idx == 0 && lane_predicate) {
-        CollectiveMainloop::prefetch_tma_descriptors(mainloop_params);
-    }
-
-    // Obtain warp index
-    int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup;
-
-    PipelineParams pipeline_params;
-    pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesK;
-    int warp_group_idx = cutlass::canonical_warp_group_idx();
-    pipeline_params.role = warp_group_idx == 0
-        ? MainloopPipeline::ThreadCategory::Producer
-        : MainloopPipeline::ThreadCategory::Consumer;
-    pipeline_params.is_leader = warp_group_thread_idx == 0;
-    pipeline_params.num_consumers = NumMmaThreads;
-
-    if (warp_idx == 0 && lane_predicate) {
-        shared_storage.barrier_Q.init(1);
-    }
-
-    MainloopPipeline pipeline_k(shared_storage.pipeline_k, pipeline_params, ClusterShape{});
-    MainloopPipeline pipeline_v(shared_storage.pipeline_v, pipeline_params, ClusterShape{});
-
-    __syncthreads();
-
-    CollectiveMainloop collective_mainloop;
-
-    if (warp_group_idx == 0) {
-        cutlass::arch::warpgroup_reg_dealloc<Ktraits::kNWarps == 8 ? 56 : 24>();
-
-        int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);
-        if (warp_idx_in_warpgroup == 0) {
-            PipelineState smem_pipe_write_k = cutlass::make_producer_start_state<MainloopPipeline>();
-            PipelineState smem_pipe_write_v = cutlass::make_producer_start_state<MainloopPipeline>();
-
-            collective_mainloop.load<Ktraits::UseMoba>(
-                mainloop_params,
-                pipeline_k,
-                pipeline_v,
-                smem_pipe_write_k,
-                smem_pipe_write_v,
-                shared_storage,
-                qk_gate_topk_idx,
-                n_block_max,
-                m_block,
-                bidh,
-                bidb,
-                data_params.cu_seq_q,
-                data_params.cu_seq_k,
-                seq_len_q,
-                seq_len_k);
-        }
-    } else {
-        cutlass::arch::warpgroup_reg_alloc<Ktraits::kNWarps == 8 ? 256 : 240>();
-        typename Ktraits::TiledMma1 tiled_mma1;
-
-        collective_mainloop.mma_init();
-
-        PipelineState smem_pipe_read_k, smem_pipe_read_v;
-
-        Tensor tOrO = partition_fragment_C(tiled_mma1, select<0, 2>(TileShape_MNK{}));
-        Softmax<2 * (2 * kBlockM / NumMmaThreads)> softmax;
-
-        collective_mainloop.mma<Ktraits::UseMoba>(
-            mainloop_params,
-            pipeline_k,
-            pipeline_v,
-            smem_pipe_read_k,
-            smem_pipe_read_v,
-            tOrO,
-            softmax,
-            qk_gate_topk_idx,
-            n_block_max,
-            threadIdx.x - NumCopyThreads,
-            m_block,
-            seq_len_q,
-            seq_len_k,
-            shared_storage);
-
-        const int o_head_stride = data_params.head_num * kHeadDim;
-        const int store_offset = (data_params.cu_seq_q[bidb] + m_block * kBlockM) * o_head_stride + bidh * kHeadDim;
-
-        const int real_seq = seq_len_q - m_block * kBlockM;
-
-        collective_mainloop.store<NumMmaThreads>(
-            mainloop_params,
-            tOrO,
-            shared_storage,
-            tiled_mma1,
-            threadIdx.x - NumCopyThreads,
-            o_head_stride,
-            real_seq,
-            reinterpret_cast<Element*>(data_params.o_ptr) + store_offset);
-    }
+    const int real_seq = seq_len_q - m_block * kBlockM;
 
+    collective_mainloop.store<NumMmaThreads>(
+        mainloop_params,
+        tOrO,
+        shared_storage,
+        tiled_mma1,
+        threadIdx.x - NumCopyThreads,
+        o_head_stride,
+        real_seq,
+        reinterpret_cast<Element*>(data_params.o_ptr) + store_offset);
+  }
 }
 
+template <typename Kernel_traits>
+void run_moba_decoder_attn(moba_encoder_attn_params& params,
+                           cudaStream_t stream) {
+  using Element = typename Kernel_traits::Element;
+  using TileShape_MNK = typename Kernel_traits::TileShape_MNK;
+  using ClusterShape = typename Kernel_traits::ClusterShape_MNK;
 
-template<typename Kernel_traits>
-void run_moba_decoder_attn(moba_encoder_attn_params &params, cudaStream_t stream) {
-    using Element = typename Kernel_traits::Element;
-    using TileShape_MNK = typename Kernel_traits::TileShape_MNK;
-    using ClusterShape = typename Kernel_traits::ClusterShape_MNK;
+  using CollectiveMainloop = CollectiveMainloopAttn<Kernel_traits>;
+  constexpr int kHeadDim = Kernel_traits::kHeadDim;
 
-    using CollectiveMainloop = CollectiveMainloopAttn<Kernel_traits>;
-    constexpr int kHeadDim = Kernel_traits::kHeadDim;
+  typename CollectiveMainloop::Params mainloop_params =
+      CollectiveMainloop::to_underlying_arguments(
+          {static_cast<Element const*>(params.q_ptr),
+           get_gmem_layout<kHeadDim>(params.max_seq_q * params.batch_size,
+                                     params.head_num),
+           static_cast<Element const*>(params.k_ptr),
+           get_gmem_layout<kHeadDim>(params.max_seq_k * params.batch_size,
+                                     params.kv_head_num),
+           static_cast<Element const*>(params.v_ptr),
+           get_gmem_layout<kHeadDim>(params.max_seq_k * params.batch_size,
+                                     params.kv_head_num),
+           params.scale_softmax_log2});
 
-    typename CollectiveMainloop::Params mainloop_params =
-        CollectiveMainloop::to_underlying_arguments({
-            static_cast<Element const*>(params.q_ptr),
-            get_gmem_layout<kHeadDim>(params.max_seq_q * params.batch_size, params.head_num),
-            static_cast<Element const*>(params.k_ptr),
-            get_gmem_layout<kHeadDim>(params.max_seq_k * params.batch_size, params.kv_head_num),
-            static_cast<Element const*>(params.v_ptr),
-            get_gmem_layout<kHeadDim>(params.max_seq_k * params.batch_size, params.kv_head_num),
-            params.scale_softmax_log2
-        });
+  int num_blocks_m =
+      cutlass::ceil_div(params.max_seq_q, Kernel_traits::kBlockM);
 
-    int num_blocks_m = cutlass::ceil_div(params.max_seq_q, Kernel_traits::kBlockM);
+  num_blocks_m = cutlass::ceil_div(num_blocks_m, size<0>(ClusterShape{})) *
+                 size<0>(ClusterShape{});
 
-    num_blocks_m = cutlass::ceil_div(num_blocks_m, size<0>(ClusterShape{})) * size<0>(ClusterShape{});
+  void* kernel;
+  kernel = (void*)moba_encoder_attention_kernel<Kernel_traits>;
+  int smem_size = sizeof(typename Kernel_traits::SharedStorage);
 
-    void *kernel;
-    kernel = (void *)moba_encoder_attention_kernel<Kernel_traits>;
-    int smem_size = sizeof(typename Kernel_traits::SharedStorage);
+  if (smem_size >= 48 * 1024) {
+    cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  }
 
-    if (smem_size >= 48 * 1024) {
-       cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-    }
+  dim3 grid_dims;
+  grid_dims.x = num_blocks_m;
+  grid_dims.y = params.head_num;
+  grid_dims.z = params.batch_size;
 
-    dim3 grid_dims;
-    grid_dims.x = num_blocks_m;
-    grid_dims.y = params.head_num;
-    grid_dims.z = params.batch_size;
-
-    static constexpr int ctaSize = Kernel_traits::kNWarps * 32;
-    dim3 block_dims(ctaSize);
-    dim3 cluster_dims(size<0>(ClusterShape{}), size<1>(ClusterShape{}), size<2>(ClusterShape{}));
-    cutlass::ClusterLaunchParams launch_params{grid_dims, block_dims, cluster_dims, smem_size, stream};
-    cutlass::launch_kernel_on_cluster(launch_params, kernel, mainloop_params, params);
+  static constexpr int ctaSize = Kernel_traits::kNWarps * 32;
+  dim3 block_dims(ctaSize);
+  dim3 cluster_dims(size<0>(ClusterShape{}),
+                    size<1>(ClusterShape{}),
+                    size<2>(ClusterShape{}));
+  cutlass::ClusterLaunchParams launch_params{
+      grid_dims, block_dims, cluster_dims, smem_size, stream};
+  cutlass::launch_kernel_on_cluster(
+      launch_params, kernel, mainloop_params, params);
 }
 
-
 template <int kBlockM, int kBlockN, int kMaxN, typename InputType>
-void run_moba_encoder_attn_hdim128(moba_encoder_attn_params &params, cudaStream_t stream) {
+void run_moba_encoder_attn_hdim128(moba_encoder_attn_params& params,
+                                   cudaStream_t stream) {
+  constexpr static int Headdim = 128;
+  constexpr static int kNWarps = kBlockM / 16 + 4;
+  constexpr static int kStages = 2;
 
-    constexpr static int Headdim = 128;
-    constexpr static int kNWarps = kBlockM / 16 + 4;
-    constexpr static int kStages = 2;
-
-    using Ktraits = moba_encoder_attn_kernel_traits<Headdim, kBlockM, kBlockN, kNWarps, kStages, kMaxN, true, InputType>;
-    run_moba_decoder_attn<Ktraits>(params, stream);
+  using Ktraits = moba_encoder_attn_kernel_traits<Headdim,
+                                                  kBlockM,
+                                                  kBlockN,
+                                                  kNWarps,
+                                                  kStages,
+                                                  kMaxN,
+                                                  true,
+                                                  InputType>;
+  run_moba_decoder_attn<Ktraits>(params, stream);
 }
 
 template <typename T>
-void DispatchMobaEncoderAttn(
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& k_input,
-        const paddle::Tensor& v_input,
-        const paddle::Tensor& qk_gate_topk_idx,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::Tensor& cu_seq_q_pack,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& out,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int batch_size,
-        const int max_input_length) {
+void DispatchMobaEncoderAttn(const paddle::Tensor& q_input,
+                             const paddle::Tensor& k_input,
+                             const paddle::Tensor& v_input,
+                             const paddle::Tensor& qk_gate_topk_idx,
+                             const paddle::Tensor& cu_seq_q,
+                             const paddle::Tensor& cu_seq_k,
+                             const paddle::Tensor& cu_seq_q_pack,
+                             const paddle::Tensor& seq_len_encoder,
+                             const paddle::Tensor& seq_len_decoder,
+                             const paddle::Tensor& out,
+                             const int max_seq_q,
+                             const int max_seq_k,
+                             const int head_num,
+                             const int kv_head_num,
+                             const int head_dim,
+                             const int batch_size,
+                             const int max_input_length) {
+  constexpr int kBlockM = 128;
+  constexpr int kBlockN = 128;
+  constexpr int kMobaBlockSize = 128;
+  constexpr int kMaxN = 1024;
 
-    constexpr int kBlockM = 128;
-    constexpr int kBlockN = 128;
-    constexpr int kMobaBlockSize = 128;
-    constexpr int kMaxN = 1024;
+  using cute_type = typename cuteType<T>::type;
 
-    using cute_type = typename cuteType<T>::type;
+  moba_encoder_attn_params params;
+  memset(&params, 0, sizeof(moba_encoder_attn_params));
 
-    moba_encoder_attn_params params;
-    memset(&params, 0, sizeof(moba_encoder_attn_params));
+  params.q_ptr =
+      reinterpret_cast<cute_type*>(const_cast<T*>(q_input.data<T>()));
+  params.k_ptr =
+      reinterpret_cast<cute_type*>(const_cast<T*>(k_input.data<T>()));
+  params.v_ptr =
+      reinterpret_cast<cute_type*>(const_cast<T*>(v_input.data<T>()));
+  params.o_ptr = reinterpret_cast<cute_type*>(const_cast<T*>(out.data<T>()));
+  params.cu_seq_q = const_cast<int*>(cu_seq_q.data<int>());
+  params.cu_seq_k = const_cast<int*>(cu_seq_k.data<int>());
+  params.head_num = head_num;
+  params.kv_head_num = kv_head_num;
+  params.max_seq_q = max_seq_q;
+  params.max_seq_k = max_seq_k;
+  params.batch_size = batch_size;
+  params.gqa_group_size = head_num / kv_head_num;
+  constexpr float kLog2e = 1.4426950408889634074;
+  params.scale_softmax_log2 = 1.0f / std::sqrt(head_dim) * kLog2e;
+  params.qk_gate_topk_idx = const_cast<int*>(qk_gate_topk_idx.data<int>());
+  params.seq_len_encoder = const_cast<int*>(seq_len_encoder.data<int>());
+  params.cu_seq_q_pack = const_cast<int*>(cu_seq_q_pack.data<int>());
 
-    params.q_ptr = reinterpret_cast<cute_type*>(const_cast<T*>(q_input.data<T>()));
-    params.k_ptr = reinterpret_cast<cute_type*>(const_cast<T*>(k_input.data<T>()));
-    params.v_ptr = reinterpret_cast<cute_type*>(const_cast<T*>(v_input.data<T>()));
-    params.o_ptr = reinterpret_cast<cute_type*>(const_cast<T*>(out.data<T>()));
-    params.cu_seq_q = const_cast<int*>(cu_seq_q.data<int>());
-    params.cu_seq_k = const_cast<int*>(cu_seq_k.data<int>());
-    params.head_num = head_num;
-    params.kv_head_num = kv_head_num;
-    params.max_seq_q = max_seq_q;
-    params.max_seq_k = max_seq_k;
-    params.batch_size = batch_size;
-    params.gqa_group_size = head_num / kv_head_num;
-    constexpr float kLog2e = 1.4426950408889634074;
-    params.scale_softmax_log2 = 1.0f / std::sqrt(head_dim) * kLog2e;
-    params.qk_gate_topk_idx = const_cast<int*>(qk_gate_topk_idx.data<int>());
-    params.seq_len_encoder = const_cast<int*>(seq_len_encoder.data<int>());
-    params.cu_seq_q_pack = const_cast<int*>(cu_seq_q_pack.data<int>());
-
-    run_moba_encoder_attn_hdim128<kBlockM, kBlockN, kMaxN, cute_type>(params, out.stream());
+  run_moba_encoder_attn_hdim128<kBlockM, kBlockN, kMaxN, cute_type>(
+      params, out.stream());
 }
 
-void MobaEncoderAttn(
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& k_input,
-        const paddle::Tensor& v_input,
-        const paddle::Tensor& qk_gate_topk_idx,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::Tensor& cu_seq_q_pack,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& out,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_input_length) {
-
-    const int batch_size = seq_len_encoder.dims()[0];
-    if (q_input.dtype() == paddle::DataType::FLOAT16) {
-        return
-            DispatchMobaEncoderAttn<phi::dtype::float16>(
-                q_input,
-                k_input,
-                v_input,
-                qk_gate_topk_idx,
-                cu_seq_q,
-                cu_seq_k,
-                cu_seq_q_pack,
-                seq_len_encoder,
-                seq_len_decoder,
-                out,
-                max_seq_q,
-                max_seq_k,
-                head_num,
-                kv_head_num,
-                head_dim,
-                batch_size,
-                max_input_length);
-    } else if (q_input.dtype() == paddle::DataType::BFLOAT16) {
-        return
-            DispatchMobaEncoderAttn<phi::dtype::bfloat16>(
-                q_input,
-                k_input,
-                v_input,
-                qk_gate_topk_idx,
-                cu_seq_q,
-                cu_seq_k,
-                cu_seq_q_pack,
-                seq_len_encoder,
-                seq_len_decoder,
-                out,
-                max_seq_q,
-                max_seq_k,
-                head_num,
-                kv_head_num,
-                head_dim,
-                batch_size,
-                max_input_length);
-    }
+void MobaEncoderAttn(const paddle::Tensor& q_input,
+                     const paddle::Tensor& k_input,
+                     const paddle::Tensor& v_input,
+                     const paddle::Tensor& qk_gate_topk_idx,
+                     const paddle::Tensor& cu_seq_q,
+                     const paddle::Tensor& cu_seq_k,
+                     const paddle::Tensor& cu_seq_q_pack,
+                     const paddle::Tensor& seq_len_encoder,
+                     const paddle::Tensor& seq_len_decoder,
+                     const paddle::Tensor& out,
+                     const int max_seq_q,
+                     const int max_seq_k,
+                     const int head_num,
+                     const int kv_head_num,
+                     const int head_dim,
+                     const int max_input_length) {
+  const int batch_size = seq_len_encoder.dims()[0];
+  if (q_input.dtype() == paddle::DataType::FLOAT16) {
+    return DispatchMobaEncoderAttn<phi::dtype::float16>(q_input,
+                                                        k_input,
+                                                        v_input,
+                                                        qk_gate_topk_idx,
+                                                        cu_seq_q,
+                                                        cu_seq_k,
+                                                        cu_seq_q_pack,
+                                                        seq_len_encoder,
+                                                        seq_len_decoder,
+                                                        out,
+                                                        max_seq_q,
+                                                        max_seq_k,
+                                                        head_num,
+                                                        kv_head_num,
+                                                        head_dim,
+                                                        batch_size,
+                                                        max_input_length);
+  } else if (q_input.dtype() == paddle::DataType::BFLOAT16) {
+    return DispatchMobaEncoderAttn<phi::dtype::bfloat16>(q_input,
+                                                         k_input,
+                                                         v_input,
+                                                         qk_gate_topk_idx,
+                                                         cu_seq_q,
+                                                         cu_seq_k,
+                                                         cu_seq_q_pack,
+                                                         seq_len_encoder,
+                                                         seq_len_decoder,
+                                                         out,
+                                                         max_seq_q,
+                                                         max_seq_k,
+                                                         head_num,
+                                                         kv_head_num,
+                                                         head_dim,
+                                                         batch_size,
+                                                         max_input_length);
+  }
 }
 
-
 PD_BUILD_STATIC_OP(moba_encoder_attn)
-    .Inputs({
-        "q_input",
-        "k_input",
-        "v_input",
-        "qk_gate_topk_idx",
-        "cu_seq_q",
-        "cu_seq_k",
-        "cu_seq_q_pack",
-        "seq_len_encoder",
-        "seq_len_decoder",
-        "out"})
-    .Attrs({
-        "max_seq_q: int",
-        "max_seq_k: int",
-        "head_num: int",
-        "kv_head_num: int",
-        "head_dim: int",
-        "max_input_length: int"})
+    .Inputs({"q_input",
+             "k_input",
+             "v_input",
+             "qk_gate_topk_idx",
+             "cu_seq_q",
+             "cu_seq_k",
+             "cu_seq_q_pack",
+             "seq_len_encoder",
+             "seq_len_decoder",
+             "out"})
+    .Attrs({"max_seq_q: int",
+            "max_seq_k: int",
+            "head_num: int",
+            "kv_head_num: int",
+            "head_dim: int",
+            "max_input_length: int"})
     .Outputs({"attn_out"})
     .SetInplaceMap({{"out", "attn_out"}})
     .SetKernelFn(PD_KERNEL(MobaEncoderAttn));
diff --git a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_encoder_write_cache.cu b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_encoder_write_cache.cu
index 89666c51f3..4bacfd1e03 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_encoder_write_cache.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_encoder_write_cache.cu
@@ -20,147 +20,156 @@
 #endif
 
 template <typename T, int kBlockSize, int kHeadDim>
-__global__ void write_encoder_cachekv_c16(
-        const T * k_input,
-        const T * v_input,
-        const int * cu_seq_k,
-        const int * seq_len_encoder,
-        const int * seq_len_decoder,
-        T * cache_k,
-        T * cache_v,
-        const int * block_tables,
-        const int kv_head_num,
-        const int max_blocks_per_seq) {
+__global__ void write_encoder_cachekv_c16(const T* k_input,
+                                          const T* v_input,
+                                          const int* cu_seq_k,
+                                          const int* seq_len_encoder,
+                                          const int* seq_len_decoder,
+                                          T* cache_k,
+                                          T* cache_v,
+                                          const int* block_tables,
+                                          const int kv_head_num,
+                                          const int max_blocks_per_seq) {
+  constexpr int kPackSize = 16 / sizeof(T);
+  const int block_idx = blockIdx.x * kBlockSize;
+  int bidh = blockIdx.y;
+  const int bidb = blockIdx.z;
+  const int tidx = threadIdx.x;
+  const int row_idx = tidx / (kHeadDim / kPackSize);
+  const int col_idx = tidx % (kHeadDim / kPackSize) * kPackSize;
+  const int seq_len = seq_len_encoder[bidb];
 
-    constexpr int kPackSize = 16 / sizeof(T);
-    const int block_idx = blockIdx.x * kBlockSize;
-    int bidh = blockIdx.y;
-    const int bidb = blockIdx.z;
-    const int tidx = threadIdx.x;
-    const int row_idx = tidx / (kHeadDim / kPackSize);
-    const int col_idx = tidx % (kHeadDim / kPackSize) * kPackSize;
-    const int seq_len = seq_len_encoder[bidb];
+  if (seq_len == 0) return;
 
-    if (seq_len == 0) return;
+  const int remain_tokens = seq_len - block_idx;
 
-    const int remain_tokens = seq_len - block_idx;
+  const int32_t* block_table_now = block_tables + bidb * max_blocks_per_seq;
+  const uint32_t physical_block_number =
+      block_table_now[blockIdx.x + seq_len_decoder[bidb] / kBlockSize];
 
-    const int32_t *block_table_now = block_tables + bidb * max_blocks_per_seq;
-    const uint32_t physical_block_number = block_table_now[blockIdx.x + seq_len_decoder[bidb] / kBlockSize];
-
-    if (bidh < kv_head_num) {
-        T * cache = cache_k + physical_block_number * kv_head_num * kBlockSize * kHeadDim + bidh * kBlockSize * kHeadDim + col_idx;
-        const int base_load_idx = (block_idx + cu_seq_k[bidb]) * kv_head_num * kHeadDim + bidh * kHeadDim + col_idx;
-
-        #pragma unroll
-        for (int i = row_idx; i < kBlockSize; i += 128 / (kHeadDim / kPackSize)) {
-            if (i < remain_tokens) {
-                *reinterpret_cast<float4*>(cache + i * kHeadDim) = *reinterpret_cast<const float4*>(k_input + base_load_idx + i * kv_head_num * kHeadDim);
-            }
-        }
-    } else {
-        bidh -= kv_head_num;
-        const int base_load_idx = (block_idx + cu_seq_k[bidb]) * kv_head_num * kHeadDim + bidh * kHeadDim + col_idx;
-        T * cache = cache_v + physical_block_number * kv_head_num * kBlockSize * kHeadDim + bidh * kBlockSize * kHeadDim + col_idx;
-
-        #pragma unroll
-        for (int i = row_idx; i < kBlockSize; i += 128 / (kHeadDim / kPackSize)) {
-            if (i < remain_tokens) {
-                *reinterpret_cast<float4*>(cache + i * kHeadDim) = *reinterpret_cast<const float4*>(v_input + base_load_idx + i * kv_head_num * kHeadDim);
-            }
-        }
+  if (bidh < kv_head_num) {
+    T* cache = cache_k +
+               physical_block_number * kv_head_num * kBlockSize * kHeadDim +
+               bidh * kBlockSize * kHeadDim + col_idx;
+    const int base_load_idx =
+        (block_idx + cu_seq_k[bidb]) * kv_head_num * kHeadDim +
+        bidh * kHeadDim + col_idx;
 
+#pragma unroll
+    for (int i = row_idx; i < kBlockSize; i += 128 / (kHeadDim / kPackSize)) {
+      if (i < remain_tokens) {
+        *reinterpret_cast<float4*>(cache + i * kHeadDim) =
+            *reinterpret_cast<const float4*>(k_input + base_load_idx +
+                                             i * kv_head_num * kHeadDim);
+      }
     }
+  } else {
+    bidh -= kv_head_num;
+    const int base_load_idx =
+        (block_idx + cu_seq_k[bidb]) * kv_head_num * kHeadDim +
+        bidh * kHeadDim + col_idx;
+    T* cache = cache_v +
+               physical_block_number * kv_head_num * kBlockSize * kHeadDim +
+               bidh * kBlockSize * kHeadDim + col_idx;
+
+#pragma unroll
+    for (int i = row_idx; i < kBlockSize; i += 128 / (kHeadDim / kPackSize)) {
+      if (i < remain_tokens) {
+        *reinterpret_cast<float4*>(cache + i * kHeadDim) =
+            *reinterpret_cast<const float4*>(v_input + base_load_idx +
+                                             i * kv_head_num * kHeadDim);
+      }
+    }
+  }
 }
 
 void MobaEncoderAttnWriteCacheKv(
-        const paddle::Tensor& k_input,
-        const paddle::Tensor& v_input,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cache_k,
-        const paddle::Tensor& cache_v,
-        const paddle::Tensor& block_tables,
-        const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_zero_points,
-        const paddle::optional<paddle::Tensor>& cache_v_zero_points,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_seq_q,
-        const std::string &cache_quant_type_str) {
-
-    constexpr int kThreads = 128;
-    constexpr int kHeadDim = 128;
-    assert(kHeadDim == head_dim);
-    constexpr int kBlockSize = 64;
-    const int batch_size = block_tables.dims()[0];
-    const int max_blocks_per_seq = block_tables.dims()[1];
-    if (cache_quant_type_str == "none") {
-        dim3 grid_dims;
-        grid_dims.x = (max_seq_q + kBlockSize - 1) / kBlockSize;
-        grid_dims.y = kv_head_num * 2;
-        grid_dims.z = batch_size;
-        if (k_input.dtype() == paddle::DataType::FLOAT16) {
-            using T = phi::dtype::float16;
-            write_encoder_cachekv_c16<T, kBlockSize, kHeadDim><<<grid_dims, kThreads, 0, k_input.stream()>>>(
-                const_cast<T*>(k_input.data<T>()),
-                const_cast<T*>(v_input.data<T>()),
-                cu_seq_k.data<int>(),
-                seq_len_encoder.data<int>(),
-                seq_len_decoder.data<int>(),
-                const_cast<T*>(cache_k.data<T>()),
-                const_cast<T*>(cache_v.data<T>()),
-                block_tables.data<int>(),
-                kv_head_num,
-                max_blocks_per_seq);
-        } else if (k_input.dtype() == paddle::DataType::BFLOAT16) {
-            using T = phi::dtype::bfloat16;
-            write_encoder_cachekv_c16<T, kBlockSize, kHeadDim><<<grid_dims, kThreads, 0, k_input.stream()>>>(
-                const_cast<T*>(k_input.data<T>()),
-                const_cast<T*>(v_input.data<T>()),
-                cu_seq_k.data<int>(),
-                seq_len_encoder.data<int>(),
-                seq_len_decoder.data<int>(),
-                const_cast<T*>(cache_k.data<T>()),
-                const_cast<T*>(cache_v.data<T>()),
-                block_tables.data<int>(),
-                kv_head_num,
-                max_blocks_per_seq);
-        }
-    } else {
-        PADDLE_THROW(phi::errors::Unimplemented(
-            "Quantized cache not implemented for cache_quant_type = %s", cache_quant_type_str.c_str()));
+    const paddle::Tensor& k_input,
+    const paddle::Tensor& v_input,
+    const paddle::Tensor& cu_seq_k,
+    const paddle::Tensor& seq_len_encoder,
+    const paddle::Tensor& seq_len_decoder,
+    const paddle::Tensor& cache_k,
+    const paddle::Tensor& cache_v,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& cache_k_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_quant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zero_points,
+    const paddle::optional<paddle::Tensor>& cache_v_zero_points,
+    const int head_num,
+    const int kv_head_num,
+    const int head_dim,
+    const int max_seq_q,
+    const std::string& cache_quant_type_str) {
+  constexpr int kThreads = 128;
+  constexpr int kHeadDim = 128;
+  assert(kHeadDim == head_dim);
+  constexpr int kBlockSize = 64;
+  const int batch_size = block_tables.dims()[0];
+  const int max_blocks_per_seq = block_tables.dims()[1];
+  if (cache_quant_type_str == "none") {
+    dim3 grid_dims;
+    grid_dims.x = (max_seq_q + kBlockSize - 1) / kBlockSize;
+    grid_dims.y = kv_head_num * 2;
+    grid_dims.z = batch_size;
+    if (k_input.dtype() == paddle::DataType::FLOAT16) {
+      using T = phi::dtype::float16;
+      write_encoder_cachekv_c16<T, kBlockSize, kHeadDim>
+          <<<grid_dims, kThreads, 0, k_input.stream()>>>(
+              const_cast<T*>(k_input.data<T>()),
+              const_cast<T*>(v_input.data<T>()),
+              cu_seq_k.data<int>(),
+              seq_len_encoder.data<int>(),
+              seq_len_decoder.data<int>(),
+              const_cast<T*>(cache_k.data<T>()),
+              const_cast<T*>(cache_v.data<T>()),
+              block_tables.data<int>(),
+              kv_head_num,
+              max_blocks_per_seq);
+    } else if (k_input.dtype() == paddle::DataType::BFLOAT16) {
+      using T = phi::dtype::bfloat16;
+      write_encoder_cachekv_c16<T, kBlockSize, kHeadDim>
+          <<<grid_dims, kThreads, 0, k_input.stream()>>>(
+              const_cast<T*>(k_input.data<T>()),
+              const_cast<T*>(v_input.data<T>()),
+              cu_seq_k.data<int>(),
+              seq_len_encoder.data<int>(),
+              seq_len_decoder.data<int>(),
+              const_cast<T*>(cache_k.data<T>()),
+              const_cast<T*>(cache_v.data<T>()),
+              block_tables.data<int>(),
+              kv_head_num,
+              max_blocks_per_seq);
     }
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "Quantized cache not implemented for cache_quant_type = %s",
+        cache_quant_type_str.c_str()));
+  }
 }
 
 PD_BUILD_STATIC_OP(moba_encoder_attn_write_cache_kv)
-    .Inputs({
-        "k_input",
-        "v_input",
-        "cu_seq_k",
-        "seq_len_encoder",
-        "seq_len_decoder",
-        "cache_k",
-        "cache_v",
-        "block_tables",
-        paddle::Optional("cache_k_quant_scale"),
-        paddle::Optional("cache_v_quant_scale"),
-        paddle::Optional("cache_k_dequant_scale"),
-        paddle::Optional("cache_v_dequant_scale"),
-        paddle::Optional("cache_k_zero_points"),
-        paddle::Optional("cache_v_zero_points")})
-    .Attrs({
-        "head_num: int",
-        "kv_head_num: int",
-        "head_dim: int",
-        "max_seq_q: int",
-        "cache_quant_type_str: std::string"})
+    .Inputs({"k_input",
+             "v_input",
+             "cu_seq_k",
+             "seq_len_encoder",
+             "seq_len_decoder",
+             "cache_k",
+             "cache_v",
+             "block_tables",
+             paddle::Optional("cache_k_quant_scale"),
+             paddle::Optional("cache_v_quant_scale"),
+             paddle::Optional("cache_k_dequant_scale"),
+             paddle::Optional("cache_v_dequant_scale"),
+             paddle::Optional("cache_k_zero_points"),
+             paddle::Optional("cache_v_zero_points")})
+    .Attrs({"head_num: int",
+            "kv_head_num: int",
+            "head_dim: int",
+            "max_seq_q: int",
+            "cache_quant_type_str: std::string"})
     .Outputs({"cache_k_out", "cache_v_out"})
-    .SetInplaceMap({{"cache_k", "cache_k_out"},
-                    {"cache_v", "cache_v_out"}})
+    .SetInplaceMap({{"cache_k", "cache_k_out"}, {"cache_v", "cache_v_out"}})
     .SetKernelFn(PD_KERNEL(MobaEncoderAttnWriteCacheKv));
diff --git a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_qk_sort_encoder.cu b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_qk_sort_encoder.cu
index dd044551d4..c0e6e42dbb 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_qk_sort_encoder.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/moba_qk_sort_encoder.cu
@@ -19,327 +19,329 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-template <typename T, int knthreads, int moba_block_size, int kBlockM, int kBlockMaxN, int searchtimes>
-__global__ void qk_gate_sort_encoder_kernel(
-        const T* qk_gate_weight,
-        int * qk_gate_topk_idx,
-        const int *seq_len_encoder,
-        const int *seq_len_decoder,
-        const int* cu_seq_q,
-        const int* cu_seq_k,
-        const int* cu_seq_q_pack,
-        const int use_moba_seq_limit,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int kGqaGroupSize,
-        const int top_k_left,
-        const int top_k_right) {
+template <typename T,
+          int knthreads,
+          int moba_block_size,
+          int kBlockM,
+          int kBlockMaxN,
+          int searchtimes>
+__global__ void qk_gate_sort_encoder_kernel(const T* qk_gate_weight,
+                                            int* qk_gate_topk_idx,
+                                            const int* seq_len_encoder,
+                                            const int* seq_len_decoder,
+                                            const int* cu_seq_q,
+                                            const int* cu_seq_k,
+                                            const int* cu_seq_q_pack,
+                                            const int use_moba_seq_limit,
+                                            const int max_seq_q,
+                                            const int max_seq_k,
+                                            const int head_num,
+                                            const int kv_head_num,
+                                            const int kGqaGroupSize,
+                                            const int top_k_left,
+                                            const int top_k_right) {
+  const int bidt = blockIdx.x * kBlockM;
+  const int bidh = blockIdx.y;
+  const int bidb = blockIdx.z;
+  const int tidx = threadIdx.x;
 
-    const int bidt = blockIdx.x * kBlockM;
-    const int bidh = blockIdx.y;
-    const int bidb = blockIdx.z;
-    const int tidx = threadIdx.x;
+  constexpr int kPackSize = kBlockMaxN / knthreads;
 
-    constexpr int kPackSize = kBlockMaxN / knthreads;
+  static_assert(kBlockMaxN % knthreads == 0);
 
-    static_assert(kBlockMaxN % knthreads == 0);
+  const int seq_len_q = seq_len_encoder[bidb];
 
-    const int seq_len_q = seq_len_encoder[bidb];
+  if (seq_len_q == 0 || bidt >= seq_len_q) {
+    return;
+  }
 
-    if (seq_len_q == 0 || bidt >= seq_len_q) {
-        return;
+  const int seq_len_k = (bidt + kBlockM + seq_len_decoder[bidb]);
+
+  const int seq_len_moba = seq_len_k / moba_block_size;
+
+  using SrcType = Vec<T, kPackSize>;
+  using SrcType_f = Vec<float, kPackSize>;
+  using SrcType_i = Vec<int, kPackSize>;
+
+  SrcType src;
+  SrcType_f src_f;
+
+  SrcType_i select_idx;
+
+  select_idx.set_zero();
+
+  const int store_idx = cu_seq_q_pack[bidb] / kBlockM * head_num * kBlockMaxN +
+                        bidh * kBlockMaxN + blockIdx.x * head_num * kBlockMaxN +
+                        tidx * kPackSize;
+
+  if (seq_len_k < use_moba_seq_limit) {
+#pragma unroll
+    for (int i = 0; i < kPackSize; i++) {
+      select_idx.data.elt[i] = 1;
+    }
+    select_idx.store_to(qk_gate_topk_idx + store_idx);
+    return;
+  }
+
+  const int load_offset = (cu_seq_q[bidb] + bidt) * head_num * kBlockMaxN +
+                          bidh * kBlockMaxN + tidx * kPackSize;
+  const int data_len = seq_len_moba - tidx * kPackSize;
+
+#pragma unroll
+  for (int t = 0; t < kBlockM; t++) {
+    if (bidt + t >= seq_len_q) {
+      break;
+    }
+    src.load_from(qk_gate_weight + load_offset + t * head_num * kBlockMaxN);
+    float min_global = FLT_MAX;
+    float max_global = -FLT_MAX;
+#pragma unroll
+    for (int i = 0; i < kPackSize; i++) {
+      if (i < data_len) {
+        src_f.data.elt[i] = float(src.data.elt[i]);
+        min_global = min(min_global, src_f.data.elt[i]);
+      } else {
+        src_f.data.elt[i] = -FLT_MAX;
+      }
+      max_global = max(max_global, src_f.data.elt[i]);
     }
 
-    const int seq_len_k = (bidt + kBlockM + seq_len_decoder[bidb]);
+    max_global = BlockAllReduce<float, MaxOp<float>, knthreads>(max_global);
+    min_global = BlockAllReduce<float, MinOp<float>, knthreads>(min_global);
 
-    const int seq_len_moba = seq_len_k / moba_block_size;
+    float right_limit = max_global;
+    float left_limit = min_global;
 
-    using SrcType = Vec<T, kPackSize>;
-    using SrcType_f = Vec<float, kPackSize>;
-    using SrcType_i = Vec<int, kPackSize>;
+    float mid_limit;
+    int count;
 
-    SrcType src;
-    SrcType_f src_f;
-
-    SrcType_i select_idx;
-
-    select_idx.set_zero();
-
-    const int store_idx = cu_seq_q_pack[bidb] / kBlockM * head_num * kBlockMaxN + bidh * kBlockMaxN + blockIdx.x * head_num * kBlockMaxN + tidx * kPackSize;
-
-    if (seq_len_k < use_moba_seq_limit) {
-        #pragma unroll
-        for (int i = 0; i < kPackSize; i++) {
-            select_idx.data.elt[i] = 1;
-        }
-        select_idx.store_to(qk_gate_topk_idx + store_idx);
-        return;
-    }
-
-    const int load_offset = (cu_seq_q[bidb] + bidt) * head_num * kBlockMaxN + bidh * kBlockMaxN + tidx * kPackSize;
-    const int data_len = seq_len_moba - tidx * kPackSize;
-
-    #pragma unroll
-    for (int t = 0; t < kBlockM; t++) {
-        if (bidt + t >= seq_len_q) {
-            break;
-        }
-        src.load_from(qk_gate_weight + load_offset + t * head_num * kBlockMaxN);
-        float min_global = FLT_MAX;
-        float max_global = -FLT_MAX;
-        #pragma unroll
-        for (int i = 0; i < kPackSize; i++) {
-            if (i < data_len) {
-                src_f.data.elt[i] = float(src.data.elt[i]);
-                min_global = min(min_global, src_f.data.elt[i]);
-            } else {
-                src_f.data.elt[i] = -FLT_MAX;
-            }
-            max_global = max(max_global, src_f.data.elt[i]);
-        }
-
-        max_global = BlockAllReduce<float, MaxOp<float>, knthreads>(max_global);
-        min_global = BlockAllReduce<float, MinOp<float>, knthreads>(min_global);
-
-        float right_limit = max_global;
-        float left_limit = min_global;
-
-        float mid_limit;
-        int count;
-
-        if (right_limit == left_limit) {
-            mid_limit = (left_limit + right_limit) * 0.5f;
+    if (right_limit == left_limit) {
+      mid_limit = (left_limit + right_limit) * 0.5f;
+    } else {
+#pragma unroll
+      for (int i = 0; i < searchtimes; i++) {
+        mid_limit = (left_limit + right_limit) * 0.5f;
+        count = get_data_count<kPackSize, knthreads>(src_f.data.elt, mid_limit);
+        if (count < top_k_left) {
+          right_limit = mid_limit;
+        } else if (count > top_k_right) {
+          left_limit = mid_limit;
         } else {
-            #pragma unroll
-            for (int i = 0; i < searchtimes; i++) {
-                mid_limit = (left_limit + right_limit) * 0.5f;
-                count = get_data_count<kPackSize, knthreads>(src_f.data.elt, mid_limit);
-                if (count < top_k_left) {
-                    right_limit = mid_limit;
-                } else if (count > top_k_right) {
-                    left_limit = mid_limit;
-                } else {
-                    break;
-                }
-            }
-        }
-
-        #pragma unroll
-        for (int i = 0; i < kPackSize; i++) {
-            if (src_f.data.elt[i] >= mid_limit) {
-                select_idx.data.elt[i] = 1;
-            }
+          break;
         }
+      }
     }
 
-    if (tidx == 0) {
-        select_idx.data.elt[0] = 1;
+#pragma unroll
+    for (int i = 0; i < kPackSize; i++) {
+      if (src_f.data.elt[i] >= mid_limit) {
+        select_idx.data.elt[i] = 1;
+      }
     }
+  }
 
-    __align__(16) __shared__ int qk_gate_mem[kBlockMaxN];
-    __align__(16) __shared__ int qk_continue_idx_mem[kBlockMaxN];
-    select_idx.store_to(qk_gate_mem + tidx * kPackSize);
+  if (tidx == 0) {
+    select_idx.data.elt[0] = 1;
+  }
 
-    __syncthreads();
+  __align__(16) __shared__ int qk_gate_mem[kBlockMaxN];
+  __align__(16) __shared__ int qk_continue_idx_mem[kBlockMaxN];
+  select_idx.store_to(qk_gate_mem + tidx * kPackSize);
 
-    if (tidx == 0) {
-        int cur_idx = 0;
+  __syncthreads();
+
+  if (tidx == 0) {
+    int cur_idx = 0;
+    int idx = -1;
+    const int last_idx = seq_len_moba - 1;
+    while (last_idx + idx >= 0 && qk_gate_mem[last_idx + idx] == 0) {
+      idx--;
+    }
+    qk_continue_idx_mem[cur_idx] = -idx;
+    cur_idx++;
+
+    for (int i = last_idx - 1; i >= 0; --i) {
+      if (qk_gate_mem[i] == 1) {
         int idx = -1;
-        const int last_idx = seq_len_moba - 1;
-        while (last_idx + idx >= 0 && qk_gate_mem[last_idx + idx] == 0) {
-            idx--;
+        while (i + idx >= 0 && qk_gate_mem[i + idx] == 0) {
+          idx--;
         }
         qk_continue_idx_mem[cur_idx] = -idx;
         cur_idx++;
-
-        for (int i = last_idx - 1; i >= 0; --i) {
-            if (qk_gate_mem[i] == 1) {
-                int idx = -1;
-                while (i + idx >= 0 && qk_gate_mem[i + idx] == 0) {
-                    idx--;
-                }
-                qk_continue_idx_mem[cur_idx] = -idx;
-                cur_idx++;
-            }
-        }
-        qk_continue_idx_mem[cur_idx] = 10000000;
+      }
     }
+    qk_continue_idx_mem[cur_idx] = 10000000;
+  }
 
-    __syncthreads();
+  __syncthreads();
 
-    *reinterpret_cast<SrcType_i *>(qk_gate_topk_idx + store_idx) = reinterpret_cast<SrcType_i *>(qk_continue_idx_mem)[tidx];
+  *reinterpret_cast<SrcType_i*>(qk_gate_topk_idx + store_idx) =
+      reinterpret_cast<SrcType_i*>(qk_continue_idx_mem)[tidx];
 }
 
 template <int kBlockM, int kMaxN, int moba_block_size, typename T>
-void qk_gate_sort_encoder(
-        const T* qk_gate_weight,
-        int * qk_gate_topk_idx,
-        const int *seq_len_encoder,
-        const int *seq_len_decoder,
-        const int* cu_seq_q,
-        const int* cu_seq_k,
-        const int* cu_seq_q_pack,
-        const int use_moba_seq_limit,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int batch_size,
-        const int top_k_left,
-        const int top_k_right,
-        cudaStream_t stream) {
+void qk_gate_sort_encoder(const T* qk_gate_weight,
+                          int* qk_gate_topk_idx,
+                          const int* seq_len_encoder,
+                          const int* seq_len_decoder,
+                          const int* cu_seq_q,
+                          const int* cu_seq_k,
+                          const int* cu_seq_q_pack,
+                          const int use_moba_seq_limit,
+                          const int max_seq_q,
+                          const int max_seq_k,
+                          const int head_num,
+                          const int kv_head_num,
+                          const int batch_size,
+                          const int top_k_left,
+                          const int top_k_right,
+                          cudaStream_t stream) {
+  constexpr int kPackSize = 16 / sizeof(T);
 
-    constexpr int kPackSize = 16 / sizeof(T);
+  const int gqa_group_size = head_num / kv_head_num;
+  const int knthreads = kMaxN / kPackSize;
+  const int searchtimes = 6;
 
-    const int gqa_group_size = head_num / kv_head_num;
-    const int knthreads = kMaxN / kPackSize;
-    const int searchtimes = 6;
+  dim3 grid_dims;
+  grid_dims.x = (max_seq_q + kBlockM - 1) / kBlockM;
+  grid_dims.y = head_num;
+  grid_dims.z = batch_size;
 
-    dim3 grid_dims;
-    grid_dims.x = (max_seq_q + kBlockM - 1) / kBlockM;
-    grid_dims.y = head_num;
-    grid_dims.z = batch_size;
+  constexpr auto kernel = qk_gate_sort_encoder_kernel<T,
+                                                      knthreads,
+                                                      moba_block_size,
+                                                      kBlockM,
+                                                      kMaxN,
+                                                      searchtimes>;
 
-    constexpr auto kernel = qk_gate_sort_encoder_kernel<T, knthreads, moba_block_size, kBlockM, kMaxN, searchtimes>;
-
-    kernel<<<grid_dims, knthreads, 0, stream>>>(
-        qk_gate_weight,
-        qk_gate_topk_idx,
-        seq_len_encoder,
-        seq_len_decoder,
-        cu_seq_q,
-        cu_seq_k,
-        cu_seq_q_pack,
-        use_moba_seq_limit,
-        max_seq_q,
-        max_seq_k,
-        head_num,
-        kv_head_num,
-        gqa_group_size,
-        top_k_left,
-        top_k_right);
+  kernel<<<grid_dims, knthreads, 0, stream>>>(qk_gate_weight,
+                                              qk_gate_topk_idx,
+                                              seq_len_encoder,
+                                              seq_len_decoder,
+                                              cu_seq_q,
+                                              cu_seq_k,
+                                              cu_seq_q_pack,
+                                              use_moba_seq_limit,
+                                              max_seq_q,
+                                              max_seq_k,
+                                              head_num,
+                                              kv_head_num,
+                                              gqa_group_size,
+                                              top_k_left,
+                                              top_k_right);
 }
 template <typename T>
 std::vector<paddle::Tensor> DispatchQkSortEncoder(
-        const paddle::Tensor& qk_gate_weight,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::Tensor& cu_seq_q_pack,
-        const paddle::Tensor& q_pack_tokens,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int top_k_left,
-        const int top_k_right,
-        const int use_moba_seq_limit) {
-    constexpr int kBlockM = 128;
-    constexpr int kBlockN = 128;
-    constexpr int kMobaBlockSize = 128;
-    constexpr int kMaxN = 1024;
-    using cute_type = typename cuteType<T>::type;
-    const int batch_size = seq_len_encoder.dims()[0];
+    const paddle::Tensor& qk_gate_weight,
+    const paddle::Tensor& seq_len_encoder,
+    const paddle::Tensor& seq_len_decoder,
+    const paddle::Tensor& cu_seq_q,
+    const paddle::Tensor& cu_seq_k,
+    const paddle::Tensor& cu_seq_q_pack,
+    const paddle::Tensor& q_pack_tokens,
+    const int max_seq_q,
+    const int max_seq_k,
+    const int head_num,
+    const int kv_head_num,
+    const int top_k_left,
+    const int top_k_right,
+    const int use_moba_seq_limit) {
+  constexpr int kBlockM = 128;
+  constexpr int kBlockN = 128;
+  constexpr int kMobaBlockSize = 128;
+  constexpr int kMaxN = 1024;
+  using cute_type = typename cuteType<T>::type;
+  const int batch_size = seq_len_encoder.dims()[0];
 
-    paddle::Tensor qk_gate_topk_idx = paddle::empty({q_pack_tokens.data<int>()[0] / kBlockM, head_num, kMaxN}, paddle::DataType::INT32, qk_gate_weight.place());
+  paddle::Tensor qk_gate_topk_idx =
+      paddle::empty({q_pack_tokens.data<int>()[0] / kBlockM, head_num, kMaxN},
+                    paddle::DataType::INT32,
+                    qk_gate_weight.place());
 
-    qk_gate_sort_encoder<kBlockM, kMaxN, kMobaBlockSize, cute_type>(
-            reinterpret_cast<const cute_type *>(qk_gate_weight.data<T>()),
-            qk_gate_topk_idx.data<int>(),
-            seq_len_encoder.data<int>(),
-            seq_len_decoder.data<int>(),
-            cu_seq_q.data<int>(),
-            cu_seq_k.data<int>(),
-            cu_seq_q_pack.data<int>(),
-            use_moba_seq_limit,
-            max_seq_q,
-            max_seq_k,
-            head_num,
-            kv_head_num,
-            batch_size,
-            top_k_left,
-            top_k_right,
-            qk_gate_weight.stream());
+  qk_gate_sort_encoder<kBlockM, kMaxN, kMobaBlockSize, cute_type>(
+      reinterpret_cast<const cute_type*>(qk_gate_weight.data<T>()),
+      qk_gate_topk_idx.data<int>(),
+      seq_len_encoder.data<int>(),
+      seq_len_decoder.data<int>(),
+      cu_seq_q.data<int>(),
+      cu_seq_k.data<int>(),
+      cu_seq_q_pack.data<int>(),
+      use_moba_seq_limit,
+      max_seq_q,
+      max_seq_k,
+      head_num,
+      kv_head_num,
+      batch_size,
+      top_k_left,
+      top_k_right,
+      qk_gate_weight.stream());
 
-    return {qk_gate_topk_idx};
+  return {qk_gate_topk_idx};
 }
 
-
-std::vector<paddle::Tensor> QkSortEncoder(
-        const paddle::Tensor& qk_gate_weight,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::Tensor& cu_seq_q_pack,
-        const paddle::Tensor& q_pack_tokens,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int top_k_left,
-        const int top_k_right,
-        const int use_moba_seq_limit) {
-    if (qk_gate_weight.dtype() == paddle::DataType::FLOAT16) {
-        return std::move(
-            DispatchQkSortEncoder<phi::dtype::float16>(
-                qk_gate_weight,
-                seq_len_encoder,
-                seq_len_decoder,
-                cu_seq_q,
-                cu_seq_k,
-                cu_seq_q_pack,
-                q_pack_tokens,
-                max_seq_q,
-                max_seq_k,
-                head_num,
-                kv_head_num,
-                top_k_left,
-                top_k_right,
-                use_moba_seq_limit
-            )
-        );
-    } else if (qk_gate_weight.dtype() == paddle::DataType::BFLOAT16) {
-        return std::move(
-            DispatchQkSortEncoder<phi::dtype::bfloat16>(
-                qk_gate_weight,
-                seq_len_encoder,
-                seq_len_decoder,
-                cu_seq_q,
-                cu_seq_k,
-                cu_seq_q_pack,
-                q_pack_tokens,
-                max_seq_q,
-                max_seq_k,
-                head_num,
-                kv_head_num,
-                top_k_left,
-                top_k_right,
-                use_moba_seq_limit
-            )
-        );
-    }
+std::vector<paddle::Tensor> QkSortEncoder(const paddle::Tensor& qk_gate_weight,
+                                          const paddle::Tensor& seq_len_encoder,
+                                          const paddle::Tensor& seq_len_decoder,
+                                          const paddle::Tensor& cu_seq_q,
+                                          const paddle::Tensor& cu_seq_k,
+                                          const paddle::Tensor& cu_seq_q_pack,
+                                          const paddle::Tensor& q_pack_tokens,
+                                          const int max_seq_q,
+                                          const int max_seq_k,
+                                          const int head_num,
+                                          const int kv_head_num,
+                                          const int top_k_left,
+                                          const int top_k_right,
+                                          const int use_moba_seq_limit) {
+  if (qk_gate_weight.dtype() == paddle::DataType::FLOAT16) {
+    return std::move(
+        DispatchQkSortEncoder<phi::dtype::float16>(qk_gate_weight,
+                                                   seq_len_encoder,
+                                                   seq_len_decoder,
+                                                   cu_seq_q,
+                                                   cu_seq_k,
+                                                   cu_seq_q_pack,
+                                                   q_pack_tokens,
+                                                   max_seq_q,
+                                                   max_seq_k,
+                                                   head_num,
+                                                   kv_head_num,
+                                                   top_k_left,
+                                                   top_k_right,
+                                                   use_moba_seq_limit));
+  } else if (qk_gate_weight.dtype() == paddle::DataType::BFLOAT16) {
+    return std::move(
+        DispatchQkSortEncoder<phi::dtype::bfloat16>(qk_gate_weight,
+                                                    seq_len_encoder,
+                                                    seq_len_decoder,
+                                                    cu_seq_q,
+                                                    cu_seq_k,
+                                                    cu_seq_q_pack,
+                                                    q_pack_tokens,
+                                                    max_seq_q,
+                                                    max_seq_k,
+                                                    head_num,
+                                                    kv_head_num,
+                                                    top_k_left,
+                                                    top_k_right,
+                                                    use_moba_seq_limit));
+  }
 }
 
 PD_BUILD_STATIC_OP(moba_qk_sort_encoder)
-    .Inputs({
-        "qk_gate_weight",
-        "seq_len_encoder",
-        "seq_len_decoder",
-        "cu_seq_q",
-        "cu_seq_k",
-        "cu_seq_q_pack",
-        "q_pack_tokens"})
-    .Attrs({
-        "max_seq_q: int",
-        "max_seq_k: int",
-        "head_num: int",
-        "kv_head_num: int",
-        "top_k_left: int",
-        "top_k_right: int",
-        "use_moba_seq_limit: int"})
+    .Inputs({"qk_gate_weight",
+             "seq_len_encoder",
+             "seq_len_decoder",
+             "cu_seq_q",
+             "cu_seq_k",
+             "cu_seq_q_pack",
+             "q_pack_tokens"})
+    .Attrs({"max_seq_q: int",
+            "max_seq_k: int",
+            "head_num: int",
+            "kv_head_num: int",
+            "top_k_left: int",
+            "top_k_right: int",
+            "use_moba_seq_limit: int"})
     .Outputs({"qk_gate_topk_idx"})
     .SetKernelFn(PD_KERNEL(QkSortEncoder));
diff --git a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/softmax.hpp b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/softmax.hpp
index fc7832d153..bb808c2a11 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/softmax.hpp
+++ b/custom_ops/gpu_ops/moba_attn/moba_encoder_attn/softmax.hpp
@@ -24,171 +24,223 @@
 
 using namespace cute;
 
-template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
-__device__ __forceinline__ void thread_reduce_(Tensor<Engine0, Layout0> const &tensor, Tensor<Engine1, Layout1> &summary, Operator &op) {
-    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
-    static_assert(Layout1::rank == 1, "Only support 1D Tensor");
-    CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
-    #pragma unroll
-    for (int mi = 0; mi < size<0>(tensor); mi++) {
-        summary(mi) = zero_init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0));
-        #pragma unroll
-        for (int ni = 1; ni < size<1>(tensor); ni++) {
-            summary(mi) = op(summary(mi), tensor(mi, ni));
-        }
+template <bool zero_init = true,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1,
+          typename Operator>
+__device__ __forceinline__ void thread_reduce_(
+    Tensor<Engine0, Layout0> const &tensor,
+    Tensor<Engine1, Layout1> &summary,
+    Operator &op) {
+  static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+  static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+  CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
+#pragma unroll
+  for (int mi = 0; mi < size<0>(tensor); mi++) {
+    summary(mi) = zero_init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0));
+#pragma unroll
+    for (int ni = 1; ni < size<1>(tensor); ni++) {
+      summary(mi) = op(summary(mi), tensor(mi, ni));
     }
+  }
 }
 
-template<typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
-__device__ __forceinline__ void quad_allreduce_(Tensor<Engine0, Layout0> &dst, Tensor<Engine1, Layout1> &src, Operator &op) {
-    CUTE_STATIC_ASSERT_V(size(dst) == size(src));
-    #pragma unroll
-    for (int i = 0; i < size(dst); i++){
-        dst(i) = Allreduce<4>::run(src(i), op);
-    }
+template <typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1,
+          typename Operator>
+__device__ __forceinline__ void quad_allreduce_(Tensor<Engine0, Layout0> &dst,
+                                                Tensor<Engine1, Layout1> &src,
+                                                Operator &op) {
+  CUTE_STATIC_ASSERT_V(size(dst) == size(src));
+#pragma unroll
+  for (int i = 0; i < size(dst); i++) {
+    dst(i) = Allreduce<4>::run(src(i), op);
+  }
 }
 
-template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
-__device__ __forceinline__ void reduce_(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &summary, Operator &op) {
-    thread_reduce_<zero_init>(tensor, summary, op);
-    quad_allreduce_(summary, summary, op);
+template <bool zero_init = true,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1,
+          typename Operator>
+__device__ __forceinline__ void reduce_(Tensor<Engine0, Layout0> const &tensor,
+                                        Tensor<Engine1, Layout1> &summary,
+                                        Operator &op) {
+  thread_reduce_<zero_init>(tensor, summary, op);
+  quad_allreduce_(summary, summary, op);
 }
 
-template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__device__ __forceinline__ void reduce_max(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &max){
-    MaxOp<float> max_op;
-    reduce_<zero_init>(tensor, max, max_op);
+template <bool zero_init = true,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1>
+__device__ __forceinline__ void reduce_max(
+    Tensor<Engine0, Layout0> const &tensor, Tensor<Engine1, Layout1> &max) {
+  MaxOp<float> max_op;
+  reduce_<zero_init>(tensor, max, max_op);
 }
 
-template<bool zero_init=true, bool warp_reduce=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__device__ __forceinline__ void reduce_sum(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &sum){
-    SumOp<float> sum_op;
-    thread_reduce_<zero_init>(tensor, sum, sum_op);
-    if constexpr (warp_reduce) { quad_allreduce_(sum, sum, sum_op); }
+template <bool zero_init = true,
+          bool warp_reduce = true,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1>
+__device__ __forceinline__ void reduce_sum(
+    Tensor<Engine0, Layout0> const &tensor, Tensor<Engine1, Layout1> &sum) {
+  SumOp<float> sum_op;
+  thread_reduce_<zero_init>(tensor, sum, sum_op);
+  if constexpr (warp_reduce) {
+    quad_allreduce_(sum, sum, sum_op);
+  }
 }
 
 __forceinline__ __device__ __half2 half_exp(__half2 x) {
-    uint32_t tmp_out, tmp_in;
-    tmp_in = reinterpret_cast<uint32_t&>(x);
-    asm ("ex2.approx.f16x2 %0, %1;\n"
-      : "=r"(tmp_out)
-      : "r"(tmp_in));
-    __half2 out = reinterpret_cast<__half2&>(tmp_out);
-    return out;
+  uint32_t tmp_out, tmp_in;
+  tmp_in = reinterpret_cast<uint32_t &>(x);
+  asm("ex2.approx.f16x2 %0, %1;\n" : "=r"(tmp_out) : "r"(tmp_in));
+  __half2 out = reinterpret_cast<__half2 &>(tmp_out);
+  return out;
 }
 
 // Apply the exp to all the elements.
-template <bool zero_init=false, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__forceinline__ __device__ void max_scale_exp2_sum(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> &max, Tensor<Engine1, Layout1> &sum, const float scale) {
-    static_assert(Layout0::rank == 2, "Only support 2D Tensor"); static_assert(Layout1::rank == 1, "Only support 1D Tensor"); CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
-    #pragma unroll
-    for (int mi = 0; mi < size<0>(tensor); ++mi) {
-        MaxOp<float> max_op;
-        max(mi) = zero_init ? tensor(mi, 0) : max_op(max(mi), tensor(mi, 0));
-        #pragma unroll
-        for (int ni = 1; ni < size<1>(tensor); ni++) {
-            max(mi) = max_op(max(mi), tensor(mi, ni));
-        }
-        max(mi) = Allreduce<4>::run(max(mi), max_op);
-        const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * scale;
-        sum(mi) = 0;
-        #pragma unroll
-        for (int ni = 0; ni < size<1>(tensor); ++ni)  {
-            tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
-            sum(mi) += tensor(mi, ni);
-        }
+template <bool zero_init = false,
+          typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1>
+__forceinline__ __device__ void max_scale_exp2_sum(
+    Tensor<Engine0, Layout0> &tensor,
+    Tensor<Engine1, Layout1> &max,
+    Tensor<Engine1, Layout1> &sum,
+    const float scale) {
+  static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+  static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+  CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
+#pragma unroll
+  for (int mi = 0; mi < size<0>(tensor); ++mi) {
+    MaxOp<float> max_op;
+    max(mi) = zero_init ? tensor(mi, 0) : max_op(max(mi), tensor(mi, 0));
+#pragma unroll
+    for (int ni = 1; ni < size<1>(tensor); ni++) {
+      max(mi) = max_op(max(mi), tensor(mi, ni));
     }
+    max(mi) = Allreduce<4>::run(max(mi), max_op);
+    const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * scale;
+    sum(mi) = 0;
+#pragma unroll
+    for (int ni = 0; ni < size<1>(tensor); ++ni) {
+      tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
+      sum(mi) += tensor(mi, ni);
+    }
+  }
 }
 
-
-template <typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__forceinline__ __device__ void scale_apply_exp2(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> const &max, const float scale) {
-    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
-    static_assert(Layout1::rank == 1, "Only support 1D Tensor");
-    CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
-    #pragma unroll
-    for (int mi = 0; mi < size<0>(tensor); ++mi) {
-        const float max_scaled = max(mi) * scale;
-        #pragma unroll
-        for (int ni = 0; ni < size<1>(tensor); ++ni)  {
-            tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
-        }
+template <typename Engine0,
+          typename Layout0,
+          typename Engine1,
+          typename Layout1>
+__forceinline__ __device__ void scale_apply_exp2(
+    Tensor<Engine0, Layout0> &tensor,
+    Tensor<Engine1, Layout1> const &max,
+    const float scale) {
+  static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+  static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+  CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
+#pragma unroll
+  for (int mi = 0; mi < size<0>(tensor); ++mi) {
+    const float max_scaled = max(mi) * scale;
+#pragma unroll
+    for (int ni = 0; ni < size<1>(tensor); ++ni) {
+      tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
     }
+  }
 }
 
-
 template <int kNRows>
 struct Softmax {
+  using TensorT = decltype(make_tensor<float>(Shape<Int<kNRows>>{}));
+  TensorT row_max, row_sum;
 
-    using TensorT = decltype(make_tensor<float>(Shape<Int<kNRows>>{}));
-    TensorT row_max, row_sum;
+  CUTLASS_DEVICE Softmax(){};
 
-    CUTLASS_DEVICE Softmax() {};
+  template <bool Is_first, bool Check_inf = false, typename Tensor0>
+  __forceinline__ __device__ TensorT max(Tensor0 &acc_s,
+                                         float softmax_scale_log2) {
+    Tensor scores =
+        make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+    static_assert(decltype(size<0>(scores))::value == kNRows);
+    TensorT scores_scale;
+    if constexpr (Is_first) {
+      reduce_max</*zero_init=*/true>(scores, row_max);
+      cute::fill(scores_scale, 1.f);
+    } else {
+      Tensor scores_max_prev = make_fragment_like(row_max);
+      cute::copy(row_max, scores_max_prev);
+      reduce_max</*zero_init=*/false>(scores, row_max);
+#pragma unroll
+      for (int mi = 0; mi < size(row_max); ++mi) {
+        float scores_max_cur = row_max(mi);
+        scores_scale(mi) =
+            exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2);
+        row_sum(mi) *= scores_scale(mi);
+      }
+    }
+    return scores_scale;
+  };
 
-    template<bool Is_first, bool Check_inf=false, typename Tensor0>
-    __forceinline__ __device__ TensorT max(Tensor0 &acc_s, float softmax_scale_log2) {
-        Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
-        static_assert(decltype(size<0>(scores))::value == kNRows);
-        TensorT scores_scale;
-        if constexpr (Is_first) {
-            reduce_max</*zero_init=*/true>(scores, row_max);
-            cute::fill(scores_scale, 1.f);
-        } else {
-            Tensor scores_max_prev = make_fragment_like(row_max);
-            cute::copy(row_max, scores_max_prev);
-            reduce_max</*zero_init=*/false>(scores, row_max);
-            #pragma unroll
-            for (int mi = 0; mi < size(row_max); ++mi) {
-                float scores_max_cur = row_max(mi);
-                scores_scale(mi) = exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2);
-                row_sum(mi) *= scores_scale(mi);
-            }
-        }
-        return scores_scale;
-    };
+  template <bool Is_first, typename Tensor0>
+  __forceinline__ __device__ TensorT online_softmax(Tensor0 &acc_s,
+                                                    float softmax_scale_log2) {
+    Tensor scores =
+        make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+    static_assert(decltype(size<0>(scores))::value == kNRows);
+    TensorT scores_scale;
+    if constexpr (Is_first) {
+      reduce_max</*zero_init=*/true>(scores, row_max);
+      scale_apply_exp2(scores, row_max, softmax_scale_log2);
+      reduce_sum</*zero_init=*/true, /*warp_reduce=*/false>(scores, row_sum);
+      cute::fill(scores_scale, 1.f);
+    } else {
+      scale_apply_exp2(scores, row_max, softmax_scale_log2);
+      reduce_sum</*zero_init=*/false, /*warp_reduce=*/false>(scores, row_sum);
+    }
+    return scores_scale;
+  };
 
-    template<bool Is_first, typename Tensor0>
-    __forceinline__ __device__ TensorT online_softmax(Tensor0 &acc_s, float softmax_scale_log2) {
-        Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
-        static_assert(decltype(size<0>(scores))::value == kNRows);
-        TensorT scores_scale;
-        if constexpr (Is_first) {
-            reduce_max</*zero_init=*/true>(scores, row_max);
-            scale_apply_exp2(scores, row_max, softmax_scale_log2);
-            reduce_sum</*zero_init=*/true, /*warp_reduce=*/false>(scores, row_sum);
-            cute::fill(scores_scale, 1.f);
-        } else {
-            scale_apply_exp2(scores, row_max, softmax_scale_log2);
-            reduce_sum</*zero_init=*/false, /*warp_reduce=*/false>(scores, row_sum);
-        }
-        return scores_scale;
-    };
-
-    __forceinline__ __device__ TensorT finalize(float softmax_scale_log2) {
-        SumOp<float> sum_op;
-        quad_allreduce_(row_sum, row_sum, sum_op);
-        TensorT scores_scale;
-        #pragma unroll
-        for (int mi = 0; mi < size(row_max); ++mi) {
-            float sum = row_sum(mi);
-            float inv_sum = 1.0f / sum;
-            row_sum(mi) = row_max(mi) * (softmax_scale_log2 * float(M_LN2)) + __logf(sum);
-            scores_scale(mi) = inv_sum;
-        }
-        return scores_scale;
-    };
-
-    template<typename Tensor1>
-    __forceinline__ __device__ void rescale_o(Tensor1 &acc_o, TensorT const &scores_scale) {
-        Tensor acc_o_rowcol = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
-        static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows);
-        #pragma unroll
-        for (int mi = 0; mi < size(row_max); ++mi) {
-            #pragma unroll
-            for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
-                acc_o_rowcol(mi, ni) *= scores_scale(mi);
-            }
-        }
-    };
+  __forceinline__ __device__ TensorT finalize(float softmax_scale_log2) {
+    SumOp<float> sum_op;
+    quad_allreduce_(row_sum, row_sum, sum_op);
+    TensorT scores_scale;
+#pragma unroll
+    for (int mi = 0; mi < size(row_max); ++mi) {
+      float sum = row_sum(mi);
+      float inv_sum = 1.0f / sum;
+      row_sum(mi) =
+          row_max(mi) * (softmax_scale_log2 * float(M_LN2)) + __logf(sum);
+      scores_scale(mi) = inv_sum;
+    }
+    return scores_scale;
+  };
 
+  template <typename Tensor1>
+  __forceinline__ __device__ void rescale_o(Tensor1 &acc_o,
+                                            TensorT const &scores_scale) {
+    Tensor acc_o_rowcol =
+        make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+    static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows);
+#pragma unroll
+    for (int mi = 0; mi < size(row_max); ++mi) {
+#pragma unroll
+      for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
+        acc_o_rowcol(mi, ni) *= scores_scale(mi);
+      }
+    }
+  };
 };
diff --git a/custom_ops/gpu_ops/moba_attn/moba_process/moba_get_kv_from_cache.cu b/custom_ops/gpu_ops/moba_attn/moba_process/moba_get_kv_from_cache.cu
index 66228342c7..f1f9b256f2 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_process/moba_get_kv_from_cache.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_process/moba_get_kv_from_cache.cu
@@ -21,96 +21,107 @@
 #endif
 
 template <typename T, int kBlockSize, int kHeadDim>
-__global__ void get_kv_from_cache_c16_kernel(
-        T * k_input,
-        T * v_input,
-        const int * seq_len_encoder,
-        const int * seq_len_decoder,
-        const int * cu_seq_k,
-        const T * cache_k,
-        const T * cache_v,
-        const int * block_tables,
-        const int kv_head_num,
-        const int head_dim,
-        const int batch_size,
-        const int max_input_length,
-        const int max_blocks_per_seq) {
+__global__ void get_kv_from_cache_c16_kernel(T* k_input,
+                                             T* v_input,
+                                             const int* seq_len_encoder,
+                                             const int* seq_len_decoder,
+                                             const int* cu_seq_k,
+                                             const T* cache_k,
+                                             const T* cache_v,
+                                             const int* block_tables,
+                                             const int kv_head_num,
+                                             const int head_dim,
+                                             const int batch_size,
+                                             const int max_input_length,
+                                             const int max_blocks_per_seq) {
+  const int block_idx = blockIdx.x;
+  int bidh = blockIdx.y;
+  const int bidb = blockIdx.z;
+  const int seq_len = seq_len_decoder[bidb] + seq_len_encoder[bidb];
+  const int tidx = threadIdx.x;
+  const int base_token_idx = block_idx * kBlockSize;
 
-    const int block_idx = blockIdx.x;
-    int bidh = blockIdx.y;
-    const int bidb = blockIdx.z;
-    const int seq_len = seq_len_decoder[bidb] + seq_len_encoder[bidb];
-    const int tidx = threadIdx.x;
-    const int base_token_idx = block_idx * kBlockSize;
+  if (base_token_idx >= seq_len || seq_len_encoder[bidb] == 0) {
+    return;
+  }
 
-    if (base_token_idx >= seq_len || seq_len_encoder[bidb] == 0) {
-        return;
+  constexpr int kPackSize = 16 / sizeof(T);
+
+  const int row_idx = tidx / (kHeadDim / kPackSize);
+  const int col_idx = tidx % (kHeadDim / kPackSize) * kPackSize;
+  const int physical_block_number =
+      block_tables[bidb * max_blocks_per_seq + block_idx];
+
+  const int remain_tokens = seq_len - base_token_idx;
+
+  if (bidh < kv_head_num) {
+    const int cache_offset =
+        physical_block_number * kv_head_num * kBlockSize * kHeadDim +
+        bidh * kBlockSize * kHeadDim + col_idx;
+    const int base_store_idx =
+        (base_token_idx + cu_seq_k[bidb]) * kv_head_num * kHeadDim +
+        bidh * kHeadDim + col_idx;
+#pragma unroll
+    for (int i = row_idx; i < kBlockSize; i += 128 / (kHeadDim / kPackSize)) {
+      if (i < remain_tokens) {
+        *reinterpret_cast<float4*>(k_input + base_store_idx +
+                                   i * kv_head_num * kHeadDim) =
+            *reinterpret_cast<const float4*>(cache_k + cache_offset +
+                                             i * kHeadDim);
+      }
     }
-
-    constexpr int kPackSize = 16 / sizeof(T);
-
-    const int row_idx = tidx / (kHeadDim / kPackSize);
-    const int col_idx = tidx % (kHeadDim / kPackSize) * kPackSize;
-    const int physical_block_number = block_tables[bidb * max_blocks_per_seq + block_idx];
-
-
-    const int remain_tokens = seq_len - base_token_idx;
-
-    if (bidh < kv_head_num) {
-        const int cache_offset = physical_block_number * kv_head_num * kBlockSize * kHeadDim + bidh * kBlockSize * kHeadDim + col_idx;
-        const int base_store_idx = (base_token_idx + cu_seq_k[bidb]) * kv_head_num * kHeadDim + bidh * kHeadDim + col_idx;
-        #pragma unroll
-        for (int i = row_idx; i < kBlockSize; i += 128 / (kHeadDim / kPackSize)) {
-            if (i < remain_tokens) {
-                *reinterpret_cast<float4*>(k_input + base_store_idx + i * kv_head_num * kHeadDim) = *reinterpret_cast<const float4*>(cache_k + cache_offset + i * kHeadDim);
-            }
-        }
-    } else {
-        bidh -= kv_head_num;
-        const int cache_offset = physical_block_number * kv_head_num * kBlockSize * kHeadDim + bidh * kBlockSize * kHeadDim + col_idx;
-        const int base_store_idx = (base_token_idx + cu_seq_k[bidb]) * kv_head_num * kHeadDim + bidh * kHeadDim + col_idx;
-        #pragma unroll
-        for (int i = row_idx; i < kBlockSize; i += 128 / (kHeadDim / kPackSize)) {
-            if (i < remain_tokens) {
-                *reinterpret_cast<float4*>(v_input + base_store_idx + i * kv_head_num * kHeadDim) = *reinterpret_cast<const float4*>(cache_v + cache_offset + i * kHeadDim);
-            }
-        }
+  } else {
+    bidh -= kv_head_num;
+    const int cache_offset =
+        physical_block_number * kv_head_num * kBlockSize * kHeadDim +
+        bidh * kBlockSize * kHeadDim + col_idx;
+    const int base_store_idx =
+        (base_token_idx + cu_seq_k[bidb]) * kv_head_num * kHeadDim +
+        bidh * kHeadDim + col_idx;
+#pragma unroll
+    for (int i = row_idx; i < kBlockSize; i += 128 / (kHeadDim / kPackSize)) {
+      if (i < remain_tokens) {
+        *reinterpret_cast<float4*>(v_input + base_store_idx +
+                                   i * kv_head_num * kHeadDim) =
+            *reinterpret_cast<const float4*>(cache_v + cache_offset +
+                                             i * kHeadDim);
+      }
     }
+  }
 }
 
 template <typename T>
-void get_kv_from_cache(
-        T * k_input,
-        T * v_input,
-        const int * seq_len_encoder,
-        const int * seq_len_decoder,
-        const int * cu_seq_k,
-        const void * cache_k,
-        const void * cache_v,
-        const int * block_tables,
-        const T * cache_k_dequant_scale,
-        const T * cache_v_dequant_scale,
-        const T * cache_k_zero_points,
-        const T * cache_v_zero_points,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_seq_k,
-        const int batch_size,
-        const int max_input_length,
-        const int max_blocks_per_seq,
-        const std::string &cache_quant_type_str,
-        cudaStream_t stream) {
-
-    constexpr int kThreads = 128;
-    constexpr int kHeadDim = 128;
-    assert(kHeadDim == head_dim);
-    constexpr int kBlockSize = 64;
-    if (cache_quant_type_str == "none") {
-        dim3 grid_dims;
-        grid_dims.x = (max_seq_k + kBlockSize - 1) / kBlockSize;
-        grid_dims.y = kv_head_num * 2;
-        grid_dims.z = batch_size;
-        get_kv_from_cache_c16_kernel<T, kBlockSize, kHeadDim><<<grid_dims, kThreads, 0, stream>>>(
+void get_kv_from_cache(T* k_input,
+                       T* v_input,
+                       const int* seq_len_encoder,
+                       const int* seq_len_decoder,
+                       const int* cu_seq_k,
+                       const void* cache_k,
+                       const void* cache_v,
+                       const int* block_tables,
+                       const T* cache_k_dequant_scale,
+                       const T* cache_v_dequant_scale,
+                       const T* cache_k_zero_points,
+                       const T* cache_v_zero_points,
+                       const int kv_head_num,
+                       const int head_dim,
+                       const int max_seq_k,
+                       const int batch_size,
+                       const int max_input_length,
+                       const int max_blocks_per_seq,
+                       const std::string& cache_quant_type_str,
+                       cudaStream_t stream) {
+  constexpr int kThreads = 128;
+  constexpr int kHeadDim = 128;
+  assert(kHeadDim == head_dim);
+  constexpr int kBlockSize = 64;
+  if (cache_quant_type_str == "none") {
+    dim3 grid_dims;
+    grid_dims.x = (max_seq_k + kBlockSize - 1) / kBlockSize;
+    grid_dims.y = kv_head_num * 2;
+    grid_dims.z = batch_size;
+    get_kv_from_cache_c16_kernel<T, kBlockSize, kHeadDim>
+        <<<grid_dims, kThreads, 0, stream>>>(
             k_input,
             v_input,
             seq_len_encoder,
@@ -124,169 +135,179 @@ void get_kv_from_cache(
             batch_size,
             max_input_length,
             max_blocks_per_seq);
-    } else {
-        PD_THROW("Only supported cache_quant_type_str in ['none'].");
-    }
+  } else {
+    PD_THROW("Only supported cache_quant_type_str in ['none'].");
+  }
 }
 
 void GetKVFromCache(
-        const paddle::Tensor& k_input,
-        const paddle::Tensor& v_input,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cache_k,
-        const paddle::Tensor& cache_v,
-        const paddle::Tensor& block_tables,
-        const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
-        const paddle::optional<paddle::Tensor>& cache_k_zero_points,
-        const paddle::optional<paddle::Tensor>& cache_v_zero_points,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_input_length,
-        const int max_seq_k,
-        const std::string &cache_quant_type_str) {
-
-    if (k_input.dtype() == paddle::DataType::FLOAT16) {
-        using T = phi::dtype::float16;
-        using cute_type = typename cuteType<T>::type;
-        get_kv_from_cache<cute_type>(
-            reinterpret_cast<cute_type*>(const_cast<T*>(k_input.data<T>())),
-            reinterpret_cast<cute_type*>(const_cast<T*>(v_input.data<T>())),
-            seq_len_encoder.data<int>(),
-            seq_len_decoder.data<int>(),
-            cu_seq_k.data<int>(),
-            cache_k.data(),
-            cache_v.data(),
-            block_tables.data<int>(),
-            cache_k_dequant_scale ? reinterpret_cast<cute_type*>(const_cast<T*>(cache_k_dequant_scale.get().data<T>())) : nullptr,
-            cache_v_dequant_scale ? reinterpret_cast<cute_type*>(const_cast<T*>(cache_v_dequant_scale.get().data<T>())) : nullptr,
-            cache_k_zero_points ? reinterpret_cast<cute_type*>(const_cast<T*>(cache_k_zero_points.get().data<T>())) : nullptr,
-            cache_v_zero_points ? reinterpret_cast<cute_type*>(const_cast<T*>(cache_v_zero_points.get().data<T>())) : nullptr,
-            kv_head_num,
-            head_dim,
-            max_seq_k,
-            seq_len_encoder.dims()[0],
-            max_input_length,
-            block_tables.dims()[1],
-            cache_quant_type_str,
-            k_input.stream());
-    } else if (k_input.dtype() == paddle::DataType::BFLOAT16) {
-        using T = phi::dtype::bfloat16;
-        using cute_type = typename cuteType<T>::type;
-        get_kv_from_cache<cute_type>(
-            reinterpret_cast<cute_type*>(const_cast<T*>(k_input.data<T>())),
-            reinterpret_cast<cute_type*>(const_cast<T*>(v_input.data<T>())),
-            seq_len_encoder.data<int>(),
-            seq_len_decoder.data<int>(),
-            cu_seq_k.data<int>(),
-            cache_k.data(),
-            cache_v.data(),
-            block_tables.data<int>(),
-            cache_k_dequant_scale ? reinterpret_cast<cute_type*>(const_cast<T*>(cache_k_dequant_scale.get().data<T>())) : nullptr,
-            cache_v_dequant_scale ? reinterpret_cast<cute_type*>(const_cast<T*>(cache_v_dequant_scale.get().data<T>())) : nullptr,
-            cache_k_zero_points ? reinterpret_cast<cute_type*>(const_cast<T*>(cache_k_zero_points.get().data<T>())) : nullptr,
-            cache_v_zero_points ? reinterpret_cast<cute_type*>(const_cast<T*>(cache_v_zero_points.get().data<T>())) : nullptr,
-            kv_head_num,
-            head_dim,
-            max_seq_k,
-            seq_len_encoder.dims()[0],
-            max_input_length,
-            block_tables.dims()[1],
-            cache_quant_type_str,
-            k_input.stream());
-    }
+    const paddle::Tensor& k_input,
+    const paddle::Tensor& v_input,
+    const paddle::Tensor& cu_seq_k,
+    const paddle::Tensor& seq_len_encoder,
+    const paddle::Tensor& seq_len_decoder,
+    const paddle::Tensor& cache_k,
+    const paddle::Tensor& cache_v,
+    const paddle::Tensor& block_tables,
+    const paddle::optional<paddle::Tensor>& cache_k_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_v_dequant_scale,
+    const paddle::optional<paddle::Tensor>& cache_k_zero_points,
+    const paddle::optional<paddle::Tensor>& cache_v_zero_points,
+    const int head_num,
+    const int kv_head_num,
+    const int head_dim,
+    const int max_input_length,
+    const int max_seq_k,
+    const std::string& cache_quant_type_str) {
+  if (k_input.dtype() == paddle::DataType::FLOAT16) {
+    using T = phi::dtype::float16;
+    using cute_type = typename cuteType<T>::type;
+    get_kv_from_cache<cute_type>(
+        reinterpret_cast<cute_type*>(const_cast<T*>(k_input.data<T>())),
+        reinterpret_cast<cute_type*>(const_cast<T*>(v_input.data<T>())),
+        seq_len_encoder.data<int>(),
+        seq_len_decoder.data<int>(),
+        cu_seq_k.data<int>(),
+        cache_k.data(),
+        cache_v.data(),
+        block_tables.data<int>(),
+        cache_k_dequant_scale ? reinterpret_cast<cute_type*>(const_cast<T*>(
+                                    cache_k_dequant_scale.get().data<T>()))
+                              : nullptr,
+        cache_v_dequant_scale ? reinterpret_cast<cute_type*>(const_cast<T*>(
+                                    cache_v_dequant_scale.get().data<T>()))
+                              : nullptr,
+        cache_k_zero_points ? reinterpret_cast<cute_type*>(const_cast<T*>(
+                                  cache_k_zero_points.get().data<T>()))
+                            : nullptr,
+        cache_v_zero_points ? reinterpret_cast<cute_type*>(const_cast<T*>(
+                                  cache_v_zero_points.get().data<T>()))
+                            : nullptr,
+        kv_head_num,
+        head_dim,
+        max_seq_k,
+        seq_len_encoder.dims()[0],
+        max_input_length,
+        block_tables.dims()[1],
+        cache_quant_type_str,
+        k_input.stream());
+  } else if (k_input.dtype() == paddle::DataType::BFLOAT16) {
+    using T = phi::dtype::bfloat16;
+    using cute_type = typename cuteType<T>::type;
+    get_kv_from_cache<cute_type>(
+        reinterpret_cast<cute_type*>(const_cast<T*>(k_input.data<T>())),
+        reinterpret_cast<cute_type*>(const_cast<T*>(v_input.data<T>())),
+        seq_len_encoder.data<int>(),
+        seq_len_decoder.data<int>(),
+        cu_seq_k.data<int>(),
+        cache_k.data(),
+        cache_v.data(),
+        block_tables.data<int>(),
+        cache_k_dequant_scale ? reinterpret_cast<cute_type*>(const_cast<T*>(
+                                    cache_k_dequant_scale.get().data<T>()))
+                              : nullptr,
+        cache_v_dequant_scale ? reinterpret_cast<cute_type*>(const_cast<T*>(
+                                    cache_v_dequant_scale.get().data<T>()))
+                              : nullptr,
+        cache_k_zero_points ? reinterpret_cast<cute_type*>(const_cast<T*>(
+                                  cache_k_zero_points.get().data<T>()))
+                            : nullptr,
+        cache_v_zero_points ? reinterpret_cast<cute_type*>(const_cast<T*>(
+                                  cache_v_zero_points.get().data<T>()))
+                            : nullptr,
+        kv_head_num,
+        head_dim,
+        max_seq_k,
+        seq_len_encoder.dims()[0],
+        max_input_length,
+        block_tables.dims()[1],
+        cache_quant_type_str,
+        k_input.stream());
+  }
 }
 
 __global__ void get_cur_cu_seq_len_k_kernel(
-        const int* __restrict__ seq_lens_encoder,
-        const int* __restrict__ seq_lens_decoder,
-        const int* __restrict__ seq_lens_this_time,
-        int* __restrict__ cu_seqlens_k,
-        int* __restrict__ cu_seq_q_pack,
-        int* __restrict__ q_pack_tokens,
-        const int pack_size,
-        const int bsz) {
+    const int* __restrict__ seq_lens_encoder,
+    const int* __restrict__ seq_lens_decoder,
+    const int* __restrict__ seq_lens_this_time,
+    int* __restrict__ cu_seqlens_k,
+    int* __restrict__ cu_seq_q_pack,
+    int* __restrict__ q_pack_tokens,
+    const int pack_size,
+    const int bsz) {
+  int total_tokens = 0;
+  cu_seqlens_k[0] = 0;
+  cu_seq_q_pack[0] = 0;
 
-    int total_tokens = 0;
-    cu_seqlens_k[0] = 0;
-    cu_seq_q_pack[0] = 0;
-
-    for (uint32_t bid = 0; bid < bsz; bid++) {
-        int cache_len = seq_lens_decoder[bid];
-        const int q_len = seq_lens_encoder[bid];
-        if (q_len <= 0) {
-            cache_len = 0;
-        }
-        total_tokens += (cache_len + q_len);
-        cu_seqlens_k[bid + 1] = total_tokens;
-        cu_seq_q_pack[bid + 1] = cu_seq_q_pack[bid] + (q_len + pack_size -1) / pack_size * pack_size;
+  for (uint32_t bid = 0; bid < bsz; bid++) {
+    int cache_len = seq_lens_decoder[bid];
+    const int q_len = seq_lens_encoder[bid];
+    if (q_len <= 0) {
+      cache_len = 0;
     }
-    q_pack_tokens[0] = cu_seq_q_pack[bsz];
+    total_tokens += (cache_len + q_len);
+    cu_seqlens_k[bid + 1] = total_tokens;
+    cu_seq_q_pack[bid + 1] =
+        cu_seq_q_pack[bid] + (q_len + pack_size - 1) / pack_size * pack_size;
+  }
+  q_pack_tokens[0] = cu_seq_q_pack[bsz];
 }
 
 std::vector<paddle::Tensor> GetCurCuSeqLenk(
-        const paddle::Tensor& seq_lens_encoder,
-        const paddle::Tensor& seq_lens_decoder,
-        const paddle::Tensor& seq_lens_this_time,
-        const int pack_size) {
-    auto stream = seq_lens_decoder.stream();
-    auto place = seq_lens_decoder.place();
-    int bsz = seq_lens_this_time.shape()[0];
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const int pack_size) {
+  auto stream = seq_lens_decoder.stream();
+  auto place = seq_lens_decoder.place();
+  int bsz = seq_lens_this_time.shape()[0];
 
-    paddle::Tensor cu_seq_q_pack = paddle::empty({bsz + 1}, paddle::DataType::INT32, place);
-    paddle::Tensor cu_seqlens_k = paddle::empty({bsz + 1}, paddle::DataType::INT32, place);
-    paddle::Tensor q_pack_tokens = paddle::empty({1}, paddle::DataType::INT32, place);
+  paddle::Tensor cu_seq_q_pack =
+      paddle::empty({bsz + 1}, paddle::DataType::INT32, place);
+  paddle::Tensor cu_seqlens_k =
+      paddle::empty({bsz + 1}, paddle::DataType::INT32, place);
+  paddle::Tensor q_pack_tokens =
+      paddle::empty({1}, paddle::DataType::INT32, place);
 
-    get_cur_cu_seq_len_k_kernel<<<1, 1, 0, stream>>>(
-        seq_lens_encoder.data<int>(),
-        seq_lens_decoder.data<int>(),
-        seq_lens_this_time.data<int>(),
-        cu_seqlens_k.data<int>(),
-        cu_seq_q_pack.data<int>(),
-        q_pack_tokens.data<int>(),
-        pack_size,
-        bsz
-    );
+  get_cur_cu_seq_len_k_kernel<<<1, 1, 0, stream>>>(
+      seq_lens_encoder.data<int>(),
+      seq_lens_decoder.data<int>(),
+      seq_lens_this_time.data<int>(),
+      cu_seqlens_k.data<int>(),
+      cu_seq_q_pack.data<int>(),
+      q_pack_tokens.data<int>(),
+      pack_size,
+      bsz);
 
-    auto q_pack_tokens_cpu = q_pack_tokens.copy_to(paddle::CPUPlace(), true);
-    return {cu_seq_q_pack, cu_seqlens_k, q_pack_tokens_cpu};
+  auto q_pack_tokens_cpu = q_pack_tokens.copy_to(paddle::CPUPlace(), true);
+  return {cu_seq_q_pack, cu_seqlens_k, q_pack_tokens_cpu};
 }
 
 PD_BUILD_STATIC_OP(get_kv_from_cache)
-    .Inputs({
-        "k_input",
-        "v_input",
-        "cu_seq_k",
-        "seq_len_encoder",
-        "seq_len_decoder",
-        "cache_k",
-        "cache_v",
-        "block_tables",
-        paddle::Optional("cache_k_dequant_scale"),
-        paddle::Optional("cache_v_dequant_scale"),
-        paddle::Optional("cache_k_zero_points"),
-        paddle::Optional("cache_v_zero_points")})
-    .Attrs({
-        "head_num: int",
-        "kv_head_num: int",
-        "head_dim: int",
-        "max_input_length: int",
-        "max_seq_k: int",
-        "cache_quant_type_str: std::string"})
+    .Inputs({"k_input",
+             "v_input",
+             "cu_seq_k",
+             "seq_len_encoder",
+             "seq_len_decoder",
+             "cache_k",
+             "cache_v",
+             "block_tables",
+             paddle::Optional("cache_k_dequant_scale"),
+             paddle::Optional("cache_v_dequant_scale"),
+             paddle::Optional("cache_k_zero_points"),
+             paddle::Optional("cache_v_zero_points")})
+    .Attrs({"head_num: int",
+            "kv_head_num: int",
+            "head_dim: int",
+            "max_input_length: int",
+            "max_seq_k: int",
+            "cache_quant_type_str: std::string"})
     .Outputs({"k_input_out", "v_input_out"})
-    .SetInplaceMap({{"k_input", "k_input_out"},
-                    {"v_input", "v_input_out"}})
+    .SetInplaceMap({{"k_input", "k_input_out"}, {"v_input", "v_input_out"}})
     .SetKernelFn(PD_KERNEL(GetKVFromCache));
 
 PD_BUILD_STATIC_OP(get_cur_cu_seq_len_k)
-    .Inputs({
-            "seq_lens_encoder",
-            "seq_lens_decoder",
-            "seq_lens_this_time"})
-    .Attrs({
-        "pack_size: int"})
+    .Inputs({"seq_lens_encoder", "seq_lens_decoder", "seq_lens_this_time"})
+    .Attrs({"pack_size: int"})
     .Outputs({"cu_seq_q_pack", "cu_seqlens_k", "q_pack_tokens"})
     .SetKernelFn(PD_KERNEL(GetCurCuSeqLenk));
diff --git a/custom_ops/gpu_ops/moba_attn/moba_process/moba_mlp_einsum.cu b/custom_ops/gpu_ops/moba_attn/moba_process/moba_mlp_einsum.cu
index 4d740290a5..5b05b57251 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_process/moba_mlp_einsum.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_process/moba_mlp_einsum.cu
@@ -21,204 +21,206 @@
 #endif
 
 template <typename T, int moba_block_size, int kHeadDim, int kMaxN>
-__global__ void moba_mlp_einsum_kernel(
-        const T * src_data,
-        const T * weight_data,
-        const int * seq_lens_encoder,
-        const int * seq_lens_decoder,
-        const int * cu_seq_k,
-        T * dst_data,
-        const int head_num) {
+__global__ void moba_mlp_einsum_kernel(const T* src_data,
+                                       const T* weight_data,
+                                       const int* seq_lens_encoder,
+                                       const int* seq_lens_decoder,
+                                       const int* cu_seq_k,
+                                       T* dst_data,
+                                       const int head_num) {
+  constexpr int kPackSize = 16 / sizeof(T);
+  const int block_idx = blockIdx.x;
+  const int bidh = blockIdx.y;
+  const int bidb = blockIdx.z;
+  const int tidx = threadIdx.x;
+  const int lane_id = tidx % 32;
+  const int warp_id = tidx / 32;
 
-    constexpr int kPackSize = 16 / sizeof(T);
-    const int block_idx = blockIdx.x;
-    const int bidh = blockIdx.y;
-    const int bidb = blockIdx.z;
-    const int tidx = threadIdx.x;
-    const int lane_id = tidx % 32;
-    const int warp_id = tidx / 32;
+  __align__(16) __shared__ T local_sum_mem[128 / 32 * kHeadDim];
 
-    __align__(16) __shared__ T local_sum_mem[128 / 32 * kHeadDim];
+  const int seq_len_encoder = seq_lens_encoder[bidb];
+  const int seq_len_decoder = seq_len_encoder + seq_lens_decoder[bidb];
 
-    const int seq_len_encoder = seq_lens_encoder[bidb];
-    const int seq_len_decoder = seq_len_encoder + seq_lens_decoder[bidb];
+  const int seq_len_this_block = seq_len_decoder - block_idx * moba_block_size;
 
-    const int seq_len_this_block = seq_len_decoder - block_idx * moba_block_size;
+  if (seq_len_encoder == 0 || seq_len_this_block <= 0) {
+    return;
+  }
 
-    if (seq_len_encoder == 0 || seq_len_this_block <= 0) {
-        return;
+  using SrcType = Vec<T, kPackSize>;
+
+  constexpr int tidx_per_row = kHeadDim / kPackSize;
+
+  const int row_idx = tidx / tidx_per_row;
+  const int col_idx = tidx % tidx_per_row * kPackSize;
+
+  const int src_base_idx = cu_seq_k[bidb] * head_num * kHeadDim +
+                           block_idx * moba_block_size * head_num * kHeadDim +
+                           bidh * kHeadDim + row_idx * head_num * kHeadDim +
+                           col_idx;
+  const int weight_base_idx =
+      bidh * kHeadDim * moba_block_size + row_idx * kHeadDim + col_idx;
+
+  constexpr int step = 128 / tidx_per_row;
+
+  SrcType sums, src, weight;
+
+  sums.set_zero();
+
+  for (int i = 0; i < moba_block_size; i += step) {
+    if (i >= seq_len_this_block) {
+      break;
     }
+    src.load_from(src_data + src_base_idx + i * head_num * kHeadDim);
+    weight.load_from(weight_data + weight_base_idx + i * kHeadDim);
+    sums.fma(src, weight);
+  }
 
+  SrcType neighbor;
 
-    using SrcType = Vec<T, kPackSize>;
+#pragma unroll
+  for (int i = 0; i < kPackSize; i += 2) {
+    *reinterpret_cast<int32_t*>(neighbor.data.elt + i) = __shfl_down_sync(
+        0xffffffff, *reinterpret_cast<int32_t*>(sums.data.elt + i), 16);
+  }
 
-    constexpr int tidx_per_row = kHeadDim / kPackSize;
+  sums.add(neighbor);
 
-    const int row_idx = tidx / tidx_per_row;
-    const int col_idx = tidx % tidx_per_row * kPackSize;
+  if (lane_id < 16) {
+    sums.store_to(local_sum_mem + warp_id * kHeadDim + lane_id * kPackSize);
+  }
 
-    const int src_base_idx = cu_seq_k[bidb] * head_num * kHeadDim + block_idx * moba_block_size * head_num * kHeadDim + bidh * kHeadDim + row_idx * head_num * kHeadDim + col_idx;
-    const int weight_base_idx = bidh * kHeadDim * moba_block_size + row_idx * kHeadDim + col_idx;
+  __syncthreads();
+  using pack_half =
+      std::conditional_t<std::is_same<T, phi::dtype::float16>::value,
+                         __half2,
+                         nv_bfloat162>;
+  pack_half* local_sum_mem_half = reinterpret_cast<pack_half*>(local_sum_mem);
 
-    constexpr int step = 128 / tidx_per_row;
-
-    SrcType sums, src, weight;
-
-    sums.set_zero();
-
-    for (int i = 0; i < moba_block_size; i += step) {
-        if (i >= seq_len_this_block) {
-            break;
-        }
-        src.load_from(src_data + src_base_idx + i * head_num * kHeadDim);
-        weight.load_from(weight_data + weight_base_idx + i * kHeadDim);
-        sums.fma(src, weight);
+  if (tidx < kHeadDim / 2) {
+    pack_half local_sum_half = local_sum_mem_half[tidx];
+#pragma unroll
+    for (int i = 1; i < 4; i++) {
+      local_sum_half += local_sum_mem_half[tidx + i * (kHeadDim / 2)];
     }
+    local_sum_mem_half[tidx] = local_sum_half;
+  }
 
-    SrcType neighbor;
+  __syncthreads();
 
-    #pragma unroll
-    for (int i = 0; i < kPackSize; i+=2) {
-        *reinterpret_cast<int32_t*>(neighbor.data.elt + i) = __shfl_down_sync(0xffffffff, *reinterpret_cast<int32_t*>(sums.data.elt + i), 16);
-    }
+  const int store_row_id = tidx / (kHeadDim / kPackSize);
+  const int store_col_id = tidx % (kHeadDim / kPackSize) * kPackSize;
 
-    sums.add(neighbor);
+  sums.load_from(local_sum_mem + store_col_id);
 
-    if (lane_id < 16) {
-        sums.store_to(local_sum_mem + warp_id * kHeadDim + lane_id * kPackSize);
-    }
+  const int base_store_idx =
+      bidb * kMaxN * head_num * kHeadDim +
+      (block_idx * (moba_block_size / 128) + store_row_id) * head_num *
+          kHeadDim +
+      bidh * kHeadDim + store_col_id;
 
-    __syncthreads();
-    using pack_half = std::conditional_t<std::is_same<T, phi::dtype::float16>::value, __half2, nv_bfloat162>;
-    pack_half * local_sum_mem_half = reinterpret_cast<pack_half*>(local_sum_mem);
-
-    if (tidx < kHeadDim / 2) {
-        pack_half local_sum_half = local_sum_mem_half[tidx];
-        #pragma unroll
-        for (int i = 1; i < 4; i++) {
-            local_sum_half += local_sum_mem_half[tidx + i * (kHeadDim / 2)];
-        }
-        local_sum_mem_half[tidx] = local_sum_half;
-    }
-
-    __syncthreads();
-
-    const int store_row_id = tidx / (kHeadDim / kPackSize);
-    const int store_col_id = tidx % (kHeadDim / kPackSize) * kPackSize;
-
-    sums.load_from(local_sum_mem + store_col_id);
-
-    const int base_store_idx = bidb * kMaxN * head_num * kHeadDim + (block_idx * (moba_block_size / 128) + store_row_id) * head_num * kHeadDim + bidh * kHeadDim + store_col_id;
-
-    if (store_row_id < moba_block_size / 128) {
-        sums.store_to(dst_data + base_store_idx);
-    }
+  if (store_row_id < moba_block_size / 128) {
+    sums.store_to(dst_data + base_store_idx);
+  }
 }
 
-
 template <typename T, int kHeadDim, int kMaxN>
-void moba_mlp_einsum(
-        const T * src_data,
-        const T * weight_data,
-        const int * seq_lens_encoder,
-        const int * seq_lens_decoder,
-        const int * cu_seq_k,
-        T * dst_data,
-        const int moba_block_size,
-        const int max_seq_len,
-        const int head_num,
-        const int batch_size,
-        cudaStream_t stream) {
-
-    dim3 grid_dims;
-    grid_dims.x = (max_seq_len + moba_block_size - 1) / moba_block_size;
-    grid_dims.y = head_num;
-    grid_dims.z = batch_size;
-
-    if (moba_block_size == 1024) {
-        moba_mlp_einsum_kernel<T, 1024, kHeadDim, kMaxN><<<grid_dims, 128, 0, stream>>>(
-            src_data,
-            weight_data,
-            seq_lens_encoder,
-            seq_lens_decoder,
-            cu_seq_k,
-            dst_data,
-            head_num);
-    } else if (moba_block_size == 128) {
-        moba_mlp_einsum_kernel<T, 128, kHeadDim, kMaxN><<<grid_dims, 128, 0, stream>>>(
-            src_data,
-            weight_data,
-            seq_lens_encoder,
-            seq_lens_decoder,
-            cu_seq_k,
-            dst_data,
-            head_num);
-    } else {
-        PADDLE_THROW(phi::errors::Unimplemented(
-            "MobaMlpEinsum not implemented for moba_block_size = %d", moba_block_size));
-    }
+void moba_mlp_einsum(const T* src_data,
+                     const T* weight_data,
+                     const int* seq_lens_encoder,
+                     const int* seq_lens_decoder,
+                     const int* cu_seq_k,
+                     T* dst_data,
+                     const int moba_block_size,
+                     const int max_seq_len,
+                     const int head_num,
+                     const int batch_size,
+                     cudaStream_t stream) {
+  dim3 grid_dims;
+  grid_dims.x = (max_seq_len + moba_block_size - 1) / moba_block_size;
+  grid_dims.y = head_num;
+  grid_dims.z = batch_size;
 
+  if (moba_block_size == 1024) {
+    moba_mlp_einsum_kernel<T, 1024, kHeadDim, kMaxN>
+        <<<grid_dims, 128, 0, stream>>>(src_data,
+                                        weight_data,
+                                        seq_lens_encoder,
+                                        seq_lens_decoder,
+                                        cu_seq_k,
+                                        dst_data,
+                                        head_num);
+  } else if (moba_block_size == 128) {
+    moba_mlp_einsum_kernel<T, 128, kHeadDim, kMaxN>
+        <<<grid_dims, 128, 0, stream>>>(src_data,
+                                        weight_data,
+                                        seq_lens_encoder,
+                                        seq_lens_decoder,
+                                        cu_seq_k,
+                                        dst_data,
+                                        head_num);
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "MobaMlpEinsum not implemented for moba_block_size = %d",
+        moba_block_size));
+  }
 }
 
-
 std::vector<paddle::Tensor> MobaMlpEinsum(
-        const paddle::Tensor& k_input,
-        const paddle::Tensor& attn_gate_weight,
-        const paddle::Tensor& seq_lens_encoder,
-        const paddle::Tensor& seq_lens_decoder,
-        const paddle::Tensor& cu_seq_k,
-        const int max_seq_len,
-        const int kv_head_num) {
+    const paddle::Tensor& k_input,
+    const paddle::Tensor& attn_gate_weight,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& cu_seq_k,
+    const int max_seq_len,
+    const int kv_head_num) {
+  const int kHeadDim = 128;
+  const int kMaxN = 1024;
+  const int moba_block_size = attn_gate_weight.dims()[1];
+  const int batch_size = seq_lens_encoder.dims()[0];
+  paddle::Tensor k_gate_weight =
+      paddle::zeros({batch_size, kMaxN, kv_head_num, kHeadDim},
+                    k_input.dtype(),
+                    k_input.place());
 
-    const int kHeadDim = 128;
-    const int kMaxN = 1024;
-    const int moba_block_size = attn_gate_weight.dims()[1];
-    const int batch_size = seq_lens_encoder.dims()[0];
-    paddle::Tensor k_gate_weight = paddle::zeros({batch_size, kMaxN, kv_head_num, kHeadDim}, k_input.dtype(), k_input.place());
-
-    if (k_input.dtype() == paddle::DataType::FLOAT16) {
-        using T = phi::dtype::float16;
-        moba_mlp_einsum<T, kHeadDim, kMaxN>(
-            const_cast<T*>(k_input.data<T>()),
-            const_cast<T*>(attn_gate_weight.data<T>()),
-            const_cast<int*>(seq_lens_encoder.data<int>()),
-            const_cast<int*>(seq_lens_decoder.data<int>()),
-            const_cast<int*>(cu_seq_k.data<int>()),
-            k_gate_weight.data<T>(),
-            moba_block_size,
-            max_seq_len,
-            kv_head_num,
-            batch_size,
-            k_input.stream()
-        );
-    } else if (k_input.dtype() == paddle::DataType::BFLOAT16) {
-        using T = phi::dtype::bfloat16;
-        moba_mlp_einsum<T, kHeadDim, kMaxN>(
-            const_cast<T*>(k_input.data<T>()),
-            const_cast<T*>(attn_gate_weight.data<T>()),
-            const_cast<int*>(seq_lens_encoder.data<int>()),
-            const_cast<int*>(seq_lens_decoder.data<int>()),
-            const_cast<int*>(cu_seq_k.data<int>()),
-            k_gate_weight.data<T>(),
-            moba_block_size,
-            max_seq_len,
-            kv_head_num,
-            batch_size,
-            k_input.stream()
-        );
-    }
-    return {k_gate_weight};
+  if (k_input.dtype() == paddle::DataType::FLOAT16) {
+    using T = phi::dtype::float16;
+    moba_mlp_einsum<T, kHeadDim, kMaxN>(
+        const_cast<T*>(k_input.data<T>()),
+        const_cast<T*>(attn_gate_weight.data<T>()),
+        const_cast<int*>(seq_lens_encoder.data<int>()),
+        const_cast<int*>(seq_lens_decoder.data<int>()),
+        const_cast<int*>(cu_seq_k.data<int>()),
+        k_gate_weight.data<T>(),
+        moba_block_size,
+        max_seq_len,
+        kv_head_num,
+        batch_size,
+        k_input.stream());
+  } else if (k_input.dtype() == paddle::DataType::BFLOAT16) {
+    using T = phi::dtype::bfloat16;
+    moba_mlp_einsum<T, kHeadDim, kMaxN>(
+        const_cast<T*>(k_input.data<T>()),
+        const_cast<T*>(attn_gate_weight.data<T>()),
+        const_cast<int*>(seq_lens_encoder.data<int>()),
+        const_cast<int*>(seq_lens_decoder.data<int>()),
+        const_cast<int*>(cu_seq_k.data<int>()),
+        k_gate_weight.data<T>(),
+        moba_block_size,
+        max_seq_len,
+        kv_head_num,
+        batch_size,
+        k_input.stream());
+  }
+  return {k_gate_weight};
 }
 
 PD_BUILD_STATIC_OP(moba_mlp_einsum)
-    .Inputs({
-        "k_input",
-        "attn_gate_weight",
-        "seq_lens_encoder",
-        "seq_lens_decoder",
-        "cu_seq_k"})
-    .Attrs({
-        "max_seq_len: int",
-        "kv_head_num: int"})
+    .Inputs({"k_input",
+             "attn_gate_weight",
+             "seq_lens_encoder",
+             "seq_lens_decoder",
+             "cu_seq_k"})
+    .Attrs({"max_seq_len: int", "kv_head_num: int"})
     .Outputs({"k_gate"})
     .SetKernelFn(PD_KERNEL(MobaMlpEinsum));
diff --git a/custom_ops/gpu_ops/moba_attn/moba_process/moba_qk_gemm.cu b/custom_ops/gpu_ops/moba_attn/moba_process/moba_qk_gemm.cu
index 2662fb26ba..a8eba3c43e 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_process/moba_qk_gemm.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_process/moba_qk_gemm.cu
@@ -29,441 +29,463 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-template <typename input_type, int kBlockM, int kBlockN, int kMobaBlockSize, int kMaxN, int kHeadDim, bool is_split_kv>
-__global__ void qk_gemm_kernel(
-        const input_type *q_input,
-        const input_type *k_gate_mean,
-        input_type *qk_gate_weight,
-        const int *seq_len_encoder,
-        const int *seq_len_decoder,
-        const int *cu_seq_q,
-        const int *cu_seq_k,
-        const int use_moba_seq_limit,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int kGQA_groupsize) {
+template <typename input_type,
+          int kBlockM,
+          int kBlockN,
+          int kMobaBlockSize,
+          int kMaxN,
+          int kHeadDim,
+          bool is_split_kv>
+__global__ void qk_gemm_kernel(const input_type *q_input,
+                               const input_type *k_gate_mean,
+                               input_type *qk_gate_weight,
+                               const int *seq_len_encoder,
+                               const int *seq_len_decoder,
+                               const int *cu_seq_q,
+                               const int *cu_seq_k,
+                               const int use_moba_seq_limit,
+                               const int max_seq_q,
+                               const int max_seq_k,
+                               const int head_num,
+                               const int kv_head_num,
+                               const int kGQA_groupsize) {
+  using TileShape_MNK = Shape<Int<kBlockM>, Int<kBlockN>, Int<kHeadDim>>;
 
-    using TileShape_MNK = Shape<Int<kBlockM>, Int<kBlockN>, Int<kHeadDim>>;
+  using SmemLayoutAtomQ =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               input_type,
+               decltype(cute::get<0>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutQ =
+      decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_MNK{})));
 
-    using SmemLayoutAtomQ = decltype(
-        cutlass::gemm::collective::detail::ss_smem_selector<
-            GMMA::Major::K, input_type,
-            decltype(cute::get<0>(TileShape_MNK{})),
-            decltype(cute::get<2>(TileShape_MNK{}))>());
-    using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_MNK{})));
+  using SmemLayoutAtomK =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               input_type,
+               decltype(cute::get<1>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutK =
+      decltype(tile_to_shape(SmemLayoutAtomK{}, select<1, 2>(TileShape_MNK{})));
 
-    using SmemLayoutAtomK = decltype(
-        cutlass::gemm::collective::detail::ss_smem_selector<
-            GMMA::Major::K, input_type, decltype(cute::get<1>(TileShape_MNK{})),
-            decltype(cute::get<2>(TileShape_MNK{}))>());
-    using SmemLayoutK = decltype(tile_to_shape(SmemLayoutAtomK{}, select<1, 2>(TileShape_MNK{})));
+  using SmemLayoutAtomQK =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               input_type,
+               decltype(cute::get<0>(TileShape_MNK{})),
+               decltype(cute::get<1>(TileShape_MNK{}))>());
 
-    using SmemLayoutAtomQK = decltype(
-        cutlass::gemm::collective::detail::ss_smem_selector<
-        GMMA::Major::K, input_type,
-        decltype(cute::get<0>(TileShape_MNK{})),
-        decltype(cute::get<1>(TileShape_MNK{}))>());
+  using SmemLayoutQK = decltype(tile_to_shape(SmemLayoutAtomQK{},
+                                              select<0, 1>(TileShape_MNK{})));
 
-    using SmemLayoutQK = decltype(tile_to_shape(SmemLayoutAtomQK{}, select<0, 1>(TileShape_MNK{})));
+  using MMA_Atom_Arch =
+      std::conditional_t<std::is_same_v<input_type, cutlass::half_t>,
+                         MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
+                         MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>>;
 
+  using ValLayoutMNK = std::conditional_t<is_split_kv,
+                                          Layout<Shape<_1, _4, _1>>,
+                                          Layout<Shape<_4, _1, _1>>>;
 
-    using MMA_Atom_Arch = std::conditional_t<
-        std::is_same_v<input_type, cutlass::half_t>,
-        MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
-        MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>
-    >;
+  using PermutationMNK =
+      std::conditional_t<is_split_kv, Tile<_16, _64, _16>, Tile<_64, _16, _16>>;
 
-    using ValLayoutMNK = std::conditional_t<
-        is_split_kv,
-        Layout<Shape<_1,_4,_1>>,
-        Layout<Shape<_4,_1,_1>>
-    >;
+  using TiledMma = TiledMMA<MMA_Atom_Arch, ValLayoutMNK, PermutationMNK>;
 
-    using PermutationMNK = std::conditional_t<
-        is_split_kv,
-        Tile<_16,_64,_16>,
-        Tile<_64,_16,_16>
-    >;
+  using SmemCopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, input_type>;
+  using SmemCopyAtomQK = Copy_Atom<cute::SM90_U32x4_STSM_N, input_type>;
 
-    using TiledMma = TiledMMA<
-        MMA_Atom_Arch,
-        ValLayoutMNK,
-        PermutationMNK>;
+  constexpr int kNThreads = 128;
+  constexpr int kThreadPerValue = 16 / sizeof(input_type);
+  constexpr int kThreadsPerRow = kHeadDim / kThreadPerValue;
+  constexpr int kThreadsPerRowQK = kBlockN / kThreadPerValue;
 
-    using SmemCopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, input_type>;
-    using SmemCopyAtomQK = Copy_Atom<cute::SM90_U32x4_STSM_N, input_type>;
+  using GmemLayoutAtom =
+      Layout<Shape<Int<kNThreads / kThreadsPerRow>, Int<kThreadsPerRow>>,
+             Stride<Int<kThreadsPerRow>, _1>>;
 
-    constexpr int kNThreads = 128;
-    constexpr int kThreadPerValue = 16 / sizeof(input_type);
-    constexpr int kThreadsPerRow = kHeadDim / kThreadPerValue;
-    constexpr int kThreadsPerRowQK = kBlockN / kThreadPerValue;
+  using GmemTiledCopy = decltype(make_tiled_copy(
+      Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, input_type>{},
+      GmemLayoutAtom{},
+      Layout<Shape<_1, Int<kThreadPerValue>>>{}));
 
-    using GmemLayoutAtom = Layout<
-        Shape <Int<kNThreads / kThreadsPerRow>, Int<kThreadsPerRow>>,
-        Stride<Int<kThreadsPerRow>, _1>>;
+  using GmemLayoutAtomQK =
+      Layout<Shape<Int<kNThreads / kThreadsPerRowQK>, Int<kThreadsPerRowQK>>,
+             Stride<Int<kThreadsPerRowQK>, _1>>;
 
-    using GmemTiledCopy = decltype(
-        make_tiled_copy(Copy_Atom<
-            SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, input_type>{},
-            GmemLayoutAtom{},
-            Layout<Shape<_1, Int<kThreadPerValue>>>{}));
+  using GmemTiledCopyQK = decltype(make_tiled_copy(
+      Copy_Atom<UniversalCopy<cutlass::uint128_t>, input_type>{},
+      GmemLayoutAtomQK{},
+      Layout<Shape<_1, Int<kThreadPerValue>>>{}));
 
-    using GmemLayoutAtomQK = Layout<
-        Shape <Int<kNThreads / kThreadsPerRowQK>, Int<kThreadsPerRowQK>>,
-        Stride<Int<kThreadsPerRowQK>, _1>>;
+  int mn_block = blockIdx.x;
+  const int bidb = blockIdx.y;
+  const int bidh = blockIdx.z;
+  const int bidh_k = bidh / kGQA_groupsize;
+  const int tidx = threadIdx.x;
 
-    using GmemTiledCopyQK = decltype(
-        make_tiled_copy(Copy_Atom<
-            UniversalCopy<cutlass::uint128_t>, input_type>{},
-            GmemLayoutAtomQK{},
-            Layout<Shape<_1, Int<kThreadPerValue>>>{}));
+  extern __shared__ char smem_[];
 
-    int mn_block = blockIdx.x;
-    const int bidb = blockIdx.y;
-    const int bidh = blockIdx.z;
-    const int bidh_k = bidh / kGQA_groupsize;
-    const int tidx = threadIdx.x;
+  const int seq_len_q = seq_len_encoder[bidb];
+  const int seq_len_k = seq_len_decoder[bidb];
+  const int seq_len_qk = seq_len_q + seq_len_k;
 
-    extern __shared__ char smem_[];
+  int q_head_stride;
+  const int k_head_stride = kv_head_num * kHeadDim;
+  int qk_head_stride;
+  int offset_q;
+  int offset_k;
+  int offset_qk;
+  int remain_q_seq;
 
-    const int seq_len_q = seq_len_encoder[bidb];
-    const int seq_len_k = seq_len_decoder[bidb];
-    const int seq_len_qk = seq_len_q + seq_len_k;
+  if constexpr (is_split_kv) {
+    if (seq_len_k < use_moba_seq_limit || seq_len_k == 0) {
+      return;
+    }
+    mn_block *= kBlockN;
+    q_head_stride = kHeadDim;
+    qk_head_stride = kMaxN;
+    if (mn_block >= (seq_len_k + kMobaBlockSize - 1) / kMobaBlockSize) {
+      return;
+    }
+    offset_q =
+        cu_seq_q[bidb] * head_num * kHeadDim + bidh * kGQA_groupsize * kHeadDim;
+    offset_k = (bidb * kMaxN + mn_block) * k_head_stride + bidh * kHeadDim;
+    offset_qk =
+        bidb * head_num * kMaxN + bidh * kGQA_groupsize * kMaxN + mn_block;
+    remain_q_seq = kGQA_groupsize;
+  } else {
+    if (seq_len_q == 0 || seq_len_qk < use_moba_seq_limit) {
+      return;
+    }
+    q_head_stride = head_num * kHeadDim;
+    qk_head_stride = head_num * kMaxN;
+    mn_block *= kBlockM;
+    if (mn_block >= seq_len_q) {
+      return;
+    }
+    offset_q = (cu_seq_q[bidb] + mn_block) * q_head_stride + bidh * kHeadDim;
+    offset_k = bidb * kMaxN * k_head_stride + bidh_k * kHeadDim;
+    offset_qk = (cu_seq_q[bidb] + mn_block) * qk_head_stride + bidh * kMaxN;
+    remain_q_seq = seq_len_q - mn_block;
+  }
 
-    int q_head_stride;
-    const int k_head_stride = kv_head_num * kHeadDim;
-    int qk_head_stride;
-    int offset_q;
-    int offset_k;
-    int offset_qk;
-    int remain_q_seq;
+  Tensor gQ = make_tensor(make_gmem_ptr(q_input + offset_q),
+                          Shape<Int<kBlockM>, Int<kHeadDim>>{},
+                          make_stride(q_head_stride, _1{}));
+  Tensor gK = make_tensor(make_gmem_ptr(k_gate_mean + offset_k),
+                          Shape<Int<kBlockN>, Int<kHeadDim>>{},
+                          make_stride(k_head_stride, _1{}));
+  Tensor gQK = make_tensor(make_gmem_ptr(qk_gate_weight + offset_qk),
+                           Shape<Int<kBlockM>, Int<kBlockN>>{},
+                           make_stride(qk_head_stride, _1{}));
 
-    if constexpr (is_split_kv) {
-        if (seq_len_k < use_moba_seq_limit || seq_len_k == 0) {
-            return;
-        }
-        mn_block *= kBlockN;
-        q_head_stride = kHeadDim;
-        qk_head_stride = kMaxN;
-        if (mn_block >= (seq_len_k + kMobaBlockSize - 1) / kMobaBlockSize) {
-            return;
-        }
-        offset_q = cu_seq_q[bidb] * head_num * kHeadDim + bidh * kGQA_groupsize * kHeadDim;
-        offset_k = (bidb * kMaxN + mn_block) * k_head_stride + bidh * kHeadDim;
-        offset_qk = bidb * head_num * kMaxN + bidh * kGQA_groupsize * kMaxN + mn_block;
-        remain_q_seq = kGQA_groupsize;
+  Tensor sK = make_tensor(make_smem_ptr(reinterpret_cast<input_type *>(smem_)),
+                          SmemLayoutK{});
+  Tensor sQ = make_tensor(sK.data() + size(sK), SmemLayoutQ{});
+  Tensor sQK = make_tensor(sK.data() + size(sK), SmemLayoutQK{});
+
+  auto gmem_tiled_copy = GmemTiledCopy{};
+  auto gmem_tiled_copy_qk = GmemTiledCopyQK{};
+  auto gmem_thr_copy = gmem_tiled_copy.get_thread_slice(tidx);
+  auto gmem_thr_copy_qk = gmem_tiled_copy_qk.get_thread_slice(tidx);
+
+  Tensor tQgQ = gmem_thr_copy.partition_S(gQ);
+  Tensor tQsQ = gmem_thr_copy.partition_D(sQ);
+
+  Tensor tKgK = gmem_thr_copy.partition_S(gK);
+  Tensor tKsK = gmem_thr_copy.partition_D(sK);
+
+  Tensor tQKgQK = gmem_thr_copy_qk.partition_S(gQK);
+  Tensor tQKsQK = gmem_thr_copy_qk.partition_D(sQK);
+
+  Tensor cQ = make_identity_tensor(make_shape(kBlockM, kHeadDim));
+  Tensor tQcQ = gmem_thr_copy.partition_S(cQ);
+
+  Tensor cK = make_identity_tensor(make_shape(kBlockN, kHeadDim));
+  Tensor tKcK = gmem_thr_copy.partition_S(cK);
+
+  Tensor cQK = make_identity_tensor(make_shape(kBlockM, kBlockN));
+  Tensor tQKcQK = gmem_thr_copy.partition_S(cQK);
+
+  if (remain_q_seq >= kBlockM) {
+    copy(gmem_tiled_copy, tQgQ, tQsQ, tQcQ);
+  } else {
+    copy<false>(gmem_tiled_copy, tQgQ, tQsQ, tQcQ, remain_q_seq);
+  }
+  copy(gmem_tiled_copy, tKgK, tKsK, tKcK);
+
+  cute::cp_async_fence();
+
+  TiledMma tiled_mma;
+  auto thr_mma = tiled_mma.get_thread_slice(tidx);
+  Tensor tSrQ = thr_mma.partition_fragment_A(sQ);
+  Tensor tSrK = thr_mma.partition_fragment_B(sK);
+
+  Tensor acc_s =
+      partition_fragment_C(tiled_mma, Shape<Int<kBlockM>, Int<kBlockN>>{});
+
+  auto smem_tiled_copy_Q = make_tiled_copy_A(SmemCopyAtom{}, tiled_mma);
+  auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx);
+  Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ);
+
+  auto smem_tiled_copy_K = make_tiled_copy_B(SmemCopyAtom{}, tiled_mma);
+  auto smem_thr_copy_K =
+      make_tiled_copy_B(SmemCopyAtom{}, tiled_mma).get_thread_slice(tidx);
+  Tensor tSsK = smem_thr_copy_K.partition_S(sK);
+
+  auto smem_tiled_copy_QK = make_tiled_copy_C(SmemCopyAtomQK{}, tiled_mma);
+  auto smem_thr_copy_QK = smem_tiled_copy_QK.get_thread_slice(tidx);
+  Tensor tsQK = smem_thr_copy_QK.partition_D(sQK);
+
+  const int n_blocks =
+      is_split_kv
+          ? 1
+          : cute::ceil_div(cute::ceil_div(seq_len_qk, kMobaBlockSize), kBlockN);
+
+#pragma unroll
+  for (int n_block = 0; n_block < n_blocks; ++n_block) {
+    clear(acc_s);
+    cp_async_wait<0>();
+    __syncthreads();
+    if (n_block == 0) {
+      gemm(acc_s,
+           tSrQ,
+           tSrK,
+           tSsQ,
+           tSsK,
+           tiled_mma,
+           smem_thr_copy_Q,
+           smem_thr_copy_K,
+           smem_tiled_copy_Q,
+           smem_tiled_copy_K);
     } else {
-        if (seq_len_q == 0 || seq_len_qk < use_moba_seq_limit) {
-            return;
-        }
-        q_head_stride = head_num * kHeadDim;
-        qk_head_stride = head_num * kMaxN;
-        mn_block *= kBlockM;
-        if (mn_block >= seq_len_q) {
-            return;
-        }
-        offset_q = (cu_seq_q[bidb] + mn_block) * q_head_stride + bidh * kHeadDim;
-        offset_k = bidb * kMaxN * k_head_stride + bidh_k * kHeadDim;
-        offset_qk = (cu_seq_q[bidb] + mn_block) * qk_head_stride + bidh * kMaxN;
-        remain_q_seq = seq_len_q - mn_block;
+      gemm<true>(acc_s,
+                 tSrQ,
+                 tSrK,
+                 tSsQ,
+                 tSsK,
+                 tiled_mma,
+                 smem_thr_copy_Q,
+                 smem_thr_copy_K,
+                 smem_tiled_copy_Q,
+                 smem_tiled_copy_K);
+    }
+    if constexpr (!is_split_kv) {
+      if (n_block < n_blocks - 1) {
+        __syncthreads();
+        tKgK.data() = tKgK.data() + kBlockN * k_head_stride;
+        copy(gmem_tiled_copy, tKgK, tKsK, tKcK);
+        cute::cp_async_fence();
+      }
     }
 
-    Tensor gQ = make_tensor(make_gmem_ptr(q_input + offset_q),
-        Shape<Int<kBlockM>, Int<kHeadDim>>{},
-        make_stride(q_head_stride, _1{}));
-    Tensor gK = make_tensor(make_gmem_ptr(k_gate_mean + offset_k),
-        Shape<Int<kBlockN>, Int<kHeadDim>>{},
-        make_stride(k_head_stride, _1{}));
-    Tensor gQK = make_tensor(make_gmem_ptr(qk_gate_weight + offset_qk),
-        Shape<Int<kBlockM>, Int<kBlockN>>{},
-        make_stride(qk_head_stride, _1{}));
-
-    Tensor sK = make_tensor(make_smem_ptr(reinterpret_cast<input_type *>(smem_)), SmemLayoutK{});
-    Tensor sQ = make_tensor(sK.data() + size(sK), SmemLayoutQ{});
-    Tensor sQK = make_tensor(sK.data() + size(sK), SmemLayoutQK{});
-
-    auto gmem_tiled_copy = GmemTiledCopy{};
-    auto gmem_tiled_copy_qk = GmemTiledCopyQK{};
-    auto gmem_thr_copy = gmem_tiled_copy.get_thread_slice(tidx);
-    auto gmem_thr_copy_qk = gmem_tiled_copy_qk.get_thread_slice(tidx);
-
-
-    Tensor tQgQ = gmem_thr_copy.partition_S(gQ);
-    Tensor tQsQ = gmem_thr_copy.partition_D(sQ);
-
-    Tensor tKgK = gmem_thr_copy.partition_S(gK);
-    Tensor tKsK = gmem_thr_copy.partition_D(sK);
-
-    Tensor tQKgQK = gmem_thr_copy_qk.partition_S(gQK);
-    Tensor tQKsQK = gmem_thr_copy_qk.partition_D(sQK);
-
-
-    Tensor cQ = make_identity_tensor(make_shape(kBlockM, kHeadDim));
-    Tensor tQcQ = gmem_thr_copy.partition_S(cQ);
-
-    Tensor cK = make_identity_tensor(make_shape(kBlockN, kHeadDim));
-    Tensor tKcK = gmem_thr_copy.partition_S(cK);
-
-    Tensor cQK = make_identity_tensor(make_shape(kBlockM, kBlockN));
-    Tensor tQKcQK = gmem_thr_copy.partition_S(cQK);
+    Tensor rS = convert_type<input_type>(acc_s);
+    Tensor trQK = smem_thr_copy_QK.retile_S(rS);
+    cute::copy(smem_tiled_copy_QK, trQK, tsQK);
 
+    __syncthreads();
     if (remain_q_seq >= kBlockM) {
-        copy(gmem_tiled_copy, tQgQ, tQsQ, tQcQ);
+      copy(gmem_tiled_copy_qk, tQKsQK, tQKgQK, tQKcQK);
     } else {
-        copy<false>(gmem_tiled_copy, tQgQ, tQsQ, tQcQ, remain_q_seq);
+      copy<false>(gmem_tiled_copy_qk, tQKsQK, tQKgQK, tQKcQK, remain_q_seq);
     }
-    copy(gmem_tiled_copy, tKgK, tKsK, tKcK);
-
-    cute::cp_async_fence();
-
-    TiledMma tiled_mma;
-    auto thr_mma = tiled_mma.get_thread_slice(tidx);
-    Tensor tSrQ = thr_mma.partition_fragment_A(sQ);
-    Tensor tSrK = thr_mma.partition_fragment_B(sK);
-
-    Tensor acc_s = partition_fragment_C(tiled_mma, Shape<Int<kBlockM>, Int<kBlockN>>{});
-
-    auto smem_tiled_copy_Q = make_tiled_copy_A(SmemCopyAtom{}, tiled_mma);
-    auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx);
-    Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ);
-
-    auto smem_tiled_copy_K = make_tiled_copy_B(SmemCopyAtom{}, tiled_mma);
-    auto smem_thr_copy_K = make_tiled_copy_B(SmemCopyAtom{}, tiled_mma).get_thread_slice(tidx);
-    Tensor tSsK = smem_thr_copy_K.partition_S(sK);
-
-    auto smem_tiled_copy_QK = make_tiled_copy_C(SmemCopyAtomQK{}, tiled_mma);
-    auto smem_thr_copy_QK = smem_tiled_copy_QK.get_thread_slice(tidx);
-    Tensor tsQK = smem_thr_copy_QK.partition_D(sQK);
-
-    const int n_blocks = is_split_kv ? 1 : cute::ceil_div(cute::ceil_div(seq_len_qk, kMobaBlockSize), kBlockN);
-
-    #pragma unroll
-    for (int n_block = 0; n_block < n_blocks; ++n_block) {
-        clear(acc_s);
-        cp_async_wait<0>();
-        __syncthreads();
-        if (n_block == 0) {
-            gemm(acc_s, tSrQ, tSrK, tSsQ, tSsK, tiled_mma, smem_thr_copy_Q, smem_thr_copy_K, smem_tiled_copy_Q, smem_tiled_copy_K);
-        } else {
-            gemm<true>(acc_s, tSrQ, tSrK, tSsQ, tSsK, tiled_mma, smem_thr_copy_Q, smem_thr_copy_K, smem_tiled_copy_Q, smem_tiled_copy_K);
-        }
-        if constexpr (!is_split_kv) {
-            if (n_block < n_blocks - 1) {
-                __syncthreads();
-                tKgK.data() = tKgK.data() + kBlockN * k_head_stride;
-                copy(gmem_tiled_copy, tKgK, tKsK, tKcK);
-                cute::cp_async_fence();
-            }
-        }
-
-        Tensor rS = convert_type<input_type>(acc_s);
-        Tensor trQK = smem_thr_copy_QK.retile_S(rS);
-        cute::copy(smem_tiled_copy_QK, trQK, tsQK);
-
-        __syncthreads();
-        if (remain_q_seq >= kBlockM) {
-            copy(gmem_tiled_copy_qk, tQKsQK, tQKgQK, tQKcQK);
-        } else {
-            copy<false>(gmem_tiled_copy_qk, tQKsQK, tQKgQK, tQKcQK, remain_q_seq);
-        }
-        if constexpr (!is_split_kv) {
-            __syncthreads();
-            tQKgQK.data() = tQKgQK.data() + kBlockN;
-        }
+    if constexpr (!is_split_kv) {
+      __syncthreads();
+      tQKgQK.data() = tQKgQK.data() + kBlockN;
     }
+  }
 }
 
-template <typename input_type, int kBlockM, int kBlockN, int kMobaBlockSize, int kMaxN, bool is_split_kv>
-void qk_gemm(
-        const input_type *q_input,
-        const input_type *k_gate_mean,
-        input_type *qk_gate_weight,
-        const int *seq_len_encoder,
-        const int *seq_len_decoder,
-        const int *cu_seq_q,
-        const int *cu_seq_k,
-        const int use_moba_seq_limit,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int bsz,
-        cudaStream_t stream) {
+template <typename input_type,
+          int kBlockM,
+          int kBlockN,
+          int kMobaBlockSize,
+          int kMaxN,
+          bool is_split_kv>
+void qk_gemm(const input_type *q_input,
+             const input_type *k_gate_mean,
+             input_type *qk_gate_weight,
+             const int *seq_len_encoder,
+             const int *seq_len_decoder,
+             const int *cu_seq_q,
+             const int *cu_seq_k,
+             const int use_moba_seq_limit,
+             const int max_seq_q,
+             const int max_seq_k,
+             const int head_num,
+             const int kv_head_num,
+             const int bsz,
+             cudaStream_t stream) {
+  const int gqa_group_size = head_num / kv_head_num;
 
-    const int gqa_group_size = head_num / kv_head_num;
+  dim3 grid_dims;
+  const int num_m_block = (max_seq_q + kBlockM - 1) / kBlockM;
+  const int num_n_block =
+      ((max_seq_k + kMobaBlockSize - 1) / kMobaBlockSize + kBlockN - 1) /
+      kBlockN;
 
-    dim3 grid_dims;
-    const int num_m_block = (max_seq_q + kBlockM - 1) / kBlockM;
-    const int num_n_block = ((max_seq_k + kMobaBlockSize - 1) / kMobaBlockSize + kBlockN - 1) / kBlockN;
+  if (is_split_kv) {
+    grid_dims.x = num_n_block;
+    grid_dims.z = kv_head_num;
+  } else {
+    grid_dims.x = num_m_block;
+    grid_dims.z = head_num;
+  }
+  grid_dims.y = bsz;
 
-    if (is_split_kv) {
-        grid_dims.x = num_n_block;
-        grid_dims.z = kv_head_num;
-    } else {
-        grid_dims.x = num_m_block;
-        grid_dims.z = head_num;
-    }
-    grid_dims.y = bsz;
+  constexpr int kHeadDim = 128;
+  constexpr int smemq = kBlockM * kHeadDim * sizeof(input_type);
+  constexpr int smemk = kBlockN * kHeadDim * sizeof(input_type);
+  constexpr int smemqk = kBlockM * kBlockN * sizeof(input_type);
+  const int smem_size = smemk + max(smemq, smemqk);
 
-    constexpr int kHeadDim = 128;
-    constexpr int smemq = kBlockM * kHeadDim * sizeof(input_type);
-    constexpr int smemk = kBlockN * kHeadDim * sizeof(input_type);
-    constexpr int smemqk = kBlockM * kBlockN * sizeof(input_type);
-    const int smem_size = smemk + max(smemq, smemqk);
+  auto kernel = &qk_gemm_kernel<input_type,
+                                kBlockM,
+                                kBlockN,
+                                kMobaBlockSize,
+                                kMaxN,
+                                kHeadDim,
+                                is_split_kv>;
 
-    auto kernel = &qk_gemm_kernel<input_type, kBlockM, kBlockN, kMobaBlockSize, kMaxN, kHeadDim, is_split_kv>;
+  if (smem_size >= 48 * 1024) {
+    cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  }
 
-    if (smem_size >= 48 * 1024) {
-       cudaFuncSetAttribute(
-            kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-    }
+  kernel<<<grid_dims, 128, smem_size, stream>>>(q_input,
+                                                k_gate_mean,
+                                                qk_gate_weight,
+                                                seq_len_encoder,
+                                                seq_len_decoder,
+                                                cu_seq_q,
+                                                cu_seq_k,
+                                                use_moba_seq_limit,
+                                                max_seq_q,
+                                                max_seq_k,
+                                                head_num,
+                                                kv_head_num,
+                                                gqa_group_size);
+}
 
-    kernel<<<grid_dims, 128, smem_size, stream>>>(
-        q_input,
-        k_gate_mean,
-        qk_gate_weight,
-        seq_len_encoder,
-        seq_len_decoder,
-        cu_seq_q,
-        cu_seq_k,
+template <typename T>
+std::vector<paddle::Tensor> DispatchMobaQKGemm(
+    const paddle::Tensor &q_input,
+    const paddle::Tensor &k_block_means,
+    const paddle::Tensor &seq_len_encoder,
+    const paddle::Tensor &seq_len_decoder,
+    const paddle::Tensor &cu_seq_q,
+    const paddle::Tensor &cu_seq_k,
+    const int max_seq_q,
+    const int max_seq_k,
+    const int head_num,
+    const int kv_head_num,
+    const bool is_split_kv,
+    const int use_moba_seq_limit) {
+  constexpr int kMobaBlockSize = 128;
+  constexpr int kMaxN = 1024;
+  const int batch_size = seq_len_encoder.dims()[0];
+  using cute_type = typename cuteType<T>::type;
+  if (is_split_kv) {
+    paddle::Tensor qk_gate_weight = paddle::empty(
+        {batch_size, head_num, kMaxN}, q_input.dtype(), q_input.place());
+    qk_gemm<cute_type, 16, kMobaBlockSize, kMobaBlockSize, kMaxN, true>(
+        reinterpret_cast<const cute_type *>(q_input.data<T>()),
+        reinterpret_cast<const cute_type *>(k_block_means.data<T>()),
+        reinterpret_cast<cute_type *>(qk_gate_weight.data<T>()),
+        seq_len_encoder.data<int>(),
+        seq_len_decoder.data<int>(),
+        cu_seq_q.data<int>(),
+        cu_seq_k.data<int>(),
         use_moba_seq_limit,
         max_seq_q,
         max_seq_k,
         head_num,
         kv_head_num,
-        gqa_group_size);
+        batch_size,
+        q_input.stream());
+    return {qk_gate_weight};
+  } else {
+    constexpr int kBlockM = 128;
+    constexpr int kBlockN = 128;
+    const int token_num = q_input.dims()[0];
+    paddle::Tensor qk_gate_weight = paddle::empty(
+        {token_num, head_num, kMaxN}, q_input.dtype(), q_input.place());
+    qk_gemm<cute_type, kBlockM, kBlockN, kMobaBlockSize, kMaxN, false>(
+        reinterpret_cast<cute_type *>(const_cast<T *>(q_input.data<T>())),
+        reinterpret_cast<cute_type *>(const_cast<T *>(k_block_means.data<T>())),
+        reinterpret_cast<cute_type *>(qk_gate_weight.data<T>()),
+        seq_len_encoder.data<int>(),
+        seq_len_decoder.data<int>(),
+        cu_seq_q.data<int>(),
+        cu_seq_k.data<int>(),
+        use_moba_seq_limit,
+        max_seq_q,
+        max_seq_k,
+        head_num,
+        kv_head_num,
+        batch_size,
+        q_input.stream());
+    return {qk_gate_weight};
+  }
 }
 
-
-template <typename T>
-std::vector<paddle::Tensor> DispatchMobaQKGemm(
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& k_block_means,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const bool is_split_kv,
-        const int use_moba_seq_limit) {
-
-    constexpr int kMobaBlockSize = 128;
-    constexpr int kMaxN = 1024;
-    const int batch_size = seq_len_encoder.dims()[0];
-    using cute_type = typename cuteType<T>::type;
-    if (is_split_kv) {
-        paddle::Tensor qk_gate_weight = paddle::empty({batch_size, head_num, kMaxN}, q_input.dtype(), q_input.place());
-        qk_gemm<cute_type, 16, kMobaBlockSize, kMobaBlockSize, kMaxN, true>(
-            reinterpret_cast<const cute_type*>(q_input.data<T>()),
-            reinterpret_cast<const cute_type*>(k_block_means.data<T>()),
-            reinterpret_cast<cute_type*>(qk_gate_weight.data<T>()),
-            seq_len_encoder.data<int>(),
-            seq_len_decoder.data<int>(),
-            cu_seq_q.data<int>(),
-            cu_seq_k.data<int>(),
-            use_moba_seq_limit,
-            max_seq_q,
-            max_seq_k,
-            head_num,
-            kv_head_num,
-            batch_size,
-            q_input.stream()
-        );
-        return {qk_gate_weight};
-    } else {
-        constexpr int kBlockM = 128;
-        constexpr int kBlockN = 128;
-        const int token_num = q_input.dims()[0];
-        paddle::Tensor qk_gate_weight = paddle::empty({token_num, head_num, kMaxN}, q_input.dtype(), q_input.place());
-        qk_gemm<cute_type, kBlockM, kBlockN, kMobaBlockSize, kMaxN, false>(
-            reinterpret_cast<cute_type *>(const_cast<T*>(q_input.data<T>())),
-            reinterpret_cast<cute_type *>(const_cast<T*>(k_block_means.data<T>())),
-            reinterpret_cast<cute_type *>(qk_gate_weight.data<T>()),
-            seq_len_encoder.data<int>(),
-            seq_len_decoder.data<int>(),
-            cu_seq_q.data<int>(),
-            cu_seq_k.data<int>(),
-            use_moba_seq_limit,
-            max_seq_q,
-            max_seq_k,
-            head_num,
-            kv_head_num,
-            batch_size,
-            q_input.stream());
-        return {qk_gate_weight};
-    }
-}
-
-std::vector<paddle::Tensor> MobaQKGemm(
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& k_block_means,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const bool is_split_kv,
-        const int use_moba_seq_limit) {
-
-    if (q_input.dtype() == paddle::DataType::FLOAT16) {
-        return std::move(
-            DispatchMobaQKGemm<phi::dtype::float16>(
-                q_input,
-                k_block_means,
-                seq_len_encoder,
-                seq_len_decoder,
-                cu_seq_q,
-                cu_seq_k,
-                max_seq_q,
-                max_seq_k,
-                head_num,
-                kv_head_num,
-                is_split_kv,
-                use_moba_seq_limit
-            )
-        );
-    } else if (q_input.dtype() == paddle::DataType::BFLOAT16) {
-        return std::move(
-            DispatchMobaQKGemm<phi::dtype::bfloat16>(
-                q_input,
-                k_block_means,
-                seq_len_encoder,
-                seq_len_decoder,
-                cu_seq_q,
-                cu_seq_k,
-                max_seq_q,
-                max_seq_k,
-                head_num,
-                kv_head_num,
-                is_split_kv,
-                use_moba_seq_limit
-            )
-        );
-    }
+std::vector<paddle::Tensor> MobaQKGemm(const paddle::Tensor &q_input,
+                                       const paddle::Tensor &k_block_means,
+                                       const paddle::Tensor &seq_len_encoder,
+                                       const paddle::Tensor &seq_len_decoder,
+                                       const paddle::Tensor &cu_seq_q,
+                                       const paddle::Tensor &cu_seq_k,
+                                       const int max_seq_q,
+                                       const int max_seq_k,
+                                       const int head_num,
+                                       const int kv_head_num,
+                                       const bool is_split_kv,
+                                       const int use_moba_seq_limit) {
+  if (q_input.dtype() == paddle::DataType::FLOAT16) {
+    return std::move(
+        DispatchMobaQKGemm<phi::dtype::float16>(q_input,
+                                                k_block_means,
+                                                seq_len_encoder,
+                                                seq_len_decoder,
+                                                cu_seq_q,
+                                                cu_seq_k,
+                                                max_seq_q,
+                                                max_seq_k,
+                                                head_num,
+                                                kv_head_num,
+                                                is_split_kv,
+                                                use_moba_seq_limit));
+  } else if (q_input.dtype() == paddle::DataType::BFLOAT16) {
+    return std::move(
+        DispatchMobaQKGemm<phi::dtype::bfloat16>(q_input,
+                                                 k_block_means,
+                                                 seq_len_encoder,
+                                                 seq_len_decoder,
+                                                 cu_seq_q,
+                                                 cu_seq_k,
+                                                 max_seq_q,
+                                                 max_seq_k,
+                                                 head_num,
+                                                 kv_head_num,
+                                                 is_split_kv,
+                                                 use_moba_seq_limit));
+  }
 }
 
 PD_BUILD_STATIC_OP(moba_qk_gemm)
-    .Inputs({
-        "q_input",
-        "k_block_means",
-        "seq_len_encoder",
-        "seq_len_decoder",
-        "cu_seq_q",
-        "cu_seq_k"})
-    .Attrs({
-        "max_seq_q: int",
-        "max_seq_k: int",
-        "head_num: int",
-        "kv_head_num: int",
-        "is_split_kv: bool",
-        "use_moba_seq_limit: int"})
+    .Inputs({"q_input",
+             "k_block_means",
+             "seq_len_encoder",
+             "seq_len_decoder",
+             "cu_seq_q",
+             "cu_seq_k"})
+    .Attrs({"max_seq_q: int",
+            "max_seq_k: int",
+            "head_num: int",
+            "kv_head_num: int",
+            "is_split_kv: bool",
+            "use_moba_seq_limit: int"})
     .Outputs({"qk_gate_weight"})
     .SetKernelFn(PD_KERNEL(MobaQKGemm));
diff --git a/custom_ops/gpu_ops/moba_attn/moba_process/split_qkv_and_rope.cu b/custom_ops/gpu_ops/moba_attn/moba_process/split_qkv_and_rope.cu
index c1fe07b398..ff693254ac 100644
--- a/custom_ops/gpu_ops/moba_attn/moba_process/split_qkv_and_rope.cu
+++ b/custom_ops/gpu_ops/moba_attn/moba_process/split_qkv_and_rope.cu
@@ -20,352 +20,390 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-template <typename input_type, int moba_block_size, int kBlockM, int kMaxN, int tokens_per_block, bool need_k_mean>
-__global__ void fused_block_mean_and_rope_kernel(
-        const input_type *qkv_input,
-        const input_type *qkv_bias,
-        input_type *k_gate_mean,
-        input_type *q_input,
-        input_type *k_input,
-        input_type *v_input,
-        const float *rope_sin_cos,
-        const int *seq_len_encoder,
-        const int *seq_len_decoder,
-        const int *cu_seq_q,
-        const int *cu_seq_k,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int max_input_length) {
+template <typename input_type,
+          int moba_block_size,
+          int kBlockM,
+          int kMaxN,
+          int tokens_per_block,
+          bool need_k_mean>
+__global__ void fused_block_mean_and_rope_kernel(const input_type *qkv_input,
+                                                 const input_type *qkv_bias,
+                                                 input_type *k_gate_mean,
+                                                 input_type *q_input,
+                                                 input_type *k_input,
+                                                 input_type *v_input,
+                                                 const float *rope_sin_cos,
+                                                 const int *seq_len_encoder,
+                                                 const int *seq_len_decoder,
+                                                 const int *cu_seq_q,
+                                                 const int *cu_seq_k,
+                                                 const int max_seq_q,
+                                                 const int max_seq_k,
+                                                 const int head_num,
+                                                 const int kv_head_num,
+                                                 const int max_input_length) {
+  constexpr int kPackSize = 16 / sizeof(input_type);
+  constexpr int kHeadDim = 128;
 
-    constexpr int kPackSize = 16 / sizeof(input_type);
-    constexpr int kHeadDim = 128;
+  using src_type = Vec<input_type, kPackSize>;
 
-    using src_type = Vec<input_type, kPackSize>;
+  using rope_type = Vec<float, kPackSize / 2>;
+  using pack_half =
+      std::conditional_t<std::is_same<input_type, cutlass::half_t>::value,
+                         __half2,
+                         nv_bfloat162>;
 
-    using rope_type = Vec<float, kPackSize / 2>;
-    using pack_half = std::conditional_t<std::is_same<input_type, cutlass::half_t>::value, __half2, nv_bfloat162>;
+  __align__(16) __shared__ input_type local_sum_mem[128 / 32 * kHeadDim];
 
-    __align__(16) __shared__ input_type local_sum_mem[128 / 32 * kHeadDim];
+  const int bidb = blockIdx.x;
+  const int bidh = blockIdx.y;
+  const int bidt_q = blockIdx.z * tokens_per_block;
+  const int bidt_v = blockIdx.z * tokens_per_block;
+  const int bidt_k = need_k_mean ? blockIdx.z * moba_block_size
+                                 : blockIdx.z * tokens_per_block;
+  const int tidx = threadIdx.x;
+  const int lane_id = tidx % 32;
+  const int warp_id = tidx / 32;
+  const int seq_len = seq_len_encoder[bidb];
+  const int seq_len_start = seq_len_decoder[bidb];
 
-    const int bidb = blockIdx.x;
-    const int bidh = blockIdx.y;
-    const int bidt_q = blockIdx.z * tokens_per_block;
-    const int bidt_v = blockIdx.z * tokens_per_block;
-    const int bidt_k = need_k_mean ? blockIdx.z * moba_block_size : blockIdx.z * tokens_per_block;
-    const int tidx = threadIdx.x;
-    const int lane_id = tidx % 32;
-    const int warp_id = tidx / 32;
-    const int seq_len = seq_len_encoder[bidb];
-    const int seq_len_start = seq_len_decoder[bidb];
+  if (seq_len == 0) {
+    return;
+  }
 
-    if (seq_len == 0) {
-        return;
+  const int all_head_num = head_num + 2 * kv_head_num;
+  const int hidden = all_head_num * kHeadDim;
+
+  const int row_idx = tidx / (kHeadDim / kPackSize);
+  const int col_idx = tidx % (kHeadDim / kPackSize);
+
+  const int bias_idx = bidh * kHeadDim + col_idx * kPackSize;
+
+  src_type src, src_bias;
+  rope_type sin, cos;
+
+  const bool need_add_bias = qkv_bias != nullptr;
+
+  if (need_add_bias) {
+    src_bias.load_from(qkv_bias + bias_idx);
+  }
+
+  if (bidh < head_num) {
+    const int cur_token = bidt_q + row_idx;
+    const float *cos_rope = rope_sin_cos +
+                            (cur_token + seq_len_start) * (kHeadDim / 2) +
+                            col_idx * (kPackSize / 2);
+    const float *sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
+
+    if (cur_token < seq_len) {
+      src.load_from(qkv_input + cu_seq_q[bidb] * hidden + bias_idx +
+                    cur_token * hidden);
+
+      if (need_add_bias) {
+        src.add(src_bias);
+      }
+
+      sin.load_from(sin_rope);
+      cos.load_from(cos_rope);
+      apply_rotary_embedding<input_type, kPackSize>(src, cos, sin);
+
+      src.store_to(q_input +
+                   (cu_seq_q[bidb] + cur_token) * head_num * kHeadDim +
+                   bias_idx);
     }
+  } else if (bidh < head_num + kv_head_num) {
+    if constexpr (!need_k_mean) {
+      const int cur_token = bidt_k + row_idx;
+      const float *cos_rope = rope_sin_cos +
+                              (cur_token + seq_len_start) * (kHeadDim / 2) +
+                              col_idx * (kPackSize / 2);
+      const float *sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
 
-    const int all_head_num = head_num + 2 * kv_head_num;
-    const int hidden = all_head_num * kHeadDim;
+      if (cur_token < seq_len) {
+        src.load_from(qkv_input + cu_seq_q[bidb] * hidden + bias_idx +
+                      cur_token * hidden);
 
-    const int row_idx = tidx / (kHeadDim / kPackSize);
-    const int col_idx = tidx % (kHeadDim / kPackSize);
-
-    const int bias_idx = bidh * kHeadDim + col_idx * kPackSize;
-
-    src_type src, src_bias;
-    rope_type sin, cos;
-
-    const bool need_add_bias = qkv_bias != nullptr;
-
-    if (need_add_bias) {
-        src_bias.load_from(qkv_bias + bias_idx);
-    }
-
-    if (bidh < head_num) {
-        const int cur_token = bidt_q + row_idx;
-        const float * cos_rope = rope_sin_cos + (cur_token + seq_len_start) * (kHeadDim / 2) + col_idx * (kPackSize / 2);
-        const float * sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
-
-        if (cur_token < seq_len) {
-            src.load_from(qkv_input + cu_seq_q[bidb] * hidden + bias_idx + cur_token * hidden);
-
-            if (need_add_bias) {
-                src.add(src_bias);
-            }
-
-            sin.load_from(sin_rope);
-            cos.load_from(cos_rope);
-            apply_rotary_embedding<input_type, kPackSize>(src, cos, sin);
-
-            src.store_to(q_input + (cu_seq_q[bidb] + cur_token) * head_num * kHeadDim + bias_idx);
+        if (need_add_bias) {
+          src.add(src_bias);
         }
-    } else if (bidh < head_num + kv_head_num) {
-        if constexpr (!need_k_mean) {
-            const int cur_token = bidt_k + row_idx;
-            const float * cos_rope = rope_sin_cos + (cur_token + seq_len_start) * (kHeadDim / 2) + col_idx * (kPackSize / 2);
-            const float * sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
 
-            if (cur_token < seq_len) {
-                src.load_from(qkv_input + cu_seq_q[bidb] * hidden + bias_idx + cur_token * hidden);
+        sin.load_from(sin_rope);
+        cos.load_from(cos_rope);
+        apply_rotary_embedding<input_type, kPackSize>(src, cos, sin);
 
-                if (need_add_bias) {
-                    src.add(src_bias);
-                }
-
-                sin.load_from(sin_rope);
-                cos.load_from(cos_rope);
-                apply_rotary_embedding<input_type, kPackSize>(src, cos, sin);
-
-                src.store_to(k_input + (cu_seq_k[bidb] + cur_token) * head_num * kHeadDim + bias_idx- head_num * kHeadDim);
-            }
-        } else {
-            if (bidt_k >= seq_len) {
-                return;
-            }
-
-            src_type local_sum;
-            local_sum.set_zero();
-
-            const input_type* qkv = qkv_input + cu_seq_q[bidb] * hidden + bias_idx;
-
-            for (int i = 0; i < moba_block_size; i += tokens_per_block) {
-                const int cur_token = bidt_k + i + row_idx;
-                if (cur_token < seq_len) {
-                    src.load_from(qkv + cur_token * hidden);
-
-                    if (need_add_bias) {
-                        src.add(src_bias);
-                    }
-                    const float * cos_rope = rope_sin_cos + (cur_token + seq_len_start) * (kHeadDim / 2) + col_idx * (kPackSize / 2);
-                    const float * sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
-                    sin.load_from(sin_rope);
-                    cos.load_from(cos_rope);
-
-                    apply_rotary_embedding<input_type, kPackSize>(src, cos, sin);
-
-                    src.store_to(k_input + (cu_seq_k[bidb] + cur_token) * kv_head_num * kHeadDim + bias_idx - head_num * kHeadDim);
-
-                    local_sum.add(src);
-                }
-            }
-
-            src_type neighbor;
-
-            #pragma unroll
-            for (int i = 0; i < kPackSize; i+=2) {
-                *reinterpret_cast<int32_t*>(neighbor.data.elt + i) = __shfl_down_sync(0xffffffff, *reinterpret_cast<int32_t*>(local_sum.data.elt + i), 16);
-            }
-
-            local_sum.add(neighbor);
-
-            if (lane_id < 16) {
-                local_sum.store_to(local_sum_mem + warp_id * kHeadDim + lane_id * kPackSize);
-            }
-
-            __syncthreads();
-
-            pack_half * local_sum_mem_half = reinterpret_cast<pack_half*>(local_sum_mem);
-
-            pack_half local_sum_half = local_sum_mem_half[tidx];
-
-
-            if (tidx < kHeadDim / 2) {
-
-                #pragma unroll
-                for (int i = 1; i < 4; i++) {
-                    local_sum_half += local_sum_mem_half[tidx + i * (kHeadDim / 2)];
-                }
-
-                float inv_tokens_sum = fdividef(1.0f, min(seq_len - bidt_k, moba_block_size));
-
-                local_sum_half *= float_2_half2<input_type>(inv_tokens_sum);
-
-                const int store_mean_idx = ((bidb * kMaxN + blockIdx.z + seq_len_start / moba_block_size) * kv_head_num * kHeadDim + (bidh - head_num) * kHeadDim) / 2 + tidx;
-
-                reinterpret_cast<pack_half*>(k_gate_mean)[store_mean_idx] = local_sum_half;
-            }
-        }
+        src.store_to(k_input +
+                     (cu_seq_k[bidb] + cur_token) * head_num * kHeadDim +
+                     bias_idx - head_num * kHeadDim);
+      }
     } else {
-        const int cur_token = bidt_v + row_idx;
+      if (bidt_k >= seq_len) {
+        return;
+      }
 
+      src_type local_sum;
+      local_sum.set_zero();
+
+      const input_type *qkv = qkv_input + cu_seq_q[bidb] * hidden + bias_idx;
+
+      for (int i = 0; i < moba_block_size; i += tokens_per_block) {
+        const int cur_token = bidt_k + i + row_idx;
         if (cur_token < seq_len) {
-            src.load_from(qkv_input + cu_seq_q[bidb] * hidden + bias_idx + cur_token * hidden);
-            if (need_add_bias) {
-                src.add(src_bias);
-            }
+          src.load_from(qkv + cur_token * hidden);
 
-            src.store_to(v_input + (cu_seq_k[bidb] + cur_token) * kv_head_num * kHeadDim + bias_idx - (head_num + kv_head_num) * kHeadDim);
+          if (need_add_bias) {
+            src.add(src_bias);
+          }
+          const float *cos_rope = rope_sin_cos +
+                                  (cur_token + seq_len_start) * (kHeadDim / 2) +
+                                  col_idx * (kPackSize / 2);
+          const float *sin_rope = cos_rope + max_input_length * (kHeadDim / 2);
+          sin.load_from(sin_rope);
+          cos.load_from(cos_rope);
+
+          apply_rotary_embedding<input_type, kPackSize>(src, cos, sin);
+
+          src.store_to(k_input +
+                       (cu_seq_k[bidb] + cur_token) * kv_head_num * kHeadDim +
+                       bias_idx - head_num * kHeadDim);
+
+          local_sum.add(src);
         }
+      }
+
+      src_type neighbor;
+
+#pragma unroll
+      for (int i = 0; i < kPackSize; i += 2) {
+        *reinterpret_cast<int32_t *>(neighbor.data.elt + i) = __shfl_down_sync(
+            0xffffffff,
+            *reinterpret_cast<int32_t *>(local_sum.data.elt + i),
+            16);
+      }
+
+      local_sum.add(neighbor);
+
+      if (lane_id < 16) {
+        local_sum.store_to(local_sum_mem + warp_id * kHeadDim +
+                           lane_id * kPackSize);
+      }
+
+      __syncthreads();
+
+      pack_half *local_sum_mem_half =
+          reinterpret_cast<pack_half *>(local_sum_mem);
+
+      pack_half local_sum_half = local_sum_mem_half[tidx];
+
+      if (tidx < kHeadDim / 2) {
+#pragma unroll
+        for (int i = 1; i < 4; i++) {
+          local_sum_half += local_sum_mem_half[tidx + i * (kHeadDim / 2)];
+        }
+
+        float inv_tokens_sum =
+            fdividef(1.0f, min(seq_len - bidt_k, moba_block_size));
+
+        local_sum_half *= float_2_half2<input_type>(inv_tokens_sum);
+
+        const int store_mean_idx =
+            ((bidb * kMaxN + blockIdx.z + seq_len_start / moba_block_size) *
+                 kv_head_num * kHeadDim +
+             (bidh - head_num) * kHeadDim) /
+                2 +
+            tidx;
+
+        reinterpret_cast<pack_half *>(k_gate_mean)[store_mean_idx] =
+            local_sum_half;
+      }
     }
+  } else {
+    const int cur_token = bidt_v + row_idx;
+
+    if (cur_token < seq_len) {
+      src.load_from(qkv_input + cu_seq_q[bidb] * hidden + bias_idx +
+                    cur_token * hidden);
+      if (need_add_bias) {
+        src.add(src_bias);
+      }
+
+      src.store_to(v_input +
+                   (cu_seq_k[bidb] + cur_token) * kv_head_num * kHeadDim +
+                   bias_idx - (head_num + kv_head_num) * kHeadDim);
+    }
+  }
 }
 
 template <typename input_type, int moba_block_size, int kBlockM, int kMaxN>
-void fused_block_mean_and_rope(
-        const input_type *qkv_input,
-        const input_type *qkv_bias,
-        input_type *k_gate_mean,
-        input_type *q_input,
-        input_type *k_input,
-        input_type *v_input,
-        const float *rope_sin_cos,
-        const int *seq_len_encoder,
-        const int *seq_len_decoder,
-        const int *cu_seq_q,
-        const int *cu_seq_k,
-        const int max_seq_q,
-        const int max_seq_k,
-        const int head_num,
-        const int kv_head_num,
-        const int bsz,
-        const int max_input_length,
-        cudaStream_t stream) {
+void fused_block_mean_and_rope(const input_type *qkv_input,
+                               const input_type *qkv_bias,
+                               input_type *k_gate_mean,
+                               input_type *q_input,
+                               input_type *k_input,
+                               input_type *v_input,
+                               const float *rope_sin_cos,
+                               const int *seq_len_encoder,
+                               const int *seq_len_decoder,
+                               const int *cu_seq_q,
+                               const int *cu_seq_k,
+                               const int max_seq_q,
+                               const int max_seq_k,
+                               const int head_num,
+                               const int kv_head_num,
+                               const int bsz,
+                               const int max_input_length,
+                               cudaStream_t stream) {
+  static_assert(moba_block_size >= 64, "moba_block_size must be at least 64");
+  constexpr int kPackSize = 16 / sizeof(input_type);
+  constexpr int kHeadDim = 128;
+  constexpr int kThreads = 128;
+  constexpr int tokens_per_block = kThreads / (kHeadDim / kPackSize);
+  dim3 grid_dims;
+  grid_dims.x = bsz;
+  grid_dims.y = head_num + 2 * kv_head_num;
+  grid_dims.z = (max_seq_q + tokens_per_block - 1) / tokens_per_block;
 
-    static_assert(moba_block_size >= 64, "moba_block_size must be at least 64");
-    constexpr int kPackSize = 16 / sizeof(input_type);
-    constexpr int kHeadDim = 128;
-    constexpr int kThreads = 128;
-    constexpr int tokens_per_block = kThreads / (kHeadDim / kPackSize);
-    dim3 grid_dims;
-    grid_dims.x = bsz;
-    grid_dims.y = head_num + 2 * kv_head_num;
-    grid_dims.z = (max_seq_q + tokens_per_block - 1) / tokens_per_block;
-
-    if (k_gate_mean != nullptr) {
-        fused_block_mean_and_rope_kernel<input_type, moba_block_size, kBlockM, kMaxN, tokens_per_block, true>
-        <<<grid_dims, kThreads, 0, stream>>>(
-            qkv_input,
-            qkv_bias,
-            k_gate_mean,
-            q_input,
-            k_input,
-            v_input,
-            rope_sin_cos,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            cu_seq_k,
-            max_seq_q,
-            max_seq_k,
-            head_num,
-            kv_head_num,
-            max_input_length);
-    } else {
-        fused_block_mean_and_rope_kernel<input_type, moba_block_size, kBlockM, kMaxN, tokens_per_block, false>
-        <<<grid_dims, kThreads, 0, stream>>>(
-            qkv_input,
-            qkv_bias,
-            k_gate_mean,
-            q_input,
-            k_input,
-            v_input,
-            rope_sin_cos,
-            seq_len_encoder,
-            seq_len_decoder,
-            cu_seq_q,
-            cu_seq_k,
-            max_seq_q,
-            max_seq_k,
-            head_num,
-            kv_head_num,
-            max_input_length);
-    }
+  if (k_gate_mean != nullptr) {
+    fused_block_mean_and_rope_kernel<input_type,
+                                     moba_block_size,
+                                     kBlockM,
+                                     kMaxN,
+                                     tokens_per_block,
+                                     true>
+        <<<grid_dims, kThreads, 0, stream>>>(qkv_input,
+                                             qkv_bias,
+                                             k_gate_mean,
+                                             q_input,
+                                             k_input,
+                                             v_input,
+                                             rope_sin_cos,
+                                             seq_len_encoder,
+                                             seq_len_decoder,
+                                             cu_seq_q,
+                                             cu_seq_k,
+                                             max_seq_q,
+                                             max_seq_k,
+                                             head_num,
+                                             kv_head_num,
+                                             max_input_length);
+  } else {
+    fused_block_mean_and_rope_kernel<input_type,
+                                     moba_block_size,
+                                     kBlockM,
+                                     kMaxN,
+                                     tokens_per_block,
+                                     false>
+        <<<grid_dims, kThreads, 0, stream>>>(qkv_input,
+                                             qkv_bias,
+                                             k_gate_mean,
+                                             q_input,
+                                             k_input,
+                                             v_input,
+                                             rope_sin_cos,
+                                             seq_len_encoder,
+                                             seq_len_decoder,
+                                             cu_seq_q,
+                                             cu_seq_k,
+                                             max_seq_q,
+                                             max_seq_k,
+                                             head_num,
+                                             kv_head_num,
+                                             max_input_length);
+  }
 }
 
-void FusedBlockMeanAndRope(
-        const paddle::Tensor& qkv_out,
-        const paddle::Tensor& k_block_means,
-        const paddle::Tensor& q_input,
-        const paddle::Tensor& k_input,
-        const paddle::Tensor& v_input,
-        const paddle::Tensor& rotary_embs,
-        const paddle::Tensor& seq_len_encoder,
-        const paddle::Tensor& seq_len_decoder,
-        const paddle::Tensor& cu_seq_q,
-        const paddle::Tensor& cu_seq_k,
-        const paddle::optional<paddle::Tensor>& qkv_bias,
-        const int head_num,
-        const int kv_head_num,
-        const int head_dim,
-        const int max_input_length,
-        const int max_seq_q,
-        const int max_seq_k,
-        const std::string &cache_quant_type_str) {
+void FusedBlockMeanAndRope(const paddle::Tensor &qkv_out,
+                           const paddle::Tensor &k_block_means,
+                           const paddle::Tensor &q_input,
+                           const paddle::Tensor &k_input,
+                           const paddle::Tensor &v_input,
+                           const paddle::Tensor &rotary_embs,
+                           const paddle::Tensor &seq_len_encoder,
+                           const paddle::Tensor &seq_len_decoder,
+                           const paddle::Tensor &cu_seq_q,
+                           const paddle::Tensor &cu_seq_k,
+                           const paddle::optional<paddle::Tensor> &qkv_bias,
+                           const int head_num,
+                           const int kv_head_num,
+                           const int head_dim,
+                           const int max_input_length,
+                           const int max_seq_q,
+                           const int max_seq_k,
+                           const std::string &cache_quant_type_str) {
+  constexpr int kBlockM = 128;
+  constexpr int kBlockN = 128;
+  constexpr int kMobaBlockSize = 128;
+  constexpr int kMaxN = 1024;
 
-    constexpr int kBlockM = 128;
-    constexpr int kBlockN = 128;
-    constexpr int kMobaBlockSize = 128;
-    constexpr int kMaxN = 1024;
-
-    if (k_input.dtype() == paddle::DataType::FLOAT16) {
-        using T = phi::dtype::float16;
-        using cute_type = typename cuteType<T>::type;
-        fused_block_mean_and_rope<cute_type, kMobaBlockSize, kBlockM, kMaxN>(
-            reinterpret_cast<cute_type *>(const_cast<T*>(qkv_out.data<T>())),
-            qkv_bias ? reinterpret_cast<cute_type *>(const_cast<T*>(qkv_bias.get().data<T>())) : nullptr,
-            reinterpret_cast<cute_type *>(const_cast<T*>(k_block_means.data<T>())),
-            reinterpret_cast<cute_type*>(const_cast<T*>(q_input.data<T>())),
-            reinterpret_cast<cute_type*>(const_cast<T*>(k_input.data<T>())),
-            reinterpret_cast<cute_type*>(const_cast<T*>(v_input.data<T>())),
-            rotary_embs.data<float>(),
-            seq_len_encoder.data<int>(),
-            seq_len_decoder.data<int>(),
-            cu_seq_q.data<int>(),
-            cu_seq_k.data<int>(),
-            max_seq_q,
-            max_seq_k,
-            head_num,
-            kv_head_num,
-            seq_len_encoder.dims()[0],
-            max_input_length,
-            qkv_out.stream());
-    } else if (k_input.dtype() == paddle::DataType::BFLOAT16) {
-        using T = phi::dtype::bfloat16;
-        using cute_type = typename cuteType<T>::type;
-        fused_block_mean_and_rope<cute_type, kMobaBlockSize, kBlockM, kMaxN>(
-            reinterpret_cast<cute_type *>(const_cast<T*>(qkv_out.data<T>())),
-            qkv_bias ? reinterpret_cast<cute_type *>(const_cast<T*>(qkv_bias.get().data<T>())) : nullptr,
-            reinterpret_cast<cute_type *>(const_cast<T*>(k_block_means.data<T>())),
-            reinterpret_cast<cute_type*>(const_cast<T*>(q_input.data<T>())),
-            reinterpret_cast<cute_type*>(const_cast<T*>(k_input.data<T>())),
-            reinterpret_cast<cute_type*>(const_cast<T*>(v_input.data<T>())),
-            rotary_embs.data<float>(),
-            seq_len_encoder.data<int>(),
-            seq_len_decoder.data<int>(),
-            cu_seq_q.data<int>(),
-            cu_seq_k.data<int>(),
-            max_seq_q,
-            max_seq_k,
-            head_num,
-            kv_head_num,
-            seq_len_encoder.dims()[0],
-            max_input_length,
-            qkv_out.stream());
-    }
+  if (k_input.dtype() == paddle::DataType::FLOAT16) {
+    using T = phi::dtype::float16;
+    using cute_type = typename cuteType<T>::type;
+    fused_block_mean_and_rope<cute_type, kMobaBlockSize, kBlockM, kMaxN>(
+        reinterpret_cast<cute_type *>(const_cast<T *>(qkv_out.data<T>())),
+        qkv_bias ? reinterpret_cast<cute_type *>(
+                       const_cast<T *>(qkv_bias.get().data<T>()))
+                 : nullptr,
+        reinterpret_cast<cute_type *>(const_cast<T *>(k_block_means.data<T>())),
+        reinterpret_cast<cute_type *>(const_cast<T *>(q_input.data<T>())),
+        reinterpret_cast<cute_type *>(const_cast<T *>(k_input.data<T>())),
+        reinterpret_cast<cute_type *>(const_cast<T *>(v_input.data<T>())),
+        rotary_embs.data<float>(),
+        seq_len_encoder.data<int>(),
+        seq_len_decoder.data<int>(),
+        cu_seq_q.data<int>(),
+        cu_seq_k.data<int>(),
+        max_seq_q,
+        max_seq_k,
+        head_num,
+        kv_head_num,
+        seq_len_encoder.dims()[0],
+        max_input_length,
+        qkv_out.stream());
+  } else if (k_input.dtype() == paddle::DataType::BFLOAT16) {
+    using T = phi::dtype::bfloat16;
+    using cute_type = typename cuteType<T>::type;
+    fused_block_mean_and_rope<cute_type, kMobaBlockSize, kBlockM, kMaxN>(
+        reinterpret_cast<cute_type *>(const_cast<T *>(qkv_out.data<T>())),
+        qkv_bias ? reinterpret_cast<cute_type *>(
+                       const_cast<T *>(qkv_bias.get().data<T>()))
+                 : nullptr,
+        reinterpret_cast<cute_type *>(const_cast<T *>(k_block_means.data<T>())),
+        reinterpret_cast<cute_type *>(const_cast<T *>(q_input.data<T>())),
+        reinterpret_cast<cute_type *>(const_cast<T *>(k_input.data<T>())),
+        reinterpret_cast<cute_type *>(const_cast<T *>(v_input.data<T>())),
+        rotary_embs.data<float>(),
+        seq_len_encoder.data<int>(),
+        seq_len_decoder.data<int>(),
+        cu_seq_q.data<int>(),
+        cu_seq_k.data<int>(),
+        max_seq_q,
+        max_seq_k,
+        head_num,
+        kv_head_num,
+        seq_len_encoder.dims()[0],
+        max_input_length,
+        qkv_out.stream());
+  }
 }
 
-
-
 PD_BUILD_STATIC_OP(fused_block_mean_and_rope)
-    .Inputs({
-        "qkv_out",
-        "k_block_means",
-        "q_input",
-        "k_input",
-        "v_input",
-        "rotary_embs",
-        "seq_len_encoder",
-        "seq_len_decoder",
-        "cu_seq_q",
-        "cu_seq_k",
-        paddle::Optional("qkv_bias")})
-    .Attrs({
-        "head_num: int",
-        "kv_head_num: int",
-        "head_dim: int",
-        "max_input_length: int",
-        "max_seq_q: int",
-        "max_seq_k: int",
-        "cache_quant_type_str: std::string"})
+    .Inputs({"qkv_out",
+             "k_block_means",
+             "q_input",
+             "k_input",
+             "v_input",
+             "rotary_embs",
+             "seq_len_encoder",
+             "seq_len_decoder",
+             "cu_seq_q",
+             "cu_seq_k",
+             paddle::Optional("qkv_bias")})
+    .Attrs({"head_num: int",
+            "kv_head_num: int",
+            "head_dim: int",
+            "max_input_length: int",
+            "max_seq_q: int",
+            "max_seq_k: int",
+            "cache_quant_type_str: std::string"})
     .Outputs({"q_input_out", "k_input_out", "v_input_out", "k_block_means_out"})
     .SetInplaceMap({{"q_input", "q_input_out"},
                     {"k_input", "k_input_out"},
diff --git a/custom_ops/gpu_ops/moe/fused_moe.cu b/custom_ops/gpu_ops/moe/fused_moe.cu
index a09bfa9e79..edc5843095 100644
--- a/custom_ops/gpu_ops/moe/fused_moe.cu
+++ b/custom_ops/gpu_ops/moe/fused_moe.cu
@@ -28,13 +28,13 @@ __global__ void compute_total_rows_before_expert_kernel(
     const int64_t sorted_experts_len,
     const int64_t num_experts,
     int64_t* total_rows_before_expert) {
-    // First, compute the global tid. We only need 1 thread per expert.
-    const int expert = blockIdx.x * blockDim.x + threadIdx.x;
-    if (expert >= num_experts) return;
+  // First, compute the global tid. We only need 1 thread per expert.
+  const int expert = blockIdx.x * blockDim.x + threadIdx.x;
+  if (expert >= num_experts) return;
 
-    // This should construct the last index where each expert occurs.
-    total_rows_before_expert[expert] =
-        find_total_elts_leq_target(sorted_experts, sorted_experts_len, expert);
+  // This should construct the last index where each expert occurs.
+  total_rows_before_expert[expert] =
+      find_total_elts_leq_target(sorted_experts, sorted_experts_len, expert);
 }
 
 void compute_total_rows_before_expert(int* sorted_indices,
@@ -42,11 +42,11 @@ void compute_total_rows_before_expert(int* sorted_indices,
                                       const int64_t num_experts,
                                       int64_t* total_rows_before_expert,
                                       cudaStream_t stream) {
-    const int threads = std::min(int64_t(1024), num_experts);
-    const int blocks = (num_experts + threads - 1) / threads;
+  const int threads = std::min(int64_t(1024), num_experts);
+  const int blocks = (num_experts + threads - 1) / threads;
 
-    compute_total_rows_before_expert_kernel<<<blocks, threads, 0, stream>>>(
-        sorted_indices, total_indices, num_experts, total_rows_before_expert);
+  compute_total_rows_before_expert_kernel<<<blocks, threads, 0, stream>>>(
+      sorted_indices, total_indices, num_experts, total_rows_before_expert);
 }
 
 }  // namespace phi
@@ -65,38 +65,47 @@ void FusedMoeKernel(const paddle::Tensor& input,
                     const bool group_moe,
                     const bool norm_topk_prob,
                     paddle::Tensor* output) {
-    using namespace phi;
-    typedef PDTraits<T> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+  using namespace phi;
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    auto* output_data = output->data<data_t>();
+  auto* output_data = output->data<data_t>();
 
-    auto fp16_moe_gemm_runner = MoeGemmRunner<DataType_, cutlass::WintQuantTraits<DataType_, cutlass::WintQuantMethod::kNone>>();
-    auto int8_moe_gemm_runner = MoeGemmRunner<DataType_, cutlass::WintQuantTraits<DataType_, cutlass::WintQuantMethod::kWeightOnlyInt8>>();
-    auto int4_moe_gemm_runner = MoeGemmRunner<DataType_, cutlass::WintQuantTraits<DataType_, cutlass::WintQuantMethod::kWeightOnlyInt4>>();
+  auto fp16_moe_gemm_runner = MoeGemmRunner<
+      DataType_,
+      cutlass::WintQuantTraits<DataType_, cutlass::WintQuantMethod::kNone>>();
+  auto int8_moe_gemm_runner = MoeGemmRunner<
+      DataType_,
+      cutlass::WintQuantTraits<DataType_,
+                               cutlass::WintQuantMethod::kWeightOnlyInt8>>();
+  auto int4_moe_gemm_runner = MoeGemmRunner<
+      DataType_,
+      cutlass::WintQuantTraits<DataType_,
+                               cutlass::WintQuantMethod::kWeightOnlyInt4>>();
 
-    using NvType = typename traits_::DataType;
-    auto moe_compute = MoeHelper<data_t, NvType>(quant_method,
-                                                 &fp16_moe_gemm_runner,
-                                                 &int8_moe_gemm_runner,
-                                                 &int4_moe_gemm_runner);
+  using NvType = typename traits_::DataType;
+  auto moe_compute = MoeHelper<data_t, NvType>(quant_method,
+                                               &fp16_moe_gemm_runner,
+                                               &int8_moe_gemm_runner,
+                                               &int4_moe_gemm_runner);
 
-    moe_compute.ComputeFFN(&input,
-                           &gate_weight,
-                           &up_gate_proj_weight,
-                           up_gate_proj_scale ? up_gate_proj_scale.get_ptr() : nullptr,
-                           up_gate_proj_bias ? up_gate_proj_bias.get_ptr() : nullptr,
-                           &down_proj_weight,
-                           down_proj_scale ? down_proj_scale.get_ptr() : nullptr,
-                           down_proj_bias ? down_proj_bias.get_ptr() : nullptr,
-                           nullptr,
-                           moe_topk,
-                           group_moe,
-                           norm_topk_prob,
-                           1.0,  // ComputeFFN
-                           "ffn",
-                           output);
+  moe_compute.ComputeFFN(
+      &input,
+      &gate_weight,
+      &up_gate_proj_weight,
+      up_gate_proj_scale ? up_gate_proj_scale.get_ptr() : nullptr,
+      up_gate_proj_bias ? up_gate_proj_bias.get_ptr() : nullptr,
+      &down_proj_weight,
+      down_proj_scale ? down_proj_scale.get_ptr() : nullptr,
+      down_proj_bias ? down_proj_bias.get_ptr() : nullptr,
+      nullptr,
+      moe_topk,
+      group_moe,
+      norm_topk_prob,
+      1.0,  // ComputeFFN
+      "ffn",
+      output);
 }
 
 paddle::Tensor FusedExpertMoeFunc(
@@ -112,44 +121,44 @@ paddle::Tensor FusedExpertMoeFunc(
     const int moe_topk,
     const bool norm_topk_prob,
     const bool group_moe) {
-    const auto input_type = input.dtype();
-    auto output = paddle::empty_like(input);
+  const auto input_type = input.dtype();
+  auto output = paddle::empty_like(input);
 
-    switch (input_type) {
-        case paddle::DataType::BFLOAT16:
-            FusedMoeKernel<paddle::DataType::BFLOAT16>(input,
-                                                       gate_weight,
-                                                       up_gate_proj_weight,
-                                                       up_gate_proj_scale,
-                                                       up_gate_proj_bias,
-                                                       down_proj_weight,
-                                                       down_proj_scale,
-                                                       down_proj_bias,
-                                                       quant_method,
-                                                       moe_topk,
-                                                       group_moe,
-                                                       norm_topk_prob,
-                                                       &output);
-            break;
-        case paddle::DataType::FLOAT16:
-            FusedMoeKernel<paddle::DataType::FLOAT16>(input,
-                                                      gate_weight,
-                                                      up_gate_proj_weight,
-                                                      up_gate_proj_scale,
-                                                      up_gate_proj_bias,
-                                                      down_proj_weight,
-                                                      down_proj_scale,
-                                                      down_proj_bias,
-                                                      quant_method,
-                                                      moe_topk,
-                                                      group_moe,
-                                                      norm_topk_prob,
-                                                      &output);
-            break;
-        default:
-            PD_THROW("Unsupported data type for FusedMoeKernel");
-    }
-    return output;
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      FusedMoeKernel<paddle::DataType::BFLOAT16>(input,
+                                                 gate_weight,
+                                                 up_gate_proj_weight,
+                                                 up_gate_proj_scale,
+                                                 up_gate_proj_bias,
+                                                 down_proj_weight,
+                                                 down_proj_scale,
+                                                 down_proj_bias,
+                                                 quant_method,
+                                                 moe_topk,
+                                                 group_moe,
+                                                 norm_topk_prob,
+                                                 &output);
+      break;
+    case paddle::DataType::FLOAT16:
+      FusedMoeKernel<paddle::DataType::FLOAT16>(input,
+                                                gate_weight,
+                                                up_gate_proj_weight,
+                                                up_gate_proj_scale,
+                                                up_gate_proj_bias,
+                                                down_proj_weight,
+                                                down_proj_scale,
+                                                down_proj_bias,
+                                                quant_method,
+                                                moe_topk,
+                                                group_moe,
+                                                norm_topk_prob,
+                                                &output);
+      break;
+    default:
+      PD_THROW("Unsupported data type for FusedMoeKernel");
+  }
+  return output;
 }
 
 std::vector<paddle::Tensor> FusedExpertMoe(
@@ -165,18 +174,18 @@ std::vector<paddle::Tensor> FusedExpertMoe(
     const int moe_topk,
     const bool norm_topk_prob,
     const bool group_moe) {
-    return {FusedExpertMoeFunc(input,
-                               gate_weight,
-                               up_gate_proj_weight,
-                               down_proj_weight,
-                               up_gate_proj_bias,
-                               up_gate_proj_scale,
-                               down_proj_bias,
-                               down_proj_scale,
-                               quant_method,
-                               moe_topk,
-                               norm_topk_prob,
-                               group_moe)};
+  return {FusedExpertMoeFunc(input,
+                             gate_weight,
+                             up_gate_proj_weight,
+                             down_proj_weight,
+                             up_gate_proj_bias,
+                             up_gate_proj_scale,
+                             down_proj_bias,
+                             down_proj_scale,
+                             quant_method,
+                             moe_topk,
+                             norm_topk_prob,
+                             group_moe)};
 }
 
 std::vector<std::vector<int64_t>> FusedExpertMoeInferShape(
@@ -188,7 +197,7 @@ std::vector<std::vector<int64_t>> FusedExpertMoeInferShape(
     const paddle::optional<std::vector<int64_t>>& up_gate_proj_scale_shape,
     const paddle::optional<std::vector<int64_t>>& down_proj_bias_shape,
     const paddle::optional<std::vector<int64_t>>& down_proj_scale_shape) {
-    return {input_shape};
+  return {input_shape};
 }
 
 std::vector<paddle::DataType> FusedExpertMoeInferDtype(
@@ -200,13 +209,14 @@ std::vector<paddle::DataType> FusedExpertMoeInferDtype(
     const paddle::optional<paddle::DataType>& up_gate_proj_scale_dtype,
     const paddle::optional<paddle::DataType>& down_proj_bias_dtype,
     const paddle::optional<paddle::DataType>& down_proj_scale_dtype) {
-    return {input_dtype};
+  return {input_dtype};
 }
 
 /**
  * @brief Fused Mixture-of-Experts (MoE) Operator
  *
- * This operator combines three key MoE operations into a single optimized kernel:
+ * This operator combines three key MoE operations into a single optimized
+ * kernel:
  * 1. moe_dispatch   - Routes tokens to top-k experts using gating network
  * 2. moe_ffn        - Processes tokens through parallel expert FFNs
  * 3. moe_reduce     - Combines expert outputs with routing weights
@@ -219,9 +229,10 @@ std::vector<paddle::DataType> FusedExpertMoeInferDtype(
  *   output = ∑_i^topk(softmax(gate(x))_i * FFN_i(x)
  *
  * Reference Components:
- *   moe_dispatch: Selects top-k experts per token and generates permutation indices
- *   moe_ffn:     Applies SwiGLU activation expert networks in parallel
- *   moe_reduce:  Combines weighted expert outputs and restores original token order
+ *   moe_dispatch: Selects top-k experts per token and generates permutation
+ * indices moe_ffn:     Applies SwiGLU activation expert networks in parallel
+ *   moe_reduce:  Combines weighted expert outputs and restores original token
+ * order
  *
  * Performance Notes:
  * - Recommended hidden_size multiples of 128 for optimal memory alignment
diff --git a/custom_ops/gpu_ops/moe/gptq_marlin_repack.cu b/custom_ops/gpu_ops/moe/gptq_marlin_repack.cu
index 08ddd5a74e..a37f191709 100644
--- a/custom_ops/gpu_ops/moe/gptq_marlin_repack.cu
+++ b/custom_ops/gpu_ops/moe/gptq_marlin_repack.cu
@@ -7,8 +7,10 @@ namespace MARLIN_NAMESPACE_NAME {
 template <int const num_threads, int const num_bits, bool const has_perm>
 __global__ void gptq_marlin_repack_kernel(
     uint32_t const* __restrict__ b_q_weight_ptr,
-    uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
-    int size_k, int size_n) {
+    uint32_t const* __restrict__ perm_ptr,
+    uint32_t* __restrict__ out_ptr,
+    int size_k,
+    int size_n) {
   constexpr int pack_factor = 32 / num_bits;
 
   int k_tiles = size_k / tile_k_size;
@@ -224,7 +226,8 @@ __global__ void gptq_marlin_repack_kernel(
     while (n_tile_id < n_tiles) {
 #pragma unroll
       for (int pipe = 0; pipe < repack_stages; pipe++) {
-        fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id,
+        fetch_to_shared((pipe + repack_stages - 1) % repack_stages,
+                        k_tile_id,
                         n_tile_id + pipe + repack_stages - 1);
         repack_tile(pipe, k_tile_id, n_tile_id + pipe);
         wait_for_stage();
@@ -234,85 +237,83 @@ __global__ void gptq_marlin_repack_kernel(
   }
 }
 
-}  // namespace marlin
+}  // namespace MARLIN_NAMESPACE_NAME
 
-#define CALL_IF(NUM_BITS, HAS_PERM)                                         \
-  else if (num_bits == NUM_BITS && has_perm == HAS_PERM) {                  \
-    cudaFuncSetAttribute(                                                   \
-        MARLIN_NAMESPACE_NAME::gptq_marlin_repack_kernel<MARLIN_NAMESPACE_NAME::repack_threads, NUM_BITS, \
-                                          HAS_PERM>,                        \
-        cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);       \
-    MARLIN_NAMESPACE_NAME::gptq_marlin_repack_kernel<MARLIN_NAMESPACE_NAME::repack_threads, NUM_BITS,     \
-                                      HAS_PERM>                             \
-        <<<blocks, MARLIN_NAMESPACE_NAME::repack_threads, max_shared_mem, stream>>>(       \
-            b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n);             \
+#define CALL_IF(NUM_BITS, HAS_PERM)                                        \
+  else if (num_bits == NUM_BITS && has_perm == HAS_PERM) {                 \
+    cudaFuncSetAttribute(MARLIN_NAMESPACE_NAME::gptq_marlin_repack_kernel< \
+                             MARLIN_NAMESPACE_NAME::repack_threads,        \
+                             NUM_BITS,                                     \
+                             HAS_PERM>,                                    \
+                         cudaFuncAttributeMaxDynamicSharedMemorySize,      \
+                         max_shared_mem);                                  \
+    MARLIN_NAMESPACE_NAME::gptq_marlin_repack_kernel<                      \
+        MARLIN_NAMESPACE_NAME::repack_threads,                             \
+        NUM_BITS,                                                          \
+        HAS_PERM>                                                          \
+        <<<blocks,                                                         \
+           MARLIN_NAMESPACE_NAME::repack_threads,                          \
+           max_shared_mem,                                                 \
+           stream>>>(b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n);   \
   }
 
-std::vector<paddle::Tensor> gptq_marlin_repack(paddle::Tensor& b_q_weight, paddle::Tensor& perm,
-                                 int64_t size_k, int64_t size_n,
-                                 int64_t num_bits) {
+std::vector<paddle::Tensor> gptq_marlin_repack(paddle::Tensor& b_q_weight,
+                                               paddle::Tensor& perm,
+                                               int64_t size_k,
+                                               int64_t size_n,
+                                               int64_t num_bits) {
   // Verify compatibility with marlin tile of 16x64
-  PADDLE_ENFORCE(
-      size_k % MARLIN_NAMESPACE_NAME::tile_k_size == 0,
-      "size_k = ", size_k,
-      " is not divisible by tile_k_size = ",
-      MARLIN_NAMESPACE_NAME::tile_k_size);
+  PADDLE_ENFORCE(size_k % MARLIN_NAMESPACE_NAME::tile_k_size == 0,
+                 "size_k = ",
+                 size_k,
+                 " is not divisible by tile_k_size = ",
+                 MARLIN_NAMESPACE_NAME::tile_k_size);
 
-  PADDLE_ENFORCE(
-      size_n % MARLIN_NAMESPACE_NAME::tile_n_size == 0,
-      "size_n = ", size_n,
-      " is not divisible by tile_n_size = ",
-      MARLIN_NAMESPACE_NAME::tile_n_size);
+  PADDLE_ENFORCE(size_n % MARLIN_NAMESPACE_NAME::tile_n_size == 0,
+                 "size_n = ",
+                 size_n,
+                 " is not divisible by tile_n_size = ",
+                 MARLIN_NAMESPACE_NAME::tile_n_size);
 
-  PADDLE_ENFORCE(
-      num_bits == 4 || num_bits == 8,
-      "num_bits must be 4 or 8. Got = ", num_bits);
+  PADDLE_ENFORCE(num_bits == 4 || num_bits == 8,
+                 "num_bits must be 4 or 8. Got = ",
+                 num_bits);
 
   int const pack_factor = 32 / num_bits;
 
   // Verify B
   // shape checks
-  PADDLE_ENFORCE(
-      (size_k / pack_factor) == b_q_weight.dims()[0],
-      "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.dims()[0]);
+  PADDLE_ENFORCE((size_k / pack_factor) == b_q_weight.dims()[0],
+                 "Shape mismatch: b_q_weight.size(0) = ",
+                 b_q_weight.dims()[0]);
 
-  PADDLE_ENFORCE(
-      b_q_weight.dims()[1] == size_n,
-      "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.dims()[1],
-      ", expected size_n = ", size_n);
+  PADDLE_ENFORCE(b_q_weight.dims()[1] == size_n,
+                 "Shape mismatch: b_q_weight.size(1) = ",
+                 b_q_weight.dims()[1],
+                 ", expected size_n = ",
+                 size_n);
 
-    // Verify device and strides
-    PADDLE_ENFORCE(
-        b_q_weight.is_gpu(),
-        "b_q_weight is not on GPU");
+  // Verify device and strides
+  PADDLE_ENFORCE(b_q_weight.is_gpu(), "b_q_weight is not on GPU");
 
-    PADDLE_ENFORCE(
-        b_q_weight.is_contiguous(),
-        "b_q_weight is not contiguous");
+  PADDLE_ENFORCE(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
 
-    PADDLE_ENFORCE(
-        b_q_weight.dtype() == phi::DataType::INT32,
-        "b_q_weight type is not kInt");
+  PADDLE_ENFORCE(b_q_weight.dtype() == phi::DataType::INT32,
+                 "b_q_weight type is not kInt");
 
-    PADDLE_ENFORCE(
-        perm.is_gpu(),
-        "perm is not on GPU");
+  PADDLE_ENFORCE(perm.is_gpu(), "perm is not on GPU");
 
-    PADDLE_ENFORCE(
-        perm.is_contiguous(),
-        "perm is not contiguous");
+  PADDLE_ENFORCE(perm.is_contiguous(), "perm is not contiguous");
 
-    PADDLE_ENFORCE(
-        perm.dtype() == phi::DataType::INT32,
-        "perm type is not kInt");
+  PADDLE_ENFORCE(perm.dtype() == phi::DataType::INT32, "perm type is not kInt");
 
   // Alloc buffers
   // const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight));
-  paddle::Tensor out = paddle::empty(
-      {size_k / MARLIN_NAMESPACE_NAME::tile_size, size_n * MARLIN_NAMESPACE_NAME::tile_size / pack_factor},
-      b_q_weight.dtype(),
-      b_q_weight.place());
-
+  paddle::Tensor out =
+      paddle::empty({size_k / MARLIN_NAMESPACE_NAME::tile_size,
+                     size_n * MARLIN_NAMESPACE_NAME::tile_size / pack_factor},
+                    b_q_weight.dtype(),
+                    b_q_weight.place());
 
   // Detect if there is act_order
   bool has_perm = perm.dims()[0] != 0;
@@ -330,13 +331,11 @@ std::vector<paddle::Tensor> gptq_marlin_repack(paddle::Tensor& b_q_weight, paddl
   cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
 
   int max_shared_mem = 0;
-  cudaDeviceGetAttribute(&max_shared_mem,
-                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
+  cudaDeviceGetAttribute(
+      &max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
   // TORCH_CHECK(max_shared_mem > 0);
   PADDLE_ENFORCE(
-    max_shared_mem > 0,
-    "max_shared_mem must be > 0. Got = ", max_shared_mem);
-
+      max_shared_mem > 0, "max_shared_mem must be > 0. Got = ", max_shared_mem);
 
   if (false) {
   }
@@ -347,11 +346,11 @@ std::vector<paddle::Tensor> gptq_marlin_repack(paddle::Tensor& b_q_weight, paddl
   else {
     // TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
     //             ", has_perm = ", has_perm);
-    PADDLE_ENFORCE(
-        false,
-        "Unsupported repack config: num_bits = ", num_bits,
-        ", has_perm = ", has_perm);
-
+    PADDLE_ENFORCE(false,
+                   "Unsupported repack config: num_bits = ",
+                   num_bits,
+                   ", has_perm = ",
+                   has_perm);
   }
 
   return {out};
@@ -360,9 +359,5 @@ std::vector<paddle::Tensor> gptq_marlin_repack(paddle::Tensor& b_q_weight, paddl
 PD_BUILD_STATIC_OP(gptq_marlin_repack)
     .Inputs({"b_q_weight", "perm"})
     .Outputs({"out"})
-    .Attrs({
-        "size_k: int64_t",
-        "size_n: int64_t",
-        "num_bits: int64_t"
-    })
+    .Attrs({"size_k: int64_t", "size_n: int64_t", "num_bits: int64_t"})
     .SetKernelFn(PD_KERNEL(gptq_marlin_repack));
diff --git a/custom_ops/gpu_ops/moe/group_swiglu_with_masked.cu b/custom_ops/gpu_ops/moe/group_swiglu_with_masked.cu
index 5df0e78b98..bf3ed0d8fc 100644
--- a/custom_ops/gpu_ops/moe/group_swiglu_with_masked.cu
+++ b/custom_ops/gpu_ops/moe/group_swiglu_with_masked.cu
@@ -19,115 +19,119 @@
 #pragma once
 
 template <typename index, typename T, int VecSize>
-__global__ void group_swiglu_with_masked_kernel(T* act_out,
-                                 const T* input,
-                                 const index *token_nums_per_expert,
-                                 const int64_t group_num,
-                                 const int64_t group_size,
-                                 const int64_t hidden_dim) {
-    int64_t global_idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
-    const int64_t num = group_num * group_size * hidden_dim;
-    using LoadT = AlignedVector<T, VecSize>;
-    LoadT src_vec0, src_vec1;
-    LoadT res_vec;
+__global__ void group_swiglu_with_masked_kernel(
+    T* act_out,
+    const T* input,
+    const index* token_nums_per_expert,
+    const int64_t group_num,
+    const int64_t group_size,
+    const int64_t hidden_dim) {
+  int64_t global_idx =
+      static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const int64_t num = group_num * group_size * hidden_dim;
+  using LoadT = AlignedVector<T, VecSize>;
+  LoadT src_vec0, src_vec1;
+  LoadT res_vec;
 
-    int64_t block_id = static_cast<int64_t>(blockIdx.x);
-    const int lane_idx = threadIdx.x % 32;
+  int64_t block_id = static_cast<int64_t>(blockIdx.x);
+  const int lane_idx = threadIdx.x % 32;
 
-    while(true) {
-        int dealt_group_id = -1;
-        int dealt_seq_id = -1;
-        if (lane_idx == 0 ) {
-            int cumsum1 = 0;
-            int cumsum2 = 0;
-            for (int i = 0; i < group_num; i++) {
-                int tmp = token_nums_per_expert[i];
-                cumsum2 += tmp;
-                if (block_id >= cumsum1 && block_id < cumsum2) {
-                    dealt_group_id = i;
-                    dealt_seq_id = block_id - cumsum1;
-                    break;
-                }
-                cumsum1 += tmp;
-            }
+  while (true) {
+    int dealt_group_id = -1;
+    int dealt_seq_id = -1;
+    if (lane_idx == 0) {
+      int cumsum1 = 0;
+      int cumsum2 = 0;
+      for (int i = 0; i < group_num; i++) {
+        int tmp = token_nums_per_expert[i];
+        cumsum2 += tmp;
+        if (block_id >= cumsum1 && block_id < cumsum2) {
+          dealt_group_id = i;
+          dealt_seq_id = block_id - cumsum1;
+          break;
         }
-        dealt_group_id = __shfl_sync(0xffffffff, dealt_group_id, 0);
-        dealt_seq_id =   __shfl_sync(0xffffffff, dealt_seq_id, 0);
-        if (dealt_group_id < 0) break;
-
-        const int64_t r_offset = (dealt_group_id * group_size + dealt_seq_id) * hidden_dim * 2;
-        const int64_t w_offset = (dealt_group_id * group_size + dealt_seq_id) * hidden_dim;
-
-        for (int64_t col_id = threadIdx.x * VecSize; col_id < hidden_dim; col_id += blockDim.x * VecSize) {
-
-            Load<T, VecSize>(&input[r_offset + col_id], &src_vec0);
-            Load<T, VecSize>(&input[r_offset + col_id + hidden_dim], &src_vec1);
-
-            for (int j = 0; j < VecSize; ++j) {
-                float a = static_cast<float>(src_vec0[j]);
-                float b = static_cast<float>(src_vec1[j]);
-                float res = b * a / (1.f + exp(-a));
-                res_vec[j] = static_cast<T>(res);
-            }
-
-            Store<T, VecSize>(res_vec, &act_out[w_offset + col_id]);
-        }
-        block_id += gridDim.x;
+        cumsum1 += tmp;
+      }
     }
+    dealt_group_id = __shfl_sync(0xffffffff, dealt_group_id, 0);
+    dealt_seq_id = __shfl_sync(0xffffffff, dealt_seq_id, 0);
+    if (dealt_group_id < 0) break;
+
+    const int64_t r_offset =
+        (dealt_group_id * group_size + dealt_seq_id) * hidden_dim * 2;
+    const int64_t w_offset =
+        (dealt_group_id * group_size + dealt_seq_id) * hidden_dim;
+
+    for (int64_t col_id = threadIdx.x * VecSize; col_id < hidden_dim;
+         col_id += blockDim.x * VecSize) {
+      Load<T, VecSize>(&input[r_offset + col_id], &src_vec0);
+      Load<T, VecSize>(&input[r_offset + col_id + hidden_dim], &src_vec1);
+
+      for (int j = 0; j < VecSize; ++j) {
+        float a = static_cast<float>(src_vec0[j]);
+        float b = static_cast<float>(src_vec1[j]);
+        float res = b * a / (1.f + exp(-a));
+        res_vec[j] = static_cast<T>(res);
+      }
+
+      Store<T, VecSize>(res_vec, &act_out[w_offset + col_id]);
+    }
+    block_id += gridDim.x;
+  }
 }
 
-paddle::Tensor GroupSwigluWithMasked(const paddle::Tensor& fc1_out_tensor,
-                                                  const paddle::Tensor& token_nums_per_expert
-                                                  )
-{
-    const int64_t group_num = token_nums_per_expert.shape()[0];
-    const int64_t group_size = fc1_out_tensor.shape()[1];
-    const int64_t hidden_dim = fc1_out_tensor.shape()[2] / 2;
-    auto act_out_tensor = GetEmptyTensor({group_num, group_size, hidden_dim}, fc1_out_tensor.dtype(), fc1_out_tensor.place());
+paddle::Tensor GroupSwigluWithMasked(
+    const paddle::Tensor& fc1_out_tensor,
+    const paddle::Tensor& token_nums_per_expert) {
+  const int64_t group_num = token_nums_per_expert.shape()[0];
+  const int64_t group_size = fc1_out_tensor.shape()[1];
+  const int64_t hidden_dim = fc1_out_tensor.shape()[2] / 2;
+  auto act_out_tensor = GetEmptyTensor({group_num, group_size, hidden_dim},
+                                       fc1_out_tensor.dtype(),
+                                       fc1_out_tensor.place());
 
-    constexpr int VecSize = 8;
-    PD_CHECK(fc1_out_tensor.dtype() == paddle::DataType::BFLOAT16);
-    PD_CHECK(hidden_dim % VecSize == 0);
+  constexpr int VecSize = 8;
+  PD_CHECK(fc1_out_tensor.dtype() == paddle::DataType::BFLOAT16);
+  PD_CHECK(hidden_dim % VecSize == 0);
 
-    constexpr paddle::DataType D = paddle::DataType::BFLOAT16;
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+  constexpr paddle::DataType D = paddle::DataType::BFLOAT16;
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    const int threads = 512;
-    const int blocks = 256;
+  const int threads = 512;
+  const int blocks = 256;
 
-    #define dispatch_by_index(index) {\
-    group_swiglu_with_masked_kernel<index, DataType_, VecSize><<<blocks, threads, 0, fc1_out_tensor.stream()>>>(\
-        reinterpret_cast<DataType_*>(const_cast<data_t*>(act_out_tensor.data<data_t>())),\
-        reinterpret_cast<const DataType_*>(fc1_out_tensor.data<data_t>()),\
-        token_nums_per_expert.data<index>(),\
-        group_num,\
-        group_size,\
-        hidden_dim\
-    );} while(0)
-    if (token_nums_per_expert.dtype() == paddle::DataType::INT64) {
-        dispatch_by_index(int64_t);
-    } else if(token_nums_per_expert.dtype() == paddle::DataType::INT32) {
-        dispatch_by_index(int32_t);
-    } else {
-        PD_THROW("Unsupported token_nums_per_expert's data dtype.");
-    }
+#define dispatch_by_index(index)                                               \
+  {                                                                            \
+    group_swiglu_with_masked_kernel<index, DataType_, VecSize>                 \
+        <<<blocks, threads, 0, fc1_out_tensor.stream()>>>(                     \
+            reinterpret_cast<DataType_*>(                                      \
+                const_cast<data_t*>(act_out_tensor.data<data_t>())),           \
+            reinterpret_cast<const DataType_*>(fc1_out_tensor.data<data_t>()), \
+            token_nums_per_expert.data<index>(),                               \
+            group_num,                                                         \
+            group_size,                                                        \
+            hidden_dim);                                                       \
+  }                                                                            \
+  while (0)
+  if (token_nums_per_expert.dtype() == paddle::DataType::INT64) {
+    dispatch_by_index(int64_t);
+  } else if (token_nums_per_expert.dtype() == paddle::DataType::INT32) {
+    dispatch_by_index(int32_t);
+  } else {
+    PD_THROW("Unsupported token_nums_per_expert's data dtype.");
+  }
 
-    return act_out_tensor;
+  return act_out_tensor;
 }
 
-
-
-
 std::vector<paddle::Tensor> GroupSwigluWithMaskedWrapper(
-    const paddle::Tensor& input,
-    const paddle::Tensor& token_nums_per_expert) {
-     return {GroupSwigluWithMasked(input, token_nums_per_expert)};
+    const paddle::Tensor& input, const paddle::Tensor& token_nums_per_expert) {
+  return {GroupSwigluWithMasked(input, token_nums_per_expert)};
 }
 
 PD_BUILD_STATIC_OP(group_swiglu_with_masked)
-    .Inputs({"input",
-             "token_nums_per_expert"})
+    .Inputs({"input", "token_nums_per_expert"})
     .Outputs({"output_tensor"})
     .SetKernelFn(PD_KERNEL(GroupSwigluWithMaskedWrapper));
diff --git a/custom_ops/gpu_ops/moe/group_swiglu_with_masked.h b/custom_ops/gpu_ops/moe/group_swiglu_with_masked.h
index 2abe56d663..74d43515eb 100644
--- a/custom_ops/gpu_ops/moe/group_swiglu_with_masked.h
+++ b/custom_ops/gpu_ops/moe/group_swiglu_with_masked.h
@@ -15,6 +15,6 @@
 #pragma once
 #include "helper.h"
 
-paddle::Tensor
-GroupSwigluWithMasked(const paddle::Tensor &fc1_out_tensor,
-                      const paddle::Tensor &token_nums_per_expert);
+paddle::Tensor GroupSwigluWithMasked(
+    const paddle::Tensor &fc1_out_tensor,
+    const paddle::Tensor &token_nums_per_expert);
diff --git a/custom_ops/gpu_ops/moe/moe_deepgemm_depermute.cu b/custom_ops/gpu_ops/moe/moe_deepgemm_depermute.cu
index 2378201aa6..1aa444bb89 100644
--- a/custom_ops/gpu_ops/moe/moe_deepgemm_depermute.cu
+++ b/custom_ops/gpu_ops/moe/moe_deepgemm_depermute.cu
@@ -14,146 +14,164 @@
 
 #include "helper.h"
 
-template<typename T, int VecSize, int TopK>
-__global__ void MoEDeepGEMMDePermuteKernel(T* out, const T* ffn_out, const int* permute_indices_per_token, const int64_t* topk_idx, const float* topk_weights, const int token_num, const int num_vecs, const int hidden, const int max_tokens_per_expert) {
-    AlignedVector<T, VecSize> in_vec;
+template <typename T, int VecSize, int TopK>
+__global__ void MoEDeepGEMMDePermuteKernel(T* out,
+                                           const T* ffn_out,
+                                           const int* permute_indices_per_token,
+                                           const int64_t* topk_idx,
+                                           const float* topk_weights,
+                                           const int token_num,
+                                           const int num_vecs,
+                                           const int hidden,
+                                           const int max_tokens_per_expert) {
+  AlignedVector<T, VecSize> in_vec;
 
-    AlignedVector<T, VecSize> acc_vec[TopK];
+  AlignedVector<T, VecSize> acc_vec[TopK];
 
-    const int bid = blockIdx.x;
-    const int wid = threadIdx.x / 32;
-    const int tid = threadIdx.x % 32;
-    extern __shared__ char shm[]; // TopK * hidden
-    T* shm_hidden = reinterpret_cast<T*>(shm);
+  const int bid = blockIdx.x;
+  const int wid = threadIdx.x / 32;
+  const int tid = threadIdx.x % 32;
+  extern __shared__ char shm[];  // TopK * hidden
+  T* shm_hidden = reinterpret_cast<T*>(shm);
 
-    for (int token_idx = bid; token_idx < token_num; token_idx += gridDim.x) {
-        int src_expert_id = topk_idx[token_idx * TopK + wid];
-        int src_expert_token = permute_indices_per_token[token_idx * TopK + wid];
-        float weight = topk_weights[token_idx * TopK + wid];
+  for (int token_idx = bid; token_idx < token_num; token_idx += gridDim.x) {
+    int src_expert_id = topk_idx[token_idx * TopK + wid];
+    int src_expert_token = permute_indices_per_token[token_idx * TopK + wid];
+    float weight = topk_weights[token_idx * TopK + wid];
 
-        for (int hidden_vec_id = tid; hidden_vec_id < num_vecs; hidden_vec_id += 32) {
-            Load<T, VecSize>(ffn_out + src_expert_id * max_tokens_per_expert * hidden + src_expert_token * hidden + hidden_vec_id * VecSize, &in_vec);
+    for (int hidden_vec_id = tid; hidden_vec_id < num_vecs;
+         hidden_vec_id += 32) {
+      Load<T, VecSize>(ffn_out +
+                           src_expert_id * max_tokens_per_expert * hidden +
+                           src_expert_token * hidden + hidden_vec_id * VecSize,
+                       &in_vec);
 #pragma unroll
-            for (int i = 0; i < VecSize; i++) {
-                in_vec[i] *= weight;
-            }
-            Store<T, VecSize>(in_vec, shm_hidden + wid * hidden + hidden_vec_id * VecSize);
-        }
-
-        __syncthreads();
-
-        for (int hidden_vec_id = threadIdx.x; hidden_vec_id < num_vecs; hidden_vec_id += blockDim.x) {
-#pragma unroll
-            for (int topk_id = 0; topk_id < TopK; topk_id++) {
-                Load<T, VecSize>(shm_hidden + topk_id * hidden + hidden_vec_id * VecSize, &acc_vec[topk_id]);
-            }
-#pragma unroll
-            for (int i = 0; i < VecSize; i++) {
-#pragma unroll
-                 for (int topk_id = 1; topk_id < TopK; topk_id++) {
-                    acc_vec[0][i] += acc_vec[topk_id][i];
-                 }
-            }
-            Store<T, VecSize>(acc_vec[0], out + token_idx * hidden + hidden_vec_id * VecSize);
-        }
-
+      for (int i = 0; i < VecSize; i++) {
+        in_vec[i] *= weight;
+      }
+      Store<T, VecSize>(in_vec,
+                        shm_hidden + wid * hidden + hidden_vec_id * VecSize);
     }
+
+    __syncthreads();
+
+    for (int hidden_vec_id = threadIdx.x; hidden_vec_id < num_vecs;
+         hidden_vec_id += blockDim.x) {
+#pragma unroll
+      for (int topk_id = 0; topk_id < TopK; topk_id++) {
+        Load<T, VecSize>(
+            shm_hidden + topk_id * hidden + hidden_vec_id * VecSize,
+            &acc_vec[topk_id]);
+      }
+#pragma unroll
+      for (int i = 0; i < VecSize; i++) {
+#pragma unroll
+        for (int topk_id = 1; topk_id < TopK; topk_id++) {
+          acc_vec[0][i] += acc_vec[topk_id][i];
+        }
+      }
+      Store<T, VecSize>(acc_vec[0],
+                        out + token_idx * hidden + hidden_vec_id * VecSize);
+    }
+  }
 }
 
 template <paddle::DataType D>
 std::vector<paddle::Tensor> MoEDeepGEMMDePermuteDispatch(
-    const paddle::Tensor& ffn_out, // [num_experts, max_tokens_per_expert, hidden]
-    const paddle::Tensor& permute_indices_per_token, // [token_num, topk}]
+    const paddle::Tensor&
+        ffn_out,  // [num_experts, max_tokens_per_expert, hidden]
+    const paddle::Tensor& permute_indices_per_token,  // [token_num, topk}]
     const paddle::Tensor& topk_idx,
-    const paddle::Tensor& topk_weights
-) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+    const paddle::Tensor& topk_weights) {
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    const int token_num = permute_indices_per_token.shape()[0];
-    const int max_tokens_per_expert = ffn_out.shape()[1];
-    const int hidden = ffn_out.shape()[2];
-    const int topk = permute_indices_per_token.shape()[1];
+  const int token_num = permute_indices_per_token.shape()[0];
+  const int max_tokens_per_expert = ffn_out.shape()[1];
+  const int hidden = ffn_out.shape()[2];
+  const int topk = permute_indices_per_token.shape()[1];
 
-    auto place = ffn_out.place();
-    auto stream = ffn_out.stream();
+  auto place = ffn_out.place();
+  auto stream = ffn_out.stream();
 
-    auto out = GetEmptyTensor({token_num, hidden}, ffn_out.dtype(), place);
+  auto out = GetEmptyTensor({token_num, hidden}, ffn_out.dtype(), place);
 
-    constexpr int VecSize = 16 / sizeof(data_t);
-    int blocks = 32 * topk;
-    int grids = min(132 * 4, token_num);
-    int num_vecs = hidden / VecSize;
+  constexpr int VecSize = 16 / sizeof(data_t);
+  int blocks = 32 * topk;
+  int grids = min(132 * 4, token_num);
+  int num_vecs = hidden / VecSize;
 
-    assert(blocks <= 1024);
-    int dyn_smem_size = 0;
+  assert(blocks <= 1024);
+  int dyn_smem_size = 0;
 
-    switch (topk) {
-        case 4:
-        dyn_smem_size =  topk * hidden * sizeof(DataType_);
-        if (dyn_smem_size >= (48 << 10)) {
-            cudaFuncSetAttribute(
-                MoEDeepGEMMDePermuteKernel<DataType_, VecSize, 4>,
-                cudaFuncAttributeMaxDynamicSharedMemorySize,
-                dyn_smem_size);
-        }
-        MoEDeepGEMMDePermuteKernel<DataType_, VecSize, 4><<<grids, blocks, dyn_smem_size, stream>>>(
-            reinterpret_cast<DataType_*>(out.data<data_t>()),
-            reinterpret_cast<const DataType_*>(ffn_out.data<data_t>()),
-            permute_indices_per_token.data<int32_t>(),
-            topk_idx.data<int64_t>(),
-            topk_weights.data<float>(),
-            token_num, num_vecs, hidden, max_tokens_per_expert
-        );
-        break;
+  switch (topk) {
+    case 4:
+      dyn_smem_size = topk * hidden * sizeof(DataType_);
+      if (dyn_smem_size >= (48 << 10)) {
+        cudaFuncSetAttribute(MoEDeepGEMMDePermuteKernel<DataType_, VecSize, 4>,
+                             cudaFuncAttributeMaxDynamicSharedMemorySize,
+                             dyn_smem_size);
+      }
+      MoEDeepGEMMDePermuteKernel<DataType_, VecSize, 4>
+          <<<grids, blocks, dyn_smem_size, stream>>>(
+              reinterpret_cast<DataType_*>(out.data<data_t>()),
+              reinterpret_cast<const DataType_*>(ffn_out.data<data_t>()),
+              permute_indices_per_token.data<int32_t>(),
+              topk_idx.data<int64_t>(),
+              topk_weights.data<float>(),
+              token_num,
+              num_vecs,
+              hidden,
+              max_tokens_per_expert);
+      break;
 
-        case 8:
-        dyn_smem_size =  topk * hidden * sizeof(DataType_);
-        if (dyn_smem_size >= (48 << 10)) {
-            cudaFuncSetAttribute(
-                MoEDeepGEMMDePermuteKernel<DataType_, VecSize, 8>,
-                cudaFuncAttributeMaxDynamicSharedMemorySize,
-                dyn_smem_size);
-        }
-        MoEDeepGEMMDePermuteKernel<DataType_, VecSize, 8><<<grids, blocks, topk * hidden * sizeof(DataType_), stream>>>(
-            reinterpret_cast<DataType_*>(out.data<data_t>()),
-            reinterpret_cast<const DataType_*>(ffn_out.data<data_t>()),
-            permute_indices_per_token.data<int32_t>(),
-            topk_idx.data<int64_t>(),
-            topk_weights.data<float>(),
-            token_num, num_vecs, hidden, max_tokens_per_expert
-        );
-        break;
+    case 8:
+      dyn_smem_size = topk * hidden * sizeof(DataType_);
+      if (dyn_smem_size >= (48 << 10)) {
+        cudaFuncSetAttribute(MoEDeepGEMMDePermuteKernel<DataType_, VecSize, 8>,
+                             cudaFuncAttributeMaxDynamicSharedMemorySize,
+                             dyn_smem_size);
+      }
+      MoEDeepGEMMDePermuteKernel<DataType_, VecSize, 8>
+          <<<grids, blocks, topk * hidden * sizeof(DataType_), stream>>>(
+              reinterpret_cast<DataType_*>(out.data<data_t>()),
+              reinterpret_cast<const DataType_*>(ffn_out.data<data_t>()),
+              permute_indices_per_token.data<int32_t>(),
+              topk_idx.data<int64_t>(),
+              topk_weights.data<float>(),
+              token_num,
+              num_vecs,
+              hidden,
+              max_tokens_per_expert);
+      break;
 
-        default:
-        PD_THROW("Unsupported topk");
-    }
-    return {out};
+    default:
+      PD_THROW("Unsupported topk");
+  }
+  return {out};
 }
 
-
 std::vector<paddle::Tensor> MoEDeepGEMMDePermute(
-    const paddle::Tensor& ffn_out, // [num_experts, max_tokens_per_expert, hidden]
-    const paddle::Tensor& permute_indices_per_token, // [token_num, topk}]
+    const paddle::Tensor&
+        ffn_out,  // [num_experts, max_tokens_per_expert, hidden]
+    const paddle::Tensor& permute_indices_per_token,  // [token_num, topk}]
     const paddle::Tensor& topk_idx,
-    const paddle::Tensor& topk_weights
-) {
-    switch (ffn_out.dtype()) {
-        case paddle::DataType::BFLOAT16:
-            return MoEDeepGEMMDePermuteDispatch<paddle::DataType::BFLOAT16>(
-                ffn_out, permute_indices_per_token, topk_idx, topk_weights
-            );
-        case paddle::DataType::FLOAT16:
-            return MoEDeepGEMMDePermuteDispatch<paddle::DataType::FLOAT16>(
-                ffn_out, permute_indices_per_token, topk_idx, topk_weights
-            );
-        default:
-            PD_THROW("Unsupported data type");
-    }
+    const paddle::Tensor& topk_weights) {
+  switch (ffn_out.dtype()) {
+    case paddle::DataType::BFLOAT16:
+      return MoEDeepGEMMDePermuteDispatch<paddle::DataType::BFLOAT16>(
+          ffn_out, permute_indices_per_token, topk_idx, topk_weights);
+    case paddle::DataType::FLOAT16:
+      return MoEDeepGEMMDePermuteDispatch<paddle::DataType::FLOAT16>(
+          ffn_out, permute_indices_per_token, topk_idx, topk_weights);
+    default:
+      PD_THROW("Unsupported data type");
+  }
 }
 
 PD_BUILD_STATIC_OP(moe_deepgemm_depermute)
-    .Inputs({"ffn_out", "permute_indices_per_token", "topk_idx", "topk_weights"})
+    .Inputs(
+        {"ffn_out", "permute_indices_per_token", "topk_idx", "topk_weights"})
     .Outputs({"out"})
     .SetKernelFn(PD_KERNEL(MoEDeepGEMMDePermute));
diff --git a/custom_ops/gpu_ops/moe/moe_deepgemm_permute.cu b/custom_ops/gpu_ops/moe/moe_deepgemm_permute.cu
index ec44a5bfca..8c0f36a5e8 100644
--- a/custom_ops/gpu_ops/moe/moe_deepgemm_permute.cu
+++ b/custom_ops/gpu_ops/moe/moe_deepgemm_permute.cu
@@ -15,29 +15,41 @@
 #include "helper.h"
 
 // topk warps
-template<typename T, int VecSize>
-__global__ void MoEDeepGEMMPermuteKernel(T* out, int* token_nums_per_expert, int* permute_indices_per_token, const T* x, const int64_t* topk_idx, const int token_num, const int topk, const int num_vecs, const int hidden, const int max_tokens_per_expert) {
+template <typename T, int VecSize>
+__global__ void MoEDeepGEMMPermuteKernel(T* out,
+                                         int* token_nums_per_expert,
+                                         int* permute_indices_per_token,
+                                         const T* x,
+                                         const int64_t* topk_idx,
+                                         const int token_num,
+                                         const int topk,
+                                         const int num_vecs,
+                                         const int hidden,
+                                         const int max_tokens_per_expert) {
+  AlignedVector<T, VecSize> in_vec;
 
-    AlignedVector<T, VecSize> in_vec;
-
-    const int bid = blockIdx.x;
-    const int wid = threadIdx.x / 32;
-    const int tid = threadIdx.x % 32;
-    for (int token_idx = bid; token_idx < token_num; token_idx += gridDim.x) {
-        const int tgt_expert_id = topk_idx[token_idx * topk + wid];
-        int tgt_expert_token;
-        if (tid == 0) {
-            tgt_expert_token = atomicAdd(token_nums_per_expert + tgt_expert_id, 1);
-            permute_indices_per_token[token_idx * topk + wid] = tgt_expert_token;
-        }
-        tgt_expert_token = __shfl_sync(0xFFFFFFFF, tgt_expert_token, 0);
-
-
-        for (int hidden_vec_id = tid; hidden_vec_id < num_vecs; hidden_vec_id += 32) {
-            Load<T, VecSize>(x + token_idx * hidden + hidden_vec_id * VecSize, &in_vec);
-            Store<T, VecSize>(in_vec, out + tgt_expert_id * max_tokens_per_expert * hidden + tgt_expert_token * hidden + hidden_vec_id * VecSize);
-        }
+  const int bid = blockIdx.x;
+  const int wid = threadIdx.x / 32;
+  const int tid = threadIdx.x % 32;
+  for (int token_idx = bid; token_idx < token_num; token_idx += gridDim.x) {
+    const int tgt_expert_id = topk_idx[token_idx * topk + wid];
+    int tgt_expert_token;
+    if (tid == 0) {
+      tgt_expert_token = atomicAdd(token_nums_per_expert + tgt_expert_id, 1);
+      permute_indices_per_token[token_idx * topk + wid] = tgt_expert_token;
     }
+    tgt_expert_token = __shfl_sync(0xFFFFFFFF, tgt_expert_token, 0);
+
+    for (int hidden_vec_id = tid; hidden_vec_id < num_vecs;
+         hidden_vec_id += 32) {
+      Load<T, VecSize>(x + token_idx * hidden + hidden_vec_id * VecSize,
+                       &in_vec);
+      Store<T, VecSize>(in_vec,
+                        out + tgt_expert_id * max_tokens_per_expert * hidden +
+                            tgt_expert_token * hidden +
+                            hidden_vec_id * VecSize);
+    }
+  }
 }
 
 template <paddle::DataType D>
@@ -45,71 +57,78 @@ std::vector<paddle::Tensor> MoEDeepGEMMPermuteDispatch(
     const paddle::Tensor& x,
     const paddle::Tensor& topk_idx,
     const int num_experts,
-    const int max_tokens_per_expert
-) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+    const int max_tokens_per_expert) {
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    const int token_num = x.shape()[0];
-    const int hidden = x.shape()[1];
-    const int topk = topk_idx.shape()[1];
+  const int token_num = x.shape()[0];
+  const int hidden = x.shape()[1];
+  const int topk = topk_idx.shape()[1];
 
-    auto place = x.place();
-    auto stream = x.stream();
+  auto place = x.place();
+  auto stream = x.stream();
 
-    auto token_nums_per_expert = GetEmptyTensor({num_experts}, paddle::DataType::INT32, place);
-    auto permute_indices_per_token = GetEmptyTensor({token_num, topk}, paddle::DataType::INT32, place);
+  auto token_nums_per_expert =
+      GetEmptyTensor({num_experts}, paddle::DataType::INT32, place);
+  auto permute_indices_per_token =
+      GetEmptyTensor({token_num, topk}, paddle::DataType::INT32, place);
 
-    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(token_nums_per_expert.data<int32_t>(), 0, num_experts * sizeof(int32_t), stream));
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaMemsetAsync(token_nums_per_expert.data<int32_t>(),
+                      0,
+                      num_experts * sizeof(int32_t),
+                      stream));
 
-    auto permute_output = GetEmptyTensor({num_experts, max_tokens_per_expert, hidden}, x.dtype(), place);
+  auto permute_output = GetEmptyTensor(
+      {num_experts, max_tokens_per_expert, hidden}, x.dtype(), place);
 
-    auto permute_output_data = permute_output.data<data_t>();
+  auto permute_output_data = permute_output.data<data_t>();
 
-    constexpr int VecSize = 16 / sizeof(data_t);
+  constexpr int VecSize = 16 / sizeof(data_t);
 
-    int blocks = 32 * topk;
-    int grids = min(132 * 4, token_num);
-    int num_vecs = hidden / VecSize;
+  int blocks = 32 * topk;
+  int grids = min(132 * 4, token_num);
+  int num_vecs = hidden / VecSize;
 
-    assert(blocks <= 1024);
+  assert(blocks <= 1024);
 
-    MoEDeepGEMMPermuteKernel<DataType_, VecSize><<<grids, blocks, 0, stream>>>(
-        reinterpret_cast<DataType_*>(permute_output_data),
-        token_nums_per_expert.data<int32_t>(),
-        permute_indices_per_token.data<int32_t>(),
-        reinterpret_cast<const DataType_ *>(x.data<data_t>()),
-        topk_idx.data<int64_t>(),
-        token_num, topk, num_vecs,
-        hidden, max_tokens_per_expert
-    );
+  MoEDeepGEMMPermuteKernel<DataType_, VecSize><<<grids, blocks, 0, stream>>>(
+      reinterpret_cast<DataType_*>(permute_output_data),
+      token_nums_per_expert.data<int32_t>(),
+      permute_indices_per_token.data<int32_t>(),
+      reinterpret_cast<const DataType_*>(x.data<data_t>()),
+      topk_idx.data<int64_t>(),
+      token_num,
+      topk,
+      num_vecs,
+      hidden,
+      max_tokens_per_expert);
 
-    return {permute_output, token_nums_per_expert, permute_indices_per_token};
+  return {permute_output, token_nums_per_expert, permute_indices_per_token};
 }
 
 std::vector<paddle::Tensor> MoEDeepGEMMPermute(
     const paddle::Tensor& x,
     const paddle::Tensor& topk_idx,
     const int num_experts,
-    const int max_tokens_per_expert
-) {
-    switch (x.dtype()) {
-        case paddle::DataType::BFLOAT16:
-            return MoEDeepGEMMPermuteDispatch<paddle::DataType::BFLOAT16>(
-                x, topk_idx, num_experts, max_tokens_per_expert
-            );
-        case paddle::DataType::FLOAT16:
-            return MoEDeepGEMMPermuteDispatch<paddle::DataType::FLOAT16>(
-                x, topk_idx, num_experts, max_tokens_per_expert
-            );
-        default:
-            PD_THROW("Unsupported data type");
-    }
+    const int max_tokens_per_expert) {
+  switch (x.dtype()) {
+    case paddle::DataType::BFLOAT16:
+      return MoEDeepGEMMPermuteDispatch<paddle::DataType::BFLOAT16>(
+          x, topk_idx, num_experts, max_tokens_per_expert);
+    case paddle::DataType::FLOAT16:
+      return MoEDeepGEMMPermuteDispatch<paddle::DataType::FLOAT16>(
+          x, topk_idx, num_experts, max_tokens_per_expert);
+    default:
+      PD_THROW("Unsupported data type");
+  }
 }
 
 PD_BUILD_STATIC_OP(moe_deepgemm_permute)
     .Inputs({"x", "topk_idx"})
-    .Outputs({"permute_output", "token_nums_per_expert", "permute_indices_per_token"})
+    .Outputs({"permute_output",
+              "token_nums_per_expert",
+              "permute_indices_per_token"})
     .Attrs({"num_experts: int", "max_tokens_per_expert: int"})
     .SetKernelFn(PD_KERNEL(MoEDeepGEMMPermute));
diff --git a/custom_ops/gpu_ops/moe/moe_reduce.cu b/custom_ops/gpu_ops/moe/moe_reduce.cu
index 9a7bad1475..0fa9077428 100644
--- a/custom_ops/gpu_ops/moe/moe_reduce.cu
+++ b/custom_ops/gpu_ops/moe/moe_reduce.cu
@@ -27,8 +27,10 @@ void MoeReduceKernel(const paddle::Tensor &ffn_out,
                      const paddle::Tensor &top_k_indices,
                      const paddle::optional<paddle::Tensor> &down_proj_bias,
                      const bool norm_topk_prob,
-                     const float routed_scaling_factor, const int num_rows,
-                     const int hidden_size, const int topk,
+                     const float routed_scaling_factor,
+                     const int num_rows,
+                     const int hidden_size,
+                     const int topk,
                      paddle::Tensor *output) {
   using namespace phi;
   typedef PDTraits<T> traits_;
@@ -37,19 +39,29 @@ void MoeReduceKernel(const paddle::Tensor &ffn_out,
   auto stream = ffn_out.stream();
 
   finalize_moe_routing_kernelLauncher(
-      ffn_out.data<data_t>(), output->data<data_t>(),
+      ffn_out.data<data_t>(),
+      output->data<data_t>(),
       down_proj_bias ? down_proj_bias->data<data_t>() : nullptr,
-      top_k_weight.data<float>(), permute_indices_per_token.data<int32_t>(),
-      top_k_indices.data<int>(), num_rows, hidden_size, topk,
-      static_cast<int>(1), norm_topk_prob, routed_scaling_factor, stream);
+      top_k_weight.data<float>(),
+      permute_indices_per_token.data<int32_t>(),
+      top_k_indices.data<int>(),
+      num_rows,
+      hidden_size,
+      topk,
+      static_cast<int>(1),
+      norm_topk_prob,
+      routed_scaling_factor,
+      stream);
 }
 
 paddle::Tensor MoeExpertReduceFunc(
-    const paddle::Tensor &ffn_out, const paddle::Tensor &top_k_weight,
+    const paddle::Tensor &ffn_out,
+    const paddle::Tensor &top_k_weight,
     const paddle::Tensor &permute_indices_per_token,
     const paddle::Tensor &top_k_indices,
     const paddle::optional<paddle::Tensor> &down_proj_bias,
-    const bool norm_topk_prob, const float routed_scaling_factor) {
+    const bool norm_topk_prob,
+    const float routed_scaling_factor) {
   const auto input_type = ffn_out.dtype();
   auto place = ffn_out.place();
 
@@ -59,38 +71,57 @@ paddle::Tensor MoeExpertReduceFunc(
 
   auto output = GetEmptyTensor({num_rows, hidden_size}, input_type, place);
 
-  if(num_rows == 0){
+  if (num_rows == 0) {
     return output;
   }
 
   switch (input_type) {
-  case paddle::DataType::BFLOAT16:
-    MoeReduceKernel<paddle::DataType::BFLOAT16>(
-        ffn_out, top_k_weight, permute_indices_per_token, top_k_indices,
-        down_proj_bias, norm_topk_prob, routed_scaling_factor, num_rows, hidden_size,
-        topk, &output);
-    break;
-  case paddle::DataType::FLOAT16:
-    MoeReduceKernel<paddle::DataType::BFLOAT16>(
-        ffn_out, top_k_weight, permute_indices_per_token, top_k_indices,
-        down_proj_bias, norm_topk_prob, routed_scaling_factor, num_rows, hidden_size,
-        topk, &output);
-    break;
-  default:
-    PD_THROW("Unsupported data type for MoeDispatchKernel");
+    case paddle::DataType::BFLOAT16:
+      MoeReduceKernel<paddle::DataType::BFLOAT16>(ffn_out,
+                                                  top_k_weight,
+                                                  permute_indices_per_token,
+                                                  top_k_indices,
+                                                  down_proj_bias,
+                                                  norm_topk_prob,
+                                                  routed_scaling_factor,
+                                                  num_rows,
+                                                  hidden_size,
+                                                  topk,
+                                                  &output);
+      break;
+    case paddle::DataType::FLOAT16:
+      MoeReduceKernel<paddle::DataType::BFLOAT16>(ffn_out,
+                                                  top_k_weight,
+                                                  permute_indices_per_token,
+                                                  top_k_indices,
+                                                  down_proj_bias,
+                                                  norm_topk_prob,
+                                                  routed_scaling_factor,
+                                                  num_rows,
+                                                  hidden_size,
+                                                  topk,
+                                                  &output);
+      break;
+    default:
+      PD_THROW("Unsupported data type for MoeDispatchKernel");
   }
   return output;
 }
 
-std::vector<paddle::Tensor>
-MoeExpertReduce(const paddle::Tensor &ffn_out,
-                const paddle::Tensor &top_k_weight,
-                const paddle::Tensor &permute_indices_per_token,
-                const paddle::Tensor &top_k_indices,
-                const paddle::optional<paddle::Tensor> &down_proj_bias,
-                const bool norm_topk_prob, const float routed_scaling_factor) {
-  return {MoeExpertReduceFunc(ffn_out, top_k_weight, permute_indices_per_token,
-                              top_k_indices, down_proj_bias, norm_topk_prob,
+std::vector<paddle::Tensor> MoeExpertReduce(
+    const paddle::Tensor &ffn_out,
+    const paddle::Tensor &top_k_weight,
+    const paddle::Tensor &permute_indices_per_token,
+    const paddle::Tensor &top_k_indices,
+    const paddle::optional<paddle::Tensor> &down_proj_bias,
+    const bool norm_topk_prob,
+    const float routed_scaling_factor) {
+  return {MoeExpertReduceFunc(ffn_out,
+                              top_k_weight,
+                              permute_indices_per_token,
+                              top_k_indices,
+                              down_proj_bias,
+                              norm_topk_prob,
                               routed_scaling_factor)};
 }
 
@@ -115,7 +146,6 @@ std::vector<paddle::DataType> MoeExpertReduceInferDtype(
   return {ffn_out_dtype};
 }
 
-
 /**
  * @brief Mixture of Experts (MoE) Expert Reduce Operator
  *
@@ -131,9 +161,8 @@ std::vector<paddle::DataType> MoeExpertReduceInferDtype(
  *   - top_k_weight: Routing weights for top-k experts per token
  *                  Shape: [total_tokens, moe_topk]
  *                  dtype: float32
- *   - permute_indices_per_token: Indices mapping for reconstructing original order
- *                               Shape: [moe_topk, total_tokens]
- *                               dtype: int32
+ *   - permute_indices_per_token: Indices mapping for reconstructing original
+ * order Shape: [moe_topk, total_tokens] dtype: int32
  *   - top_k_indices: Indices of selected top-k experts for each token
  *                   Shape: [total_tokens, moe_topk]
  *                   dtype: int32
@@ -157,8 +186,11 @@ std::vector<paddle::DataType> MoeExpertReduceInferDtype(
  * - For optimal performance, hidden_size should be a multiple of 128
  */
 PD_BUILD_STATIC_OP(moe_expert_reduce)
-    .Inputs({"ffn_out", "top_k_weight", "permute_indices_per_token",
-             "top_k_indices", paddle::Optional("down_proj_bias")})
+    .Inputs({"ffn_out",
+             "top_k_weight",
+             "permute_indices_per_token",
+             "top_k_indices",
+             paddle::Optional("down_proj_bias")})
     .Outputs({"output"})
     .Attrs({"norm_topk_prob:bool", "routed_scaling_factor:float"})
     .SetKernelFn(PD_KERNEL(MoeExpertReduce))
diff --git a/custom_ops/gpu_ops/moe/moe_redundant_topk_select.cu b/custom_ops/gpu_ops/moe/moe_redundant_topk_select.cu
index 2d87c6bae5..1a6e64f0c0 100644
--- a/custom_ops/gpu_ops/moe/moe_redundant_topk_select.cu
+++ b/custom_ops/gpu_ops/moe/moe_redundant_topk_select.cu
@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-
 // Ignore CUTLASS warnings about type punning
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wstrict-aliasing"
@@ -27,35 +26,42 @@
 using namespace phi;
 
 template <typename T>
-void moe_redundant_topk_select_kernel(const T* input,
-                            const T* bias,
-                            T* output,
-                            T* softmax,
-                            const int* expert_id_to_ep_rank_array,
-                            const int* expert_in_rank_num_list,
-                            int* tokens_per_expert_stats_list,
-                            int64_t* indices,
-                            int64_t* indices_tmp,
-                            int* source_row,
-                            T* softmax_max_prob,
-                            const int64_t num_rows,
-                            const int64_t num_experts,
-                            const int64_t k,
-                            const int redundant_ep_rank_num_plus_one,
-                            cudaStream_t stream,
-                            const bool apply_norm_weight = false,
-                            const bool enable_softmax_top_k_fused = false
-                          ) {
+void moe_redundant_topk_select_kernel(
+    const T* input,
+    const T* bias,
+    T* output,
+    T* softmax,
+    const int* expert_id_to_ep_rank_array,
+    const int* expert_in_rank_num_list,
+    int* tokens_per_expert_stats_list,
+    int64_t* indices,
+    int64_t* indices_tmp,
+    int* source_row,
+    T* softmax_max_prob,
+    const int64_t num_rows,
+    const int64_t num_experts,
+    const int64_t k,
+    const int redundant_ep_rank_num_plus_one,
+    cudaStream_t stream,
+    const bool apply_norm_weight = false,
+    const bool enable_softmax_top_k_fused = false) {
   static constexpr int WARPS_PER_TB = 4;
 
-  #define LAUNCH_TOPK_GATING_SOFTMAX_HELPER(N)                                   \
-  case N: {                                                                    \
-    topk_gating_softmax_launcher_helper<T, N, WARPS_PER_TB>(                   \
-        input, bias, output, indices, source_row, num_rows, num_experts, k, stream); \
-    break;                                                                     \
+#define LAUNCH_TOPK_GATING_SOFTMAX_HELPER(N)                             \
+  case N: {                                                              \
+    topk_gating_softmax_launcher_helper<T, N, WARPS_PER_TB>(input,       \
+                                                            bias,        \
+                                                            output,      \
+                                                            indices,     \
+                                                            source_row,  \
+                                                            num_rows,    \
+                                                            num_experts, \
+                                                            k,           \
+                                                            stream);     \
+    break;                                                               \
   }
   int64_t tem_num_experts = num_experts;
-  if(bias != nullptr || apply_norm_weight)  tem_num_experts = 0;
+  if (bias != nullptr || apply_norm_weight) tem_num_experts = 0;
   switch (tem_num_experts) {
     LAUNCH_TOPK_GATING_SOFTMAX_HELPER(2)
     LAUNCH_TOPK_GATING_SOFTMAX_HELPER(4)
@@ -70,60 +76,60 @@ void moe_redundant_topk_select_kernel(const T* input,
       static constexpr int TPB = 256;
       const auto config_topk = Get1DBlocksAnd2DGridsMoe(num_rows);
       if (!enable_softmax_top_k_fused) {
-          moe_softmax<T, TPB><<<config_topk.block_per_grid, TPB, 0, stream>>>(
-              input, softmax, num_experts, num_rows);
-          if (apply_norm_weight) {
-            moe_redundant_top_k_normed<T, TPB>
-            <<<config_topk.block_per_grid, TPB, k * sizeof(T), stream>>>(softmax,
-                                                             bias,
-                                                             expert_id_to_ep_rank_array,
-                                                             expert_in_rank_num_list,
-                                                             tokens_per_expert_stats_list,
-                                                             output,
-                                                             indices,
-                                                             indices_tmp,
-                                                             source_row,
-                                                             num_experts,
-                                                             k,
-                                                             num_rows,
-                                                             redundant_ep_rank_num_plus_one);
-          } else {
-            moe_top_k<T, TPB>
-                <<<config_topk.block_per_grid, TPB, 0, stream>>>(softmax,
-                                                                  bias,
-                                                                  output,
-                                                                  indices,
-                                                                  source_row,
-                                                                  num_experts,
-                                                                  k,
-                                                                  num_rows);
-          }
+        moe_softmax<T, TPB><<<config_topk.block_per_grid, TPB, 0, stream>>>(
+            input, softmax, num_experts, num_rows);
+        if (apply_norm_weight) {
+          moe_redundant_top_k_normed<T, TPB>
+              <<<config_topk.block_per_grid, TPB, k * sizeof(T), stream>>>(
+                  softmax,
+                  bias,
+                  expert_id_to_ep_rank_array,
+                  expert_in_rank_num_list,
+                  tokens_per_expert_stats_list,
+                  output,
+                  indices,
+                  indices_tmp,
+                  source_row,
+                  num_experts,
+                  k,
+                  num_rows,
+                  redundant_ep_rank_num_plus_one);
+        } else {
+          moe_top_k<T, TPB>
+              <<<config_topk.block_per_grid, TPB, 0, stream>>>(softmax,
+                                                               bias,
+                                                               output,
+                                                               indices,
+                                                               source_row,
+                                                               num_experts,
+                                                               k,
+                                                               num_rows);
+        }
+      } else {
+        assert(k <= TPB);
+        if (apply_norm_weight) {
+          moe_softmax_top_k_fused<T, TPB, true>
+              <<<config_topk.block_per_grid, TPB, k * sizeof(T), stream>>>(
+                  input,
+                  bias,
+                  output,
+                  indices,
+                  source_row,
+                  num_experts,
+                  k,
+                  num_rows);
+        } else {
+          moe_softmax_top_k_fused<T, TPB, false>
+              <<<config_topk.block_per_grid, TPB, 0, stream>>>(input,
+                                                               bias,
+                                                               output,
+                                                               indices,
+                                                               source_row,
+                                                               num_experts,
+                                                               k,
+                                                               num_rows);
+        }
       }
-      else {
-          assert(k<=TPB);
-          if (apply_norm_weight) {
-            moe_softmax_top_k_fused<T, TPB, true>
-                <<<config_topk.block_per_grid, TPB, k * sizeof(T), stream>>>(input,
-                                                                 bias,
-                                                                 output,
-                                                                 indices,
-                                                                 source_row,
-                                                                 num_experts,
-                                                                 k,
-                                                                 num_rows);
-          } else {
-            moe_softmax_top_k_fused<T, TPB, false>
-                <<<config_topk.block_per_grid, TPB, 0, stream>>>(input,
-                                                                  bias,
-                                                                  output,
-                                                                  indices,
-                                                                  source_row,
-                                                                  num_experts,
-                                                                  k,
-                                                                  num_rows);
-          }
-      }
-
     }
   }
 }
@@ -195,24 +201,25 @@ std::vector<paddle::Tensor> MoERedundantTopKSelectKernel(
     softmax_out_ = nullptr;
   }
 
-  moe_redundant_topk_select_kernel<float>(gating_logits.data<float>(),
-                                bias ? bias.get().data<float>() : nullptr,
-                                topk_weights.data<float>(),
-                                softmax_out_,
-                                expert_id_to_ep_rank_array.data<int>(),
-                                expert_in_rank_num_list.data<int>(),
-                                tokens_per_expert_stats_list.data<int>(),
-                                topk_ids_data,
-                                topk_ids_tmp_data,
-                                source_rows_,
-                                softmax_max_prob,
-                                num_rows,
-                                expert_num,
-                                moe_topk,
-                                redundant_ep_rank_num_plus_one,
-                                stream,
-                                apply_norm_weight,
-                                enable_softmax_top_k_fused);
+  moe_redundant_topk_select_kernel<float>(
+      gating_logits.data<float>(),
+      bias ? bias.get().data<float>() : nullptr,
+      topk_weights.data<float>(),
+      softmax_out_,
+      expert_id_to_ep_rank_array.data<int>(),
+      expert_in_rank_num_list.data<int>(),
+      tokens_per_expert_stats_list.data<int>(),
+      topk_ids_data,
+      topk_ids_tmp_data,
+      source_rows_,
+      softmax_max_prob,
+      num_rows,
+      expert_num,
+      moe_topk,
+      redundant_ep_rank_num_plus_one,
+      stream,
+      apply_norm_weight,
+      enable_softmax_top_k_fused);
   return {topk_ids, topk_weights};
 }
 
@@ -235,8 +242,7 @@ std::vector<std::vector<int64_t>> MoERedundantTopKSelectKernelInferShape(
   }
   const int num_rows = token_rows;
 
-  return {{num_rows, moe_topk},
-          {num_rows, moe_topk}};
+  return {{num_rows, moe_topk}, {num_rows, moe_topk}};
 }
 
 std::vector<paddle::DataType> MoERedundantTopKSelectKernelInferDtype(
@@ -249,18 +255,22 @@ std::vector<paddle::DataType> MoERedundantTopKSelectKernelInferDtype(
     const bool apply_norm_weight,
     const bool enable_softmax_top_k_fused,
     const int redundant_ep_rank_num_plus_one) {
-  return {paddle::DataType::INT64,
-          paddle::DataType::FLOAT32};
+  return {paddle::DataType::INT64, paddle::DataType::FLOAT32};
 }
 
-
 PD_BUILD_STATIC_OP(moe_redundant_topk_select)
-    .Inputs({"gating_logits", "expert_id_to_ep_rank_array", "expert_in_rank_num_list", "tokens_per_expert_stats_list", paddle::Optional("bias")})
-    .Outputs({"topk_ids",
-              "topk_weights",
-              "tokens_per_expert_stats_list_out"})
-    .Attrs({"moe_topk:int", "apply_norm_weight:bool", "enable_softmax_top_k_fused:bool", "redundant_ep_rank_num_plus_one:int"})
-    .SetInplaceMap({{"tokens_per_expert_stats_list", "tokens_per_expert_stats_list_out"}})
+    .Inputs({"gating_logits",
+             "expert_id_to_ep_rank_array",
+             "expert_in_rank_num_list",
+             "tokens_per_expert_stats_list",
+             paddle::Optional("bias")})
+    .Outputs({"topk_ids", "topk_weights", "tokens_per_expert_stats_list_out"})
+    .Attrs({"moe_topk:int",
+            "apply_norm_weight:bool",
+            "enable_softmax_top_k_fused:bool",
+            "redundant_ep_rank_num_plus_one:int"})
+    .SetInplaceMap({{"tokens_per_expert_stats_list",
+                     "tokens_per_expert_stats_list_out"}})
     .SetKernelFn(PD_KERNEL(MoERedundantTopKSelectKernel))
     .SetInferShapeFn(PD_INFER_SHAPE(MoERedundantTopKSelectKernelInferShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(MoERedundantTopKSelectKernelInferDtype));
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_gemm.cu b/custom_ops/gpu_ops/moe/moe_wna16_marlin_gemm.cu
index e4c9082513..8b83d44df2 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_gemm.cu
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_gemm.cu
@@ -20,7 +20,7 @@
  */
 
 #ifndef MARLIN_NAMESPACE_NAME
-  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
 #endif
 #include "paddle/phi/core/enforce.h"
 #include "paddle/phi/api/include/api.h"
@@ -50,26 +50,42 @@ using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);
 
 template <int moe_block_size>
 __global__ void permute_cols_kernel(
-    int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr,
+    int4 const* __restrict__ a_int4_ptr,
+    int const* __restrict__ perm_int_ptr,
     int4* __restrict__ out_int4_ptr,
     const int32_t* __restrict__ sorted_token_ids_ptr,
     const int32_t* __restrict__ expert_ids_ptr,
-    const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m,
-    int size_k, int top_k) {};
+    const int32_t* __restrict__ num_tokens_past_padded_ptr,
+    int size_m,
+    int size_k,
+    int top_k){};
 
 }  // namespace marlin
 
 MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
-    MARLIN_NAMESPACE_NAME::Tensor& a, paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& c_or_none,
-    MARLIN_NAMESPACE_NAME::Tensor& b_q_weight, MARLIN_NAMESPACE_NAME::Tensor& b_scales,
+    MARLIN_NAMESPACE_NAME::Tensor& a,
+    paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& c_or_none,
+    MARLIN_NAMESPACE_NAME::Tensor& b_q_weight,
+    MARLIN_NAMESPACE_NAME::Tensor& b_scales,
     paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& b_zeros_or_none,
     paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& g_idx_or_none,
-    paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& perm_or_none, MARLIN_NAMESPACE_NAME::Tensor& workspace,
-    MARLIN_NAMESPACE_NAME::Tensor& sorted_token_ids, MARLIN_NAMESPACE_NAME::Tensor& expert_ids,
-    MARLIN_NAMESPACE_NAME::Tensor& num_tokens_past_padded, MARLIN_NAMESPACE_NAME::Tensor& topk_weights,
-    int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
-    MARLIN_NAMESPACE_NAME::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
-    int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
+    paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& perm_or_none,
+    MARLIN_NAMESPACE_NAME::Tensor& workspace,
+    MARLIN_NAMESPACE_NAME::Tensor& sorted_token_ids,
+    MARLIN_NAMESPACE_NAME::Tensor& expert_ids,
+    MARLIN_NAMESPACE_NAME::Tensor& num_tokens_past_padded,
+    MARLIN_NAMESPACE_NAME::Tensor& topk_weights,
+    int64_t moe_block_size,
+    int64_t top_k,
+    bool mul_topk_weights,
+    bool is_ep,
+    MARLIN_NAMESPACE_NAME::ScalarTypeId const& b_q_type_id,
+    int64_t size_m,
+    int64_t size_n,
+    int64_t size_k,
+    bool is_k_full,
+    bool use_atomic_add,
+    bool use_fp32_reduce,
     bool is_zp_float) {
   // TORCH_CHECK_NOT_IMPLEMENTED(false,
   //                             "marlin_gemm(..) requires CUDA_ARCH >= 8.0");
@@ -82,12 +98,15 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
 // on the given "perm" indices.
 template <int moe_block_size>
 __global__ void permute_cols_kernel(
-    int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr,
+    int4 const* __restrict__ a_int4_ptr,
+    int const* __restrict__ perm_int_ptr,
     int4* __restrict__ out_int4_ptr,
     const int32_t* __restrict__ sorted_token_ids_ptr,
     const int32_t* __restrict__ expert_ids_ptr,
-    const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m,
-    int size_k, int top_k) {
+    const int32_t* __restrict__ num_tokens_past_padded_ptr,
+    int size_m,
+    int size_k,
+    int top_k) {
   int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
   int num_moe_blocks = div_ceil(num_tokens_past_padded, moe_block_size);
   int32_t block_sorted_ids[moe_block_size];
@@ -181,9 +200,14 @@ typedef struct {
   thread_config_t tb_cfg;
 } exec_config_t;
 
-int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
-                          int prob_n, int prob_k, int num_bits, int group_size,
-                          bool has_act_order, bool is_k_full) {
+int get_scales_cache_size(thread_config_t const& th_config,
+                          int prob_m,
+                          int prob_n,
+                          int prob_k,
+                          int num_bits,
+                          int group_size,
+                          bool has_act_order,
+                          bool is_k_full) {
   bool cache_scales_chunk = has_act_order && !is_k_full;
 
   int tb_n = th_config.thread_n;
@@ -201,8 +225,9 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
 
   if (cache_scales_chunk) {
     int load_groups =
-        tb_groups * pipe_stages * 2;     // Chunk size is 2x pipeline over dim K
-    load_groups = std::max(load_groups, 32);  // We load at least 32 scale groups
+        tb_groups * pipe_stages * 2;  // Chunk size is 2x pipeline over dim K
+    load_groups =
+        std::max(load_groups, 32);  // We load at least 32 scale groups
     return load_groups * tb_n * 2;
   } else {
     int tb_scales = tb_groups * tb_n * 2;
@@ -211,10 +236,17 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
   }
 }
 
-int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
-                          int thread_m_blocks, int prob_m, int prob_n,
-                          int prob_k, int num_bits, int group_size,
-                          bool has_act_order, bool is_k_full, int has_zp,
+int get_kernel_cache_size(thread_config_t const& th_config,
+                          bool m_block_size_8,
+                          int thread_m_blocks,
+                          int prob_m,
+                          int prob_n,
+                          int prob_k,
+                          int num_bits,
+                          int group_size,
+                          bool has_act_order,
+                          bool is_k_full,
+                          int has_zp,
                           int is_zp_float) {
   int pack_factor = 32 / num_bits;
 
@@ -229,9 +261,14 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
   int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
   int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
   int sh_red_size = tb_m * (tb_n + 8) * 2;
-  int sh_s_size =
-      get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
-                            group_size, has_act_order, is_k_full);
+  int sh_s_size = get_scales_cache_size(th_config,
+                                        prob_m,
+                                        prob_n,
+                                        prob_k,
+                                        num_bits,
+                                        group_size,
+                                        has_act_order,
+                                        is_k_full);
   int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0;
   int sh_zp_size = 0;
   if (has_zp) {
@@ -249,10 +286,18 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
   return total_size;
 }
 
-bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
-                     int thread_m_blocks, int prob_m, int prob_n, int prob_k,
-                     int num_bits, int group_size, bool has_act_order,
-                     bool is_k_full, int has_zp, int is_zp_float,
+bool is_valid_config(thread_config_t const& th_config,
+                     bool m_block_size_8,
+                     int thread_m_blocks,
+                     int prob_m,
+                     int prob_n,
+                     int prob_k,
+                     int num_bits,
+                     int group_size,
+                     bool has_act_order,
+                     bool is_k_full,
+                     int has_zp,
+                     int is_zp_float,
                      int max_shared_mem) {
   // Sanity
   if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
@@ -276,137 +321,162 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
   }
 
   // Check that pipeline fits into cache
-  int cache_size = get_kernel_cache_size(
-      th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
-      num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);
+  int cache_size = get_kernel_cache_size(th_config,
+                                         m_block_size_8,
+                                         thread_m_blocks,
+                                         prob_m,
+                                         prob_n,
+                                         prob_k,
+                                         num_bits,
+                                         group_size,
+                                         has_act_order,
+                                         is_k_full,
+                                         has_zp,
+                                         is_zp_float);
   return cache_size <= max_shared_mem;
 }
 
-  #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
-                  M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT)    \
-    else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS &&       \
-             thread_n_blocks == THREAD_N_BLOCKS &&                           \
-             thread_k_blocks == THREAD_K_BLOCKS &&                           \
-             m_block_size_8 == M_BLOCK_SIZE_8 &&                             \
-             group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS &&   \
-             is_zp_float == IS_ZP_FLOAT) {                                   \
-      kernel = Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS,   \
-                      THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8,      \
-                      pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>;               \
-    }
+#define _GET_IF(W_TYPE,                                                        \
+                THREAD_M_BLOCKS,                                               \
+                THREAD_N_BLOCKS,                                               \
+                THREAD_K_BLOCKS,                                               \
+                M_BLOCK_SIZE_8,                                                \
+                GROUP_BLOCKS,                                                  \
+                NUM_THREADS,                                                   \
+                IS_ZP_FLOAT)                                                   \
+  else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS &&           \
+           thread_n_blocks == THREAD_N_BLOCKS &&                               \
+           thread_k_blocks == THREAD_K_BLOCKS &&                               \
+           m_block_size_8 == M_BLOCK_SIZE_8 && group_blocks == GROUP_BLOCKS && \
+           num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) {         \
+    kernel = Marlin<scalar_t,                                                  \
+                    W_TYPE.id(),                                               \
+                    NUM_THREADS,                                               \
+                    THREAD_M_BLOCKS,                                           \
+                    THREAD_N_BLOCKS,                                           \
+                    THREAD_K_BLOCKS,                                           \
+                    M_BLOCK_SIZE_8,                                            \
+                    pipe_stages,                                               \
+                    GROUP_BLOCKS,                                              \
+                    IS_ZP_FLOAT>;                                              \
+  }
 
-  // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
-  //         this is the most common cases
-  // BIGGROUP: cases for big group size (group_blocks in [-1, 8])
-  // FZP: cases for float-zero-point (is_zp_float = true)
-  // ACT: cases for act order case (group_blocks == 0)
-  // FP4: cases for nvfp4(e2m1) (group_blocks == 1)
-  #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false)   \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false)   \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false)   \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+// COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
+//         this is the most common cases
+// BIGGROUP: cases for big group size (group_blocks in [-1, 8])
+// FZP: cases for float-zero-point (is_zp_float = true)
+// ACT: cases for act order case (group_blocks == 0)
+// FP4: cases for nvfp4(e2m1) (group_blocks == 1)
+#define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false)   \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false)   \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false)   \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
 
-  #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
-                                                                          \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
-                                                                          \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+#define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+                                                                        \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+                                                                        \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
 
-  #define COMMON_GET_IF(W_TYPE)            \
-    COMMON_GET_IF_M1(W_TYPE, 8, 8, 256)    \
-    COMMON_GET_IF_M1(W_TYPE, 8, 4, 128)    \
-    COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \
-    COMMON_GET_IF_M234(W_TYPE, 8, 4, 128)
+#define COMMON_GET_IF(W_TYPE)            \
+  COMMON_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+  COMMON_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+  COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \
+  COMMON_GET_IF_M234(W_TYPE, 8, 4, 128)
 
-  #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false)   \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+#define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false)   \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
 
-  #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)   \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
-                                                                          \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+#define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)   \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+                                                                        \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
 
-  #define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)        \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
+#define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)        \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
 
-  #define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
+#define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
 
-  #define FP4_GET_IF(W_TYPE)            \
-    FP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
-    FP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
-    FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
-    FP4_GET_IF_M234(W_TYPE, 8, 4, 128)
+#define FP4_GET_IF(W_TYPE)            \
+  FP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+  FP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+  FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
+  FP4_GET_IF_M234(W_TYPE, 8, 4, 128)
 
-  #define BIGGROUP_GET_IF(W_TYPE)            \
-    BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256)    \
-    BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128)    \
-    BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
-    BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)
+#define BIGGROUP_GET_IF(W_TYPE)            \
+  BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+  BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+  BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
+  BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)
 
-  // We currently have 4-bit models only with group_blocks == 4
-  #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
+// We currently have 4-bit models only with group_blocks == 4
+#define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
 
-  #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
+#define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
 
-  #define FZP_GET_IF(W_TYPE)            \
-    FZP_GET_IF_M1(W_TYPE, 8, 8, 256)    \
-    FZP_GET_IF_M1(W_TYPE, 8, 4, 128)    \
-    FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \
-    FZP_GET_IF_M234(W_TYPE, 8, 4, 128)
+#define FZP_GET_IF(W_TYPE)            \
+  FZP_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+  FZP_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+  FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \
+  FZP_GET_IF_M234(W_TYPE, 8, 4, 128)
 
-  // We currently have 4-bit models only with group_blocks == 4
-  #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)        \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
+// We currently have 4-bit models only with group_blocks == 4
+#define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)        \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
 
-  #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
-    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
-    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
+#define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+  _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
+  _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
 
-  #define ACT_GET_IF(W_TYPE)            \
-    ACT_GET_IF_M1(W_TYPE, 8, 8, 256)    \
-    ACT_GET_IF_M1(W_TYPE, 8, 4, 128)    \
-    ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
-    ACT_GET_IF_M234(W_TYPE, 8, 4, 128)
+#define ACT_GET_IF(W_TYPE)            \
+  ACT_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+  ACT_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+  ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
+  ACT_GET_IF_M234(W_TYPE, 8, 4, 128)
 
 template <typename scalar_t>
 MarlinFuncPtr get_marlin_kernel(const MARLIN_NAMESPACE_NAME::ScalarType q_type,
-                                int thread_m_blocks, int thread_n_blocks,
-                                int thread_k_blocks, bool m_block_size_8,
-                                bool has_act_order, bool has_zp,
-                                int group_blocks, int num_threads,
+                                int thread_m_blocks,
+                                int thread_n_blocks,
+                                int thread_k_blocks,
+                                bool m_block_size_8,
+                                bool has_act_order,
+                                bool has_zp,
+                                int group_blocks,
+                                int num_threads,
                                 bool is_zp_float) {
   int num_bits = q_type.size_bits();
   auto kernel = MarlinDefault;
@@ -428,12 +498,20 @@ MarlinFuncPtr get_marlin_kernel(const MARLIN_NAMESPACE_NAME::ScalarType q_type,
 }
 
 template <typename scalar_t>
-exec_config_t determine_exec_config(const MARLIN_NAMESPACE_NAME::ScalarType& q_type, int prob_m,
-                                    int prob_n, int prob_k, int thread_m_blocks,
-                                    bool m_block_size_8, int num_bits,
-                                    int group_size, bool has_act_order,
-                                    bool is_k_full, bool has_zp,
-                                    bool is_zp_float, int max_shared_mem) {
+exec_config_t determine_exec_config(
+    const MARLIN_NAMESPACE_NAME::ScalarType& q_type,
+    int prob_m,
+    int prob_n,
+    int prob_k,
+    int thread_m_blocks,
+    bool m_block_size_8,
+    int num_bits,
+    int group_size,
+    bool has_act_order,
+    bool is_k_full,
+    bool has_zp,
+    bool is_zp_float,
+    int max_shared_mem) {
   exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
   thread_config_t* thread_configs = thread_m_blocks > 1
                                         ? large_batch_thread_configs
@@ -448,25 +526,50 @@ exec_config_t determine_exec_config(const MARLIN_NAMESPACE_NAME::ScalarType& q_t
   for (int i = 0; i < thread_configs_size; i++) {
     thread_config_t th_config = thread_configs[i];
 
-    if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m,
-                         prob_n, prob_k, num_bits, group_size, has_act_order,
-                         is_k_full, has_zp, is_zp_float, max_shared_mem)) {
+    if (!is_valid_config(th_config,
+                         m_block_size_8,
+                         thread_m_blocks,
+                         prob_m,
+                         prob_n,
+                         prob_k,
+                         num_bits,
+                         group_size,
+                         has_act_order,
+                         is_k_full,
+                         has_zp,
+                         is_zp_float,
+                         max_shared_mem)) {
       continue;
     }
 
-    int cache_size = get_kernel_cache_size(
-        th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
-        num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);
+    int cache_size = get_kernel_cache_size(th_config,
+                                           m_block_size_8,
+                                           thread_m_blocks,
+                                           prob_m,
+                                           prob_n,
+                                           prob_k,
+                                           num_bits,
+                                           group_size,
+                                           has_act_order,
+                                           is_k_full,
+                                           has_zp,
+                                           is_zp_float);
 
     int group_blocks = 0;
     if (!has_act_order) {
       group_blocks = group_size == -1 ? -1 : (group_size / 16);
     }
 
-    auto kernel = get_marlin_kernel<scalar_t>(
-        q_type, thread_m_blocks, th_config.thread_n / 16,
-        th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp,
-        group_blocks, th_config.num_threads, is_zp_float);
+    auto kernel = get_marlin_kernel<scalar_t>(q_type,
+                                              thread_m_blocks,
+                                              th_config.thread_n / 16,
+                                              th_config.thread_k / 16,
+                                              m_block_size_8,
+                                              has_act_order,
+                                              has_zp,
+                                              group_blocks,
+                                              th_config.num_threads,
+                                              is_zp_float);
 
     if (kernel == MarlinDefault) continue;
 
@@ -478,7 +581,7 @@ exec_config_t determine_exec_config(const MARLIN_NAMESPACE_NAME::ScalarType& q_t
       cudaFuncGetAttributes(&attr, reinterpret_cast<const void*>(kernel));
       int reg_size = std::max(attr.numRegs, 1) * th_config.num_threads * 4;
       int allow_count = std::min(device_max_reg_size / reg_size,
-                            max_shared_mem / (cache_size + 1024));
+                                 max_shared_mem / (cache_size + 1024));
       allow_count = std::max(std::min(allow_count, 4), 1);
       if (allow_count > count) {
         count = allow_count;
@@ -491,43 +594,80 @@ exec_config_t determine_exec_config(const MARLIN_NAMESPACE_NAME::ScalarType& q_t
 }
 
 template <typename scalar_t>
-void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
-               void* s2, void* zp, void* g_idx, void* perm, void* a_tmp,
-               void* sorted_token_ids, void* expert_ids,
-               void* num_tokens_past_padded, void* topk_weights,
-               int moe_block_size, int top_k, bool mul_topk_weights, bool is_ep,
-               int prob_m, int prob_n, int prob_k, void* workspace,
-               MARLIN_NAMESPACE_NAME::ScalarType const& q_type, bool has_act_order,
-               bool is_k_full, bool has_zp, int num_groups, int group_size,
-               int dev, cudaStream_t stream, int thread_k, int thread_n,
-               int sms, bool use_atomic_add, bool use_fp32_reduce,
+void marlin_mm(const void* A,
+               const void* B,
+               void* C,
+               void* C_tmp,
+               void* s,
+               void* s2,
+               void* zp,
+               void* g_idx,
+               void* perm,
+               void* a_tmp,
+               void* sorted_token_ids,
+               void* expert_ids,
+               void* num_tokens_past_padded,
+               void* topk_weights,
+               int moe_block_size,
+               int top_k,
+               bool mul_topk_weights,
+               bool is_ep,
+               int prob_m,
+               int prob_n,
+               int prob_k,
+               void* workspace,
+               MARLIN_NAMESPACE_NAME::ScalarType const& q_type,
+               bool has_act_order,
+               bool is_k_full,
+               bool has_zp,
+               int num_groups,
+               int group_size,
+               int dev,
+               cudaStream_t stream,
+               int thread_k,
+               int thread_n,
+               int sms,
+               bool use_atomic_add,
+               bool use_fp32_reduce,
                bool is_zp_float) {
   int thread_m_blocks = div_ceil(moe_block_size, 16);
   bool m_block_size_8 = moe_block_size == 8;
 
   if (has_zp) {
-    PADDLE_ENFORCE(
-        q_type == MARLIN_NAMESPACE_NAME::kU4 || q_type == MARLIN_NAMESPACE_NAME::kU8,
-        "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
+    PADDLE_ENFORCE(q_type == MARLIN_NAMESPACE_NAME::kU4 ||
+                       q_type == MARLIN_NAMESPACE_NAME::kU8,
+                   "q_type must be u4 or u8 when has_zp = True. Got = ",
+                   q_type.str());
   } else {
     PADDLE_ENFORCE(
-        q_type == MARLIN_NAMESPACE_NAME::kU4B8 || q_type == MARLIN_NAMESPACE_NAME::kU8B128 ||
-            q_type == MARLIN_NAMESPACE_NAME::kFE4M3fn || q_type == MARLIN_NAMESPACE_NAME::kFE2M1f,
+        q_type == MARLIN_NAMESPACE_NAME::kU4B8 ||
+            q_type == MARLIN_NAMESPACE_NAME::kU8B128 ||
+            q_type == MARLIN_NAMESPACE_NAME::kFE4M3fn ||
+            q_type == MARLIN_NAMESPACE_NAME::kFE2M1f,
         "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when "
         "has_zp = False. Got = ",
         q_type.str());
   }
 
-  PADDLE_ENFORCE(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
-              ", ", prob_n, ", ", prob_k, "]");
+  PADDLE_ENFORCE(prob_m > 0 && prob_n > 0 && prob_k > 0,
+                 "Invalid MNK = [",
+                 prob_m,
+                 ", ",
+                 prob_n,
+                 ", ",
+                 prob_k,
+                 "]");
 
   int group_blocks = 0;
   if (has_act_order) {
     if (is_k_full) {
       PADDLE_ENFORCE(group_size != -1, "group_size = ", group_size);
       group_blocks = group_size / 16;
-      PADDLE_ENFORCE(prob_k % group_blocks == 0, "prob_k = ", prob_k,
-                  " is not divisible by group_blocks = ", group_blocks);
+      PADDLE_ENFORCE(prob_k % group_blocks == 0,
+                     "prob_k = ",
+                     prob_k,
+                     " is not divisible by group_blocks = ",
+                     group_blocks);
     } else {
       PADDLE_ENFORCE(group_size == 0, "group_size = ", group_size);
       group_blocks = 0;
@@ -537,8 +677,11 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
       group_blocks = -1;
     } else {
       group_blocks = group_size / 16;
-      PADDLE_ENFORCE(prob_k % group_blocks == 0, "prob_k = ", prob_k,
-                  " is not divisible by group_blocks = ", group_blocks);
+      PADDLE_ENFORCE(prob_k % group_blocks == 0,
+                     "prob_k = ",
+                     prob_k,
+                     " is not divisible by group_blocks = ",
+                     group_blocks);
     }
   }
 
@@ -592,9 +735,11 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
   }
 
   int max_shared_mem = 0;
-  cudaDeviceGetAttribute(&max_shared_mem,
-                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
-  PADDLE_ENFORCE(max_shared_mem > 0, "max_shared_mem should > 0 ! max_shared_mem = ", max_shared_mem);
+  cudaDeviceGetAttribute(
+      &max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
+  PADDLE_ENFORCE(max_shared_mem > 0,
+                 "max_shared_mem should > 0 ! max_shared_mem = ",
+                 max_shared_mem);
 
   // Set thread config
   exec_config_t exec_cfg;
@@ -602,16 +747,31 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
   if (thread_k != -1 && thread_n != -1) {
     thread_tfg = thread_config_t{thread_k, thread_n, default_threads};
     exec_cfg = exec_config_t{1, thread_tfg};
-    PADDLE_ENFORCE(prob_n % thread_n == 0, "prob_n = ", prob_n,
-                " is not divisible by thread_n = ", thread_n);
-    PADDLE_ENFORCE(prob_k % thread_k == 0, "prob_k = ", prob_k,
-                " is not divisible by thread_k = ", thread_k);
+    PADDLE_ENFORCE(prob_n % thread_n == 0,
+                   "prob_n = ",
+                   prob_n,
+                   " is not divisible by thread_n = ",
+                   thread_n);
+    PADDLE_ENFORCE(prob_k % thread_k == 0,
+                   "prob_k = ",
+                   prob_k,
+                   " is not divisible by thread_k = ",
+                   thread_k);
   } else {
     // Auto config
-    exec_cfg = determine_exec_config<scalar_t>(
-        q_type, prob_m, prob_n, prob_k, thread_m_blocks, m_block_size_8,
-        num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
-        max_shared_mem);
+    exec_cfg = determine_exec_config<scalar_t>(q_type,
+                                               prob_m,
+                                               prob_n,
+                                               prob_k,
+                                               thread_m_blocks,
+                                               m_block_size_8,
+                                               num_bits,
+                                               group_size,
+                                               has_act_order,
+                                               is_k_full,
+                                               has_zp,
+                                               is_zp_float,
+                                               max_shared_mem);
     thread_tfg = exec_cfg.tb_cfg;
   }
 
@@ -625,34 +785,86 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
   int thread_k_blocks = thread_k / 16;
   int thread_n_blocks = thread_n / 16;
 
-  PADDLE_ENFORCE(
-      is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
-                      prob_n, prob_k, num_bits, group_size, has_act_order,
-                      is_k_full, has_zp, is_zp_float, max_shared_mem),
-      "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
-      ", thread_k = ", thread_tfg.thread_k,
-      ", thread_n = ", thread_tfg.thread_n,
-      ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
-      prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
-      ", group_size = ", group_size, ", has_act_order = ", has_act_order,
-      ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
-      ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);
+  PADDLE_ENFORCE(is_valid_config(thread_tfg,
+                                 m_block_size_8,
+                                 thread_m_blocks,
+                                 prob_m,
+                                 prob_n,
+                                 prob_k,
+                                 num_bits,
+                                 group_size,
+                                 has_act_order,
+                                 is_k_full,
+                                 has_zp,
+                                 is_zp_float,
+                                 max_shared_mem),
+                 "Invalid thread config: thread_m_blocks = ",
+                 thread_m_blocks,
+                 ", thread_k = ",
+                 thread_tfg.thread_k,
+                 ", thread_n = ",
+                 thread_tfg.thread_n,
+                 ", num_threads = ",
+                 thread_tfg.num_threads,
+                 " for MKN = [",
+                 prob_m,
+                 ", ",
+                 prob_k,
+                 ", ",
+                 prob_n,
+                 "] and num_bits = ",
+                 num_bits,
+                 ", group_size = ",
+                 group_size,
+                 ", has_act_order = ",
+                 has_act_order,
+                 ", is_k_full = ",
+                 is_k_full,
+                 ", has_zp = ",
+                 has_zp,
+                 ", is_zp_float = ",
+                 is_zp_float,
+                 ", max_shared_mem = ",
+                 max_shared_mem);
 
-  auto kernel = get_marlin_kernel<scalar_t>(
-      q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8,
-      has_act_order, has_zp, group_blocks, num_threads, is_zp_float);
+  auto kernel = get_marlin_kernel<scalar_t>(q_type,
+                                            thread_m_blocks,
+                                            thread_n_blocks,
+                                            thread_k_blocks,
+                                            m_block_size_8,
+                                            has_act_order,
+                                            has_zp,
+                                            group_blocks,
+                                            num_threads,
+                                            is_zp_float);
 
   if (kernel == MarlinDefault) {
-    PADDLE_ENFORCE(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
-                ", ", prob_k, "]", ", has_act_order = ", has_act_order,
-                ", num_groups = ", num_groups, ", group_size = ", group_size,
-                ", thread_m_blocks = ", thread_m_blocks,
-                ", thread_n_blocks = ", thread_n_blocks,
-                ", thread_k_blocks = ", thread_k_blocks,
-                ", num_bits = ", num_bits);
+    PADDLE_ENFORCE(false,
+                   "Unsupported shapes: MNK = [",
+                   prob_m,
+                   ", ",
+                   prob_n,
+                   ", ",
+                   prob_k,
+                   "]",
+                   ", has_act_order = ",
+                   has_act_order,
+                   ", num_groups = ",
+                   num_groups,
+                   ", group_size = ",
+                   group_size,
+                   ", thread_m_blocks = ",
+                   thread_m_blocks,
+                   ", thread_n_blocks = ",
+                   thread_n_blocks,
+                   ", thread_k_blocks = ",
+                   thread_k_blocks,
+                   ", num_bits = ",
+                   num_bits);
   }
 
-  cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel), cudaFuncAttributeMaxDynamicSharedMemorySize,
+  cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel),
+                       cudaFuncAttributeMaxDynamicSharedMemorySize,
                        max_shared_mem);
   // avoid ">>>" being formatted to "> > >"
   // clang-format off
@@ -688,50 +900,82 @@ ConvertOptionalPaddleTensorToDetailTensor(
 }
 
 MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
-    MARLIN_NAMESPACE_NAME::Tensor& a, paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& c_or_none,
-    MARLIN_NAMESPACE_NAME::Tensor& b_q_weight, MARLIN_NAMESPACE_NAME::Tensor& b_scales,
+    MARLIN_NAMESPACE_NAME::Tensor& a,
+    paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& c_or_none,
+    MARLIN_NAMESPACE_NAME::Tensor& b_q_weight,
+    MARLIN_NAMESPACE_NAME::Tensor& b_scales,
     paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& global_scale_or_none,
     paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& b_zeros_or_none,
     paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& g_idx_or_none,
-    paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& perm_or_none, MARLIN_NAMESPACE_NAME::Tensor& workspace,
-    MARLIN_NAMESPACE_NAME::Tensor& sorted_token_ids, MARLIN_NAMESPACE_NAME::Tensor& expert_ids,
-    MARLIN_NAMESPACE_NAME::Tensor& num_tokens_past_padded, MARLIN_NAMESPACE_NAME::Tensor& topk_weights,
-    int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
-    MARLIN_NAMESPACE_NAME::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
-    int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
+    paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> const& perm_or_none,
+    MARLIN_NAMESPACE_NAME::Tensor& workspace,
+    MARLIN_NAMESPACE_NAME::Tensor& sorted_token_ids,
+    MARLIN_NAMESPACE_NAME::Tensor& expert_ids,
+    MARLIN_NAMESPACE_NAME::Tensor& num_tokens_past_padded,
+    MARLIN_NAMESPACE_NAME::Tensor& topk_weights,
+    int64_t moe_block_size,
+    int64_t top_k,
+    bool mul_topk_weights,
+    bool is_ep,
+    MARLIN_NAMESPACE_NAME::ScalarTypeId const& b_q_type_id,
+    int64_t size_m,
+    int64_t size_n,
+    int64_t size_k,
+    bool is_k_full,
+    bool use_atomic_add,
+    bool use_fp32_reduce,
     bool is_zp_float) {
-  MARLIN_NAMESPACE_NAME::ScalarType const b_q_type = MARLIN_NAMESPACE_NAME::ScalarType::from_id(b_q_type_id);
+  MARLIN_NAMESPACE_NAME::ScalarType const b_q_type =
+      MARLIN_NAMESPACE_NAME::ScalarType::from_id(b_q_type_id);
   int pack_factor = 32 / b_q_type.size_bits();
 
   if (moe_block_size != 8) {
     PADDLE_ENFORCE(moe_block_size % 16 == 0,
-                "unsupported moe_block_size=", moe_block_size);
+                   "unsupported moe_block_size=",
+                   moe_block_size);
     PADDLE_ENFORCE(moe_block_size >= 16 && moe_block_size <= 64,
-                "unsupported moe_block_size=", moe_block_size);
+                   "unsupported moe_block_size=",
+                   moe_block_size);
   }
 
   // Verify A
-  PADDLE_ENFORCE(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0),
-              ", size_m = ", size_m);
-  PADDLE_ENFORCE(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1),
-              ", size_k = ", size_k);
+  PADDLE_ENFORCE(a.size(0) == size_m,
+                 "Shape mismatch: a.size(0) = ",
+                 a.size(0),
+                 ", size_m = ",
+                 size_m);
+  PADDLE_ENFORCE(a.size(1) == size_k,
+                 "Shape mismatch: a.size(1) = ",
+                 a.size(1),
+                 ", size_k = ",
+                 size_k);
 
   // Verify B
+  PADDLE_ENFORCE(size_k % MARLIN_NAMESPACE_NAME::tile_size == 0,
+                 "size_k = ",
+                 size_k,
+                 " is not divisible by tile_size = ",
+                 MARLIN_NAMESPACE_NAME::tile_size);
   PADDLE_ENFORCE(
-      size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k,
-      " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
-  PADDLE_ENFORCE((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1),
-              "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
-              ", size_k = ", size_k,
-              ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
-  PADDLE_ENFORCE(
-      b_q_weight.size(2) % MARLIN_NAMESPACE_NAME::tile_size == 0,
-      "b_q_weight.size(2) = ", b_q_weight.size(2),
-      " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
+      (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1),
+      "Shape mismatch: b_q_weight.size(1) = ",
+      b_q_weight.size(1),
+      ", size_k = ",
+      size_k,
+      ", tile_size = ",
+      MARLIN_NAMESPACE_NAME::tile_size);
+  PADDLE_ENFORCE(b_q_weight.size(2) % MARLIN_NAMESPACE_NAME::tile_size == 0,
+                 "b_q_weight.size(2) = ",
+                 b_q_weight.size(2),
+                 " is not divisible by tile_size = ",
+                 MARLIN_NAMESPACE_NAME::tile_size);
   int actual_size_n =
       (b_q_weight.size(2) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor;
-  PADDLE_ENFORCE(size_n == actual_size_n, "size_n = ", size_n,
-              ", actual_size_n = ", actual_size_n);
+  PADDLE_ENFORCE(size_n == actual_size_n,
+                 "size_n = ",
+                 size_n,
+                 ", actual_size_n = ",
+                 actual_size_n);
 
   // Verify device and strides
   PADDLE_ENFORCE(a.is_gpu(), "A is not on GPU");
@@ -766,12 +1010,18 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
     PADDLE_ENFORCE(c.is_gpu(), "c is not on GPU");
     PADDLE_ENFORCE(c.is_contiguous(), "c is not contiguous");
     PADDLE_ENFORCE(c.size(0) == size_m * top_k,
-                "Shape mismatch: c.size(0) = ", c.size(0),
-                ", size_m * topk = ", size_m * top_k);
-    PADDLE_ENFORCE(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1),
-                ", size_n = ", size_n);
+                   "Shape mismatch: c.size(0) = ",
+                   c.size(0),
+                   ", size_m * topk = ",
+                   size_m * top_k);
+    PADDLE_ENFORCE(c.size(1) == size_n,
+                   "Shape mismatch: c.size(1) = ",
+                   c.size(1),
+                   ", size_n = ",
+                   size_n);
   } else {
-    c = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({size_m * top_k, size_n}, a.dtype(), phi::GPUPlace(device_id)));
+    c = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
+        {size_m * top_k, size_n}, a.dtype(), phi::GPUPlace(device_id)));
   }
 
   // Alloc C tmp buffer that is going to be used for the global reduce
@@ -782,9 +1032,13 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
         (long)size_n * sorted_token_ids.size(0),
         (long)sms * 4 * moe_block_size * MARLIN_NAMESPACE_NAME::max_thread_n);
     if (moe_block_size == 8) max_c_tmp_size *= 2;
-    c_tmp = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({max_c_tmp_size}, MARLIN_NAMESPACE_NAME::kFloat32, phi::GPUPlace(device_id)));
+    c_tmp = ConvertPaddleTensorToDetailTensor(
+        paddle::experimental::empty({max_c_tmp_size},
+                                    MARLIN_NAMESPACE_NAME::kFloat32,
+                                    phi::GPUPlace(device_id)));
   } else {
-    c_tmp = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({0}, MARLIN_NAMESPACE_NAME::kFloat32, phi::GPUPlace(device_id)));
+    c_tmp = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
+        {0}, MARLIN_NAMESPACE_NAME::kFloat32, phi::GPUPlace(device_id)));
   }
 
   // Detect groupsize and act_order
@@ -792,8 +1046,11 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
 
   int rank = b_scales.dim();
   PADDLE_ENFORCE(rank == 3, "b_scales rank = ", rank, " is not 3");
-  PADDLE_ENFORCE(b_scales.size(2) == size_n, "b_scales dim 2 = ", b_scales.size(2),
-              " is not size_n = ", size_n);
+  PADDLE_ENFORCE(b_scales.size(2) == size_n,
+                 "b_scales dim 2 = ",
+                 b_scales.size(2),
+                 " is not size_n = ",
+                 size_n);
   int num_groups = b_scales.size(1);
 
   bool has_act_order = false;
@@ -809,40 +1066,52 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
 
     // Verify g_idx and perm
     PADDLE_ENFORCE((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
-                    (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
-                "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
-                " and perm.size(-1) = ", perm.size(-1),
-                ", where size_k = ", size_k);
+                       (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
+                   "Unexpected g_idx.size(-1) = ",
+                   g_idx.size(-1),
+                   " and perm.size(-1) = ",
+                   perm.size(-1),
+                   ", where size_k = ",
+                   size_k);
 
     has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
   } else {
-    g_idx = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
+    g_idx = ConvertPaddleTensorToDetailTensor(
+        paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
 
-    perm = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
-
-    a_tmp = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
+    perm = ConvertPaddleTensorToDetailTensor(
+        paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
 
+    a_tmp = ConvertPaddleTensorToDetailTensor(
+        paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
   }
 
   if (has_act_order) {
-    a_tmp = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({size_m * top_k, size_k}, a.dtype(), phi::GPUPlace(device_id)));
+    a_tmp = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
+        {size_m * top_k, size_k}, a.dtype(), phi::GPUPlace(device_id)));
 
     if (is_k_full) {
       PADDLE_ENFORCE(num_groups > 1, "For act_order, num_groups must be > 1");
-      PADDLE_ENFORCE(size_k % num_groups == 0, "size_k = ", size_k,
-                  ", is not divisible by num_groups = ", num_groups);
+      PADDLE_ENFORCE(size_k % num_groups == 0,
+                     "size_k = ",
+                     size_k,
+                     ", is not divisible by num_groups = ",
+                     num_groups);
       group_size = size_k / num_groups;
     } else {
       group_size = 0;
     }
 
   } else {
-    a_tmp = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
+    a_tmp = ConvertPaddleTensorToDetailTensor(
+        paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
 
     if (num_groups > 1) {
-      PADDLE_ENFORCE(
-          size_k % num_groups == 0, "size_k = ", size_k,
-          ", is not divisible by b_scales.size(1) = ", b_scales.size(1));
+      PADDLE_ENFORCE(size_k % num_groups == 0,
+                     "size_k = ",
+                     size_k,
+                     ", is not divisible by b_scales.size(1) = ",
+                     b_scales.size(1));
       group_size = size_k / num_groups;
     } else {
       group_size = -1;
@@ -853,12 +1122,14 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
   if (global_scale_or_none) {
     global_scale = global_scale_or_none.get();
     PADDLE_ENFORCE(b_q_type == MARLIN_NAMESPACE_NAME::kFE2M1f,
-                "global_scale can only be used for float4_e2m1f.");
+                   "global_scale can only be used for float4_e2m1f.");
   } else {
-    global_scale = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
+    global_scale = ConvertPaddleTensorToDetailTensor(
+        paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
 
-    PADDLE_ENFORCE(!(b_q_type == MARLIN_NAMESPACE_NAME::kFE2M1f),
-                "the global_scale parameter must be passed for float4_e2m1f.");
+    PADDLE_ENFORCE(
+        !(b_q_type == MARLIN_NAMESPACE_NAME::kFE2M1f),
+        "the global_scale parameter must be passed for float4_e2m1f.");
   }
 
   MARLIN_NAMESPACE_NAME::Tensor b_zeros;
@@ -867,27 +1138,32 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
     PADDLE_ENFORCE(b_zeros.is_gpu(), "b_zeros is not on GPU");
     PADDLE_ENFORCE(b_zeros.is_contiguous(), "b_zeros is not contiguous");
   } else {
-    b_zeros = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
+    b_zeros = ConvertPaddleTensorToDetailTensor(
+        paddle::experimental::empty({0}, a.dtype(), phi::GPUPlace(device_id)));
   }
   bool has_zp = b_zeros.numel() > 0;
 
   if (has_zp) {
-    PADDLE_ENFORCE(
-        b_q_type == MARLIN_NAMESPACE_NAME::kU4 || b_q_type == MARLIN_NAMESPACE_NAME::kU8,
-        "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
+    PADDLE_ENFORCE(b_q_type == MARLIN_NAMESPACE_NAME::kU4 ||
+                       b_q_type == MARLIN_NAMESPACE_NAME::kU8,
+                   "b_q_type must be u4 or u8 when has_zp = True. Got = ",
+                   b_q_type.str());
   } else {
-    PADDLE_ENFORCE(b_q_type == MARLIN_NAMESPACE_NAME::kU4B8 || b_q_type == MARLIN_NAMESPACE_NAME::kU8B128 ||
-                    b_q_type == MARLIN_NAMESPACE_NAME::kFE4M3fn || b_q_type == MARLIN_NAMESPACE_NAME::kFE2M1f,
-                "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or "
-                "float4_e2m1f when "
-                "has_zp = False. Got = ",
-                b_q_type.str());
+    PADDLE_ENFORCE(b_q_type == MARLIN_NAMESPACE_NAME::kU4B8 ||
+                       b_q_type == MARLIN_NAMESPACE_NAME::kU8B128 ||
+                       b_q_type == MARLIN_NAMESPACE_NAME::kFE4M3fn ||
+                       b_q_type == MARLIN_NAMESPACE_NAME::kFE2M1f,
+                   "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or "
+                   "float4_e2m1f when "
+                   "has_zp = False. Got = ",
+                   b_q_type.str());
   }
 
   if (has_zp && is_zp_float) {
-    PADDLE_ENFORCE(a.dtype() == paddle::DataType::FLOAT16,
-                "Computation type must be float16 (half) when using float zero "
-                "points.");
+    PADDLE_ENFORCE(
+        a.dtype() == paddle::DataType::FLOAT16,
+        "Computation type must be float16 (half) when using float zero "
+        "points.");
   }
 
   // Verify b_zeros
@@ -896,33 +1172,45 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
     PADDLE_ENFORCE(rank == 3, "b_zeros rank = ", rank, " is not 3");
     if (is_zp_float) {
       PADDLE_ENFORCE(b_zeros.size(2) == size_n,
-                  "b_zeros dim 2 = ", b_zeros.size(2),
-                  " is not size_n = ", size_n);
+                     "b_zeros dim 2 = ",
+                     b_zeros.size(2),
+                     " is not size_n = ",
+                     size_n);
       PADDLE_ENFORCE(num_groups == b_zeros.size(1),
-                  "b_zeros dim 1 = ", b_zeros.size(1),
-                  " is not num_groups = ", num_groups);
+                     "b_zeros dim 1 = ",
+                     b_zeros.size(1),
+                     " is not num_groups = ",
+                     num_groups);
       PADDLE_ENFORCE(num_groups != -1, "num_groups must be != -1");
     } else {
       PADDLE_ENFORCE(b_zeros.size(1) == num_groups,
-                  "b_zeros dim 1 = ", b_zeros.size(1),
-                  " is not num_groups = ", num_groups);
+                     "b_zeros dim 1 = ",
+                     b_zeros.size(1),
+                     " is not num_groups = ",
+                     num_groups);
       PADDLE_ENFORCE(b_zeros.size(2) == size_n / pack_factor,
-                  "b_zeros dim 2 = ", b_zeros.size(2),
-                  " is not size_n / pack_factor = ", size_n / pack_factor);
+                     "b_zeros dim 2 = ",
+                     b_zeros.size(2),
+                     " is not size_n / pack_factor = ",
+                     size_n / pack_factor);
     }
   }
 
   // Verify workspace size
   PADDLE_ENFORCE(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
-              "size_n = ", size_n, ", is not divisible by min_thread_n = ",
-              MARLIN_NAMESPACE_NAME::min_thread_n);
+                 "size_n = ",
+                 size_n,
+                 ", is not divisible by min_thread_n = ",
+                 MARLIN_NAMESPACE_NAME::min_thread_n);
 
   int max_n_tiles = size_n / MARLIN_NAMESPACE_NAME::min_thread_n;
   int min_workspace_size = std::min(
       max_n_tiles * (int)(sorted_token_ids.size(0) / moe_block_size), sms * 4);
   PADDLE_ENFORCE(workspace.numel() >= min_workspace_size,
-              "workspace.numel = ", workspace.numel(),
-              " is below min_workspace_size = ", min_workspace_size);
+                 "workspace.numel = ",
+                 workspace.numel(),
+                 " is below min_workspace_size = ",
+                 min_workspace_size);
 
   if (a.dtype() == paddle::DataType::FLOAT16) {
     using DataType = phi::dtype::float16;
@@ -931,20 +1219,45 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
     if (b_q_type == MARLIN_NAMESPACE_NAME::kFE2M1f) {
       scales_ptr = b_scales.data_ptr<phi::dtype::float8_e4m3fn>();
     } else {
-      scales_ptr = b_scales.data_ptr<DataType>(); // half
+      scales_ptr = b_scales.data_ptr<DataType>();  // half
     }
 
-    MARLIN_NAMESPACE_NAME::marlin_mm<half>(
-        a.data_ptr<DataType>(), b_q_weight.data_ptr(), c.data_ptr<DataType>(),
-        c_tmp.data_ptr<float>(), scales_ptr, global_scale.data_ptr<DataType>(),
-        b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(),
-        a_tmp.data_ptr<DataType>(), sorted_token_ids.data_ptr(),
-        expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
-        topk_weights.data_ptr(), moe_block_size, top_k, mul_topk_weights, is_ep,
-        size_m, size_n, size_k, workspace.data_ptr(), b_q_type, has_act_order,
-        is_k_full, has_zp, num_groups, group_size, device_id,
-        a.raw_tensor_.stream(), thread_k, thread_n, sms,
-        use_atomic_add, use_fp32_reduce, is_zp_float);
+    MARLIN_NAMESPACE_NAME::marlin_mm<half>(a.data_ptr<DataType>(),
+                                           b_q_weight.data_ptr(),
+                                           c.data_ptr<DataType>(),
+                                           c_tmp.data_ptr<float>(),
+                                           scales_ptr,
+                                           global_scale.data_ptr<DataType>(),
+                                           b_zeros.data_ptr(),
+                                           g_idx.data_ptr(),
+                                           perm.data_ptr(),
+                                           a_tmp.data_ptr<DataType>(),
+                                           sorted_token_ids.data_ptr(),
+                                           expert_ids.data_ptr(),
+                                           num_tokens_past_padded.data_ptr(),
+                                           topk_weights.data_ptr(),
+                                           moe_block_size,
+                                           top_k,
+                                           mul_topk_weights,
+                                           is_ep,
+                                           size_m,
+                                           size_n,
+                                           size_k,
+                                           workspace.data_ptr(),
+                                           b_q_type,
+                                           has_act_order,
+                                           is_k_full,
+                                           has_zp,
+                                           num_groups,
+                                           group_size,
+                                           device_id,
+                                           a.raw_tensor_.stream(),
+                                           thread_k,
+                                           thread_n,
+                                           sms,
+                                           use_atomic_add,
+                                           use_fp32_reduce,
+                                           is_zp_float);
   } else if (a.dtype() == MARLIN_NAMESPACE_NAME::kBFloat16) {
     using DataType = phi::dtype::bfloat16;
 
@@ -956,19 +1269,45 @@ MARLIN_NAMESPACE_NAME::Tensor moe_wna16_marlin_gemm(
     }
 
     MARLIN_NAMESPACE_NAME::marlin_mm<nv_bfloat16>(
-        a.data_ptr<DataType>(), b_q_weight.data_ptr(), c.data_ptr<DataType>(),
-        c_tmp.data_ptr<float>(), scales_ptr,
-        global_scale.data_ptr<DataType>(), b_zeros.data_ptr(),
-        g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<DataType>(),
-        sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
-        num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
-        moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
-        workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,
-        num_groups, group_size, device_id, a.raw_tensor_.stream(),
-        thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float);
+        a.data_ptr<DataType>(),
+        b_q_weight.data_ptr(),
+        c.data_ptr<DataType>(),
+        c_tmp.data_ptr<float>(),
+        scales_ptr,
+        global_scale.data_ptr<DataType>(),
+        b_zeros.data_ptr(),
+        g_idx.data_ptr(),
+        perm.data_ptr(),
+        a_tmp.data_ptr<DataType>(),
+        sorted_token_ids.data_ptr(),
+        expert_ids.data_ptr(),
+        num_tokens_past_padded.data_ptr(),
+        topk_weights.data_ptr(),
+        moe_block_size,
+        top_k,
+        mul_topk_weights,
+        is_ep,
+        size_m,
+        size_n,
+        size_k,
+        workspace.data_ptr(),
+        b_q_type,
+        has_act_order,
+        is_k_full,
+        has_zp,
+        num_groups,
+        group_size,
+        device_id,
+        a.raw_tensor_.stream(),
+        thread_k,
+        thread_n,
+        sms,
+        use_atomic_add,
+        use_fp32_reduce,
+        is_zp_float);
   } else {
     PADDLE_ENFORCE(false,
-                "moe_wna16_marlin_gemm only supports bfloat16 and float16");
+                   "moe_wna16_marlin_gemm only supports bfloat16 and float16");
   }
 
   return c;
@@ -1002,12 +1341,13 @@ std::vector<paddle::Tensor> MoeWna16MarlinGemmApi(
     bool is_zp_float) {
   auto a_ = ConvertPaddleTensorToDetailTensor(a);
   auto b_q_weight_ = ConvertPaddleTensorToDetailTensor(b_q_weight);
-  auto b_scales_   = ConvertPaddleTensorToDetailTensor(b_scales);
-  auto workspace_  = ConvertPaddleTensorToDetailTensor(workspace);
-  auto sorted_token_ids_    = ConvertPaddleTensorToDetailTensor(sorted_token_ids);
-  auto expert_ids_          = ConvertPaddleTensorToDetailTensor(expert_ids);
-  auto num_tokens_padded_   = ConvertPaddleTensorToDetailTensor(num_tokens_post_padded);
-  auto topk_weights_        = ConvertPaddleTensorToDetailTensor(topk_weights);
+  auto b_scales_ = ConvertPaddleTensorToDetailTensor(b_scales);
+  auto workspace_ = ConvertPaddleTensorToDetailTensor(workspace);
+  auto sorted_token_ids_ = ConvertPaddleTensorToDetailTensor(sorted_token_ids);
+  auto expert_ids_ = ConvertPaddleTensorToDetailTensor(expert_ids);
+  auto num_tokens_padded_ =
+      ConvertPaddleTensorToDetailTensor(num_tokens_post_padded);
+  auto topk_weights_ = ConvertPaddleTensorToDetailTensor(topk_weights);
 
   paddle::optional<MARLIN_NAMESPACE_NAME::Tensor> c_opt_ =
       ConvertOptionalPaddleTensorToDetailTensor(c_or_none);
@@ -1029,32 +1369,32 @@ std::vector<paddle::Tensor> MoeWna16MarlinGemmApi(
     PADDLE_ENFORCE(false, "b_q_type_str not supported!");
   }
 
-  MARLIN_NAMESPACE_NAME::Tensor out_detail = moe_wna16_marlin_gemm(
-      a_,
-      c_opt_,
-      b_q_weight_,
-      b_scales_,
-      global_scale_opt_,
-      b_zeros_opt_,
-      g_idx_opt_,
-      perm_opt_,
-      workspace_,
-      sorted_token_ids_,
-      expert_ids_,
-      num_tokens_padded_,
-      topk_weights_,
-      moe_block_size,
-      top_k,
-      mul_topk_weights,
-      is_ep,
-      b_q_type_id,
-      size_m,
-      size_n,
-      size_k,
-      is_k_full,
-      use_atomic_add,
-      use_fp32_reduce,
-      is_zp_float);
+  MARLIN_NAMESPACE_NAME::Tensor out_detail =
+      moe_wna16_marlin_gemm(a_,
+                            c_opt_,
+                            b_q_weight_,
+                            b_scales_,
+                            global_scale_opt_,
+                            b_zeros_opt_,
+                            g_idx_opt_,
+                            perm_opt_,
+                            workspace_,
+                            sorted_token_ids_,
+                            expert_ids_,
+                            num_tokens_padded_,
+                            topk_weights_,
+                            moe_block_size,
+                            top_k,
+                            mul_topk_weights,
+                            is_ep,
+                            b_q_type_id,
+                            size_m,
+                            size_n,
+                            size_k,
+                            is_k_full,
+                            use_atomic_add,
+                            use_fp32_reduce,
+                            is_zp_float);
   paddle::Tensor out = ConvertDetailTensorToPaddleTensor(out_detail);
   return {out};
 }
@@ -1105,16 +1445,32 @@ std::vector<paddle::DataType> MoeWna16MarlinGemmInferDtype(
   return {a_dtype};
 }
 PD_BUILD_STATIC_OP(moe_wna16_marlin_gemm)
-    .Inputs({"a", paddle::Optional("c_or_none"), "b_q_weight", "b_scales",
-             paddle::Optional("global_scale_or_none"), paddle::Optional("b_zeros_or_none"),
-             paddle::Optional("g_idx_or_none"), paddle::Optional("perm_or_none"), "workspace",
-             "sorted_token_ids", "expert_ids", "num_tokens_post_padded",
+    .Inputs({"a",
+             paddle::Optional("c_or_none"),
+             "b_q_weight",
+             "b_scales",
+             paddle::Optional("global_scale_or_none"),
+             paddle::Optional("b_zeros_or_none"),
+             paddle::Optional("g_idx_or_none"),
+             paddle::Optional("perm_or_none"),
+             "workspace",
+             "sorted_token_ids",
+             "expert_ids",
+             "num_tokens_post_padded",
              "topk_weights"})
     .Outputs({"out"})
-    .Attrs({"moe_block_size:int64_t", "top_k:int64_t", "mul_topk_weights:bool",
-            "is_ep:bool", "b_q_type_str:std::string", "size_m:int64_t",
-            "size_n:int64_t", "size_k:int64_t", "is_k_full:bool",
-            "use_atomic_add:bool", "use_fp32_reduce:bool", "is_zp_float:bool"})
+    .Attrs({"moe_block_size:int64_t",
+            "top_k:int64_t",
+            "mul_topk_weights:bool",
+            "is_ep:bool",
+            "b_q_type_str:std::string",
+            "size_m:int64_t",
+            "size_n:int64_t",
+            "size_k:int64_t",
+            "is_k_full:bool",
+            "use_atomic_add:bool",
+            "use_fp32_reduce:bool",
+            "is_zp_float:bool"})
     .SetKernelFn(PD_KERNEL(MoeWna16MarlinGemmApi))
     .SetInferShapeFn(PD_INFER_SHAPE(MoeWna16MarlinGemmInferShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(MoeWna16MarlinGemmInferDtype));
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_gemm.h b/custom_ops/gpu_ops/moe/moe_wna16_marlin_gemm.h
index 3683b92f8b..bdb6254b47 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_gemm.h
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_gemm.h
@@ -1,6 +1,6 @@
 #pragma once
 #ifndef MARLIN_NAMESPACE_NAME
-  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
 #endif
 
 #include "paddle/phi/api/include/api.h"
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/ScalarType.h b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/ScalarType.h
index f7e5b005e8..9b1d026313 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/ScalarType.h
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/ScalarType.h
@@ -41,15 +41,18 @@ class ScalarType {
     NAN_REPR_ID_MAX
   };
 
-  constexpr ScalarType(uint8_t exponent, uint8_t mantissa, bool signed_,
-                       int32_t bias, bool finite_values_only = false,
+  constexpr ScalarType(uint8_t exponent,
+                       uint8_t mantissa,
+                       bool signed_,
+                       int32_t bias,
+                       bool finite_values_only = false,
                        NanRepr nan_repr = NAN_IEEE_754)
       : exponent(exponent),
         mantissa(mantissa),
         signed_(signed_),
         bias(bias),
         finite_values_only(finite_values_only),
-        nan_repr(nan_repr) {};
+        nan_repr(nan_repr){};
 
   static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) {
     return ScalarType(0, size_bits - 1, true, bias);
@@ -67,16 +70,17 @@ class ScalarType {
   }
 
   // IEEE 754 non-compliant floating point type
-  static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
+  static constexpr ScalarType float_(uint8_t exponent,
+                                     uint8_t mantissa,
                                      bool finite_values_only,
                                      NanRepr nan_repr) {
     // PADDLE_ENFORCE(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
     // PADDLE_ENFORCE(mantissa > 0 && exponent > 0);
     // PADDLE_ENFORCE(nan_repr != NAN_IEEE_754,
-    //             "use `float_IEEE754` constructor for floating point types that "
-    //             "follow IEEE 754 conventions");
-    return ScalarType(exponent, mantissa, true, 0, finite_values_only,
-                      nan_repr);
+    //             "use `float_IEEE754` constructor for floating point types
+    //             that " "follow IEEE 754 conventions");
+    return ScalarType(
+        exponent, mantissa, true, 0, finite_values_only, nan_repr);
   }
 
   uint8_t const exponent;  // size of the exponent field (0 for integer types)
@@ -103,7 +107,9 @@ class ScalarType {
   }
 
   template <typename Fn, typename Init, typename Member, typename... Rest>
-  static constexpr auto reduce_members_helper(Fn f, Init val, Member member,
+  static constexpr auto reduce_members_helper(Fn f,
+                                              Init val,
+                                              Member member,
                                               Rest... rest) {
     auto new_val = f(val, member);
     if constexpr (sizeof...(rest) > 0) {
@@ -116,8 +122,14 @@ class ScalarType {
   template <typename Fn, typename Init>
   constexpr auto reduce_members(Fn f, Init init) const {
     // Should be in constructor order for `from_id`
-    return reduce_members_helper(f, init, exponent, mantissa, signed_, bias,
-                                 finite_values_only, nan_repr);
+    return reduce_members_helper(f,
+                                 init,
+                                 exponent,
+                                 mantissa,
+                                 signed_,
+                                 bias,
+                                 finite_values_only,
+                                 nan_repr);
   };
 
   template <typename Fn, typename Init>
@@ -194,7 +206,8 @@ class ScalarType {
  private:
   double _floating_point_max() const {
     PADDLE_ENFORCE(mantissa <= 52 && exponent <= 11,
-                "Cannot represent max/min as a double for type ", str());
+                   "Cannot represent max/min as a double for type ",
+                   str());
 
     uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
     if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
@@ -204,7 +217,8 @@ class ScalarType {
     uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
     if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
       PADDLE_ENFORCE(exponent < 11,
-                  "Cannot represent max/min as a double for type ", str());
+                     "Cannot represent max/min as a double for type ",
+                     str());
       max_exponent += 1;
     }
 
@@ -327,23 +341,32 @@ using ScalarTypeId = MARLIN_NAMESPACE_NAME::ScalarType::Id;
 //   https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L60-L70
 static inline constexpr auto kS4 = MARLIN_NAMESPACE_NAME::ScalarType::int_(4);
 static inline constexpr auto kU4 = MARLIN_NAMESPACE_NAME::ScalarType::uint(4);
-static inline constexpr auto kU4B8 = MARLIN_NAMESPACE_NAME::ScalarType::uint(4, 8);
+static inline constexpr auto kU4B8 =
+    MARLIN_NAMESPACE_NAME::ScalarType::uint(4, 8);
 static inline constexpr auto kS8 = MARLIN_NAMESPACE_NAME::ScalarType::int_(8);
 static inline constexpr auto kU8 = MARLIN_NAMESPACE_NAME::ScalarType::uint(8);
-static inline constexpr auto kU8B128 = MARLIN_NAMESPACE_NAME::ScalarType::uint(8, 128);
+static inline constexpr auto kU8B128 =
+    MARLIN_NAMESPACE_NAME::ScalarType::uint(8, 128);
 
 static inline constexpr auto kFE2M1f =
-    MARLIN_NAMESPACE_NAME::ScalarType::float_(2, 1, true, MARLIN_NAMESPACE_NAME::ScalarType::NAN_NONE);
+    MARLIN_NAMESPACE_NAME::ScalarType::float_(
+        2, 1, true, MARLIN_NAMESPACE_NAME::ScalarType::NAN_NONE);
 static inline constexpr auto kFE3M2f =
-    MARLIN_NAMESPACE_NAME::ScalarType::float_(3, 2, true, MARLIN_NAMESPACE_NAME::ScalarType::NAN_NONE);
+    MARLIN_NAMESPACE_NAME::ScalarType::float_(
+        3, 2, true, MARLIN_NAMESPACE_NAME::ScalarType::NAN_NONE);
 static inline constexpr auto kFE4M3fn =
-    MARLIN_NAMESPACE_NAME::ScalarType::float_(4, 3, true, MARLIN_NAMESPACE_NAME::ScalarType::NAN_EXTD_RANGE_MAX_MIN);
-static inline constexpr auto kFE5M2 = MARLIN_NAMESPACE_NAME::ScalarType::float_IEEE754(5, 2);
-static inline constexpr auto kFE8M7 = MARLIN_NAMESPACE_NAME::ScalarType::float_IEEE754(8, 7);
-static inline constexpr auto kFE5M10 = MARLIN_NAMESPACE_NAME::ScalarType::float_IEEE754(5, 10);
+    MARLIN_NAMESPACE_NAME::ScalarType::float_(
+        4, 3, true, MARLIN_NAMESPACE_NAME::ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE5M2 =
+    MARLIN_NAMESPACE_NAME::ScalarType::float_IEEE754(5, 2);
+static inline constexpr auto kFE8M7 =
+    MARLIN_NAMESPACE_NAME::ScalarType::float_IEEE754(8, 7);
+static inline constexpr auto kFE5M10 =
+    MARLIN_NAMESPACE_NAME::ScalarType::float_IEEE754(5, 10);
 
 // // Fixed width style names, generally following:
-// //  https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57
+// //
+// https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57
 constexpr auto kInt4 = kS4;
 constexpr auto kUint4 = kU4;
 constexpr auto kUint4b8 = kU4B8;
@@ -369,4 +392,4 @@ constexpr auto kBFloat16 = phi::DataType::BFLOAT16;
 constexpr auto kFloat32 = phi::DataType::FLOAT32;
 constexpr auto kByte = phi::DataType::INT8;
 
-};  // namespace Marlin
+};  // namespace MARLIN_NAMESPACE_NAME
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/dequant.h b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/dequant.h
index 853a42021f..049b21a51f 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/dequant.h
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/dequant.h
@@ -92,7 +92,8 @@ __device__ inline uint32_t prmt(uint32_t a) {
   return res;
 }
 
-template <typename scalar_t2, MARLIN_NAMESPACE_NAME::ScalarTypeId w_type_id,
+template <typename scalar_t2,
+          MARLIN_NAMESPACE_NAME::ScalarTypeId w_type_id,
           bool skip_flop = false>
 __device__ inline void dequant(int q, scalar_t2* frag_b);
 
@@ -106,8 +107,8 @@ __device__ inline void dequant(int q, scalar_t2* frag_b);
 // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
 //
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4B8.id(), true>(int q,
-                                                              half2* frag_b) {
+__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4B8.id(), true>(
+    int q, half2* frag_b) {
   const int MASK = 0x000f000f;
   const int EX = 0x64006400;
   // Guarantee that the `(a & b) | c` operations are LOP3s.
@@ -120,8 +121,8 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4B8.id(), true>(i
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4B8.id(), false>(int q,
-                                                               half2* frag_b) {
+__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4B8.id(), false>(
+    int q, half2* frag_b) {
   const int LO = 0x000f000f;
   const int HI = 0x00f000f0;
   const int EX = 0x64006400;
@@ -143,14 +144,14 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4B8.id(), false>(
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4.id(), true>(int q,
-                                                            half2* frag_b) {
+__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4.id(), true>(
+    int q, half2* frag_b) {
   dequant<half2, MARLIN_NAMESPACE_NAME::kU4B8.id(), true>(q, frag_b);
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4.id(), false>(int q,
-                                                             half2* frag_b) {
+__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4.id(), false>(
+    int q, half2* frag_b) {
   const int LO = 0x000f000f;
   const int HI = 0x00f000f0;
   const int EX = 0x64006400;
@@ -172,7 +173,8 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU4.id(), false>(in
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4B8.id(), true>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4B8.id(), true>(
     int q, nv_bfloat162* frag_b) {
   static constexpr uint32_t MASK = 0x000f000f;
   static constexpr uint32_t EX = 0x43004300;
@@ -189,7 +191,8 @@ __device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4B8.id(),
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4B8.id(), false>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4B8.id(), false>(
     int q, nv_bfloat162* frag_b) {
   dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4B8.id(), true>(q, frag_b);
 
@@ -200,13 +203,15 @@ __device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4B8.id(),
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4.id(), true>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4.id(), true>(
     int q, nv_bfloat162* frag_b) {
   dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4B8.id(), true>(q, frag_b);
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4.id(), false>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4.id(), false>(
     int q, nv_bfloat162* frag_b) {
   dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4.id(), true>(q, frag_b);
 
@@ -225,8 +230,9 @@ __device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU4.id(), fa
 // https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
 //
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8B128.id(), true>(int q,
-                                                                half2* frag_b) {
+__device__ inline void
+dequant<half2, MARLIN_NAMESPACE_NAME::kU8B128.id(), true>(int q,
+                                                          half2* frag_b) {
   static constexpr uint32_t mask_for_elt_01 = 0x5250;
   static constexpr uint32_t mask_for_elt_23 = 0x5351;
   static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
@@ -239,8 +245,9 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8B128.id(), true>
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8B128.id(), false>(
-    int q, half2* frag_b) {
+__device__ inline void
+dequant<half2, MARLIN_NAMESPACE_NAME::kU8B128.id(), false>(int q,
+                                                           half2* frag_b) {
   dequant<half2, MARLIN_NAMESPACE_NAME::kU8B128.id(), true>(q, frag_b);
 
   static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
@@ -251,14 +258,14 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8B128.id(), false
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8.id(), true>(int q,
-                                                            half2* frag_b) {
+__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8.id(), true>(
+    int q, half2* frag_b) {
   dequant<half2, MARLIN_NAMESPACE_NAME::kU8B128.id(), true>(q, frag_b);
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8.id(), false>(int q,
-                                                             half2* frag_b) {
+__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8.id(), false>(
+    int q, half2* frag_b) {
   dequant<half2, MARLIN_NAMESPACE_NAME::kU8.id(), true>(q, frag_b);
 
   static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
@@ -269,7 +276,8 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kU8.id(), false>(in
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU8B128.id(), false>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU8B128.id(), false>(
     int q, nv_bfloat162* frag_b) {
   float fp32_intermediates[4];
   uint32_t* fp32_intermediates_casted =
@@ -287,14 +295,15 @@ __device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU8B128.id()
   fp32_intermediates[3] -= 8388736.f;
 
   uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
-  bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
-                                   fp32_intermediates_casted[1], 0x7632);
-  bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
-                                   fp32_intermediates_casted[3], 0x7632);
+  bf16_result_ptr[0] = __byte_perm(
+      fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632);
+  bf16_result_ptr[1] = __byte_perm(
+      fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632);
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU8.id(), false>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU8.id(), false>(
     int q, nv_bfloat162* frag_b) {
   float fp32_intermediates[4];
   uint32_t* fp32_intermediates_casted =
@@ -312,15 +321,16 @@ __device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kU8.id(), fa
   fp32_intermediates[3] -= 8388608.f;
 
   uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
-  bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
-                                   fp32_intermediates_casted[1], 0x7632);
-  bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
-                                   fp32_intermediates_casted[3], 0x7632);
+  bf16_result_ptr[0] = __byte_perm(
+      fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632);
+  bf16_result_ptr[1] = __byte_perm(
+      fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632);
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), true>(
-    int q, half2* frag_b) {
+__device__ inline void
+dequant<half2, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), true>(int q,
+                                                           half2* frag_b) {
   // Constants for FP8 (E4M3) and FP16 formats
   constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
   constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT;
@@ -337,8 +347,9 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), true
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), false>(
-    int q, half2* frag_b) {
+__device__ inline void
+dequant<half2, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), false>(int q,
+                                                            half2* frag_b) {
   dequant<half2, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), true>(q, frag_b);
 
   // Constants for FP8 (E4M3) and FP16 formats
@@ -355,7 +366,8 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), fals
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), true>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), true>(
     int q, nv_bfloat162* frag_b) {
   // Constants for FP8 (E4M3) and BF16 formats
   constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
@@ -374,7 +386,8 @@ __device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), false>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), false>(
     int q, nv_bfloat162* frag_b) {
   dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(), true>(q, frag_b);
 
@@ -396,8 +409,9 @@ __device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE4M3fn.id(
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), true>(int q,
-                                                                half2* frag_b) {
+__device__ inline void
+dequant<half2, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), true>(int q,
+                                                          half2* frag_b) {
   // Constants for FP4 (E2M1) and FP16 formats
   constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
   constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT;
@@ -414,8 +428,9 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), true>
 }
 
 template <>
-__device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), false>(
-    int q, half2* frag_b) {
+__device__ inline void
+dequant<half2, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), false>(int q,
+                                                           half2* frag_b) {
   dequant<half2, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), true>(q, frag_b);
 
   // Constants for FP4 (E2M1) and FP16 formats
@@ -432,7 +447,8 @@ __device__ inline void dequant<half2, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), false
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), true>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), true>(
     int q, nv_bfloat162* frag_b) {
   // Constants for FP4 (E2M1) and FP16 formats
   constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
@@ -450,7 +466,8 @@ __device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE2M1f.id()
 }
 
 template <>
-__device__ inline void dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), false>(
+__device__ inline void
+dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), false>(
     int q, nv_bfloat162* frag_b) {
   dequant<nv_bfloat162, MARLIN_NAMESPACE_NAME::kFE2M1f.id(), true>(q, frag_b);
 
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/kernel.h b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/kernel.h
index ea3644fb74..6c417b77e8 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/kernel.h
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/kernel.h
@@ -1,6 +1,6 @@
 
 #ifndef MARLIN_NAMESPACE_NAME
-  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
 #endif
 #include "moe/moe_wna16_marlin_utils/marlin.cuh"
 #include "moe/moe_wna16_marlin_utils/marlin_dtypes.cuh"
@@ -22,7 +22,8 @@
 
 namespace MARLIN_NAMESPACE_NAME {
 template <typename scalar_t,  // compute dtype, half or nv_float16
-          const MARLIN_NAMESPACE_NAME::ScalarTypeId w_type_id,  // weight ScalarType id
+          const MARLIN_NAMESPACE_NAME::ScalarTypeId w_type_id,  // weight
+                                                                // ScalarType id
           const int threads,          // number of threads in a threadblock
           const int thread_m_blocks,  // number of 16x16 blocks in the m
                                       // dimension (batchsize) of the
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin.cuh b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin.cuh
index 1ac0517a5d..05649312a0 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin.cuh
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin.cuh
@@ -10,7 +10,7 @@
 #include <iostream>
 
 #ifndef MARLIN_NAMESPACE_NAME
-  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
 #endif
 
 namespace MARLIN_NAMESPACE_NAME {
@@ -51,10 +51,10 @@ using I4 = Vec<int, 4>;
 
 constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
 
-
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
 
-__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
+__device__ inline void cp_async4_pred(void* smem_ptr,
+                                      const void* glob_ptr,
                                       bool pred = true) {
   const int BYTES = 16;
   uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
@@ -64,7 +64,9 @@ __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
       "   setp.ne.b32 p, %0, 0;\n"
       "   @p cp.async.cg.shared.global [%1], [%2], %3;\n"
       "}\n" ::"r"((int)pred),
-      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+      "r"(smem),
+      "l"(glob_ptr),
+      "n"(BYTES));
 }
 
 __device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
@@ -74,7 +76,8 @@ __device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
       "{\n"
       "   cp.async.cg.shared.global [%0], [%1], %2;\n"
       "}\n" ::"r"(smem),
-      "l"(glob_ptr), "n"(BYTES));
+      "l"(glob_ptr),
+      "n"(BYTES));
 }
 
 __device__ inline void cp_async_fence() {
@@ -93,7 +96,6 @@ inline void cp_async_fence() {}
 template <int n>
 inline void cp_async_wait() {}
 
-
 #endif
 
 }  // namespace MARLIN_NAMESPACE_NAME
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin_dtypes.cuh b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin_dtypes.cuh
index 216f40acde..3e4874bc3a 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin_dtypes.cuh
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin_dtypes.cuh
@@ -6,7 +6,7 @@
 #include <cuda_bf16.h>
 
 #ifndef MARLIN_NAMESPACE_NAME
-  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
 #endif
 
 namespace MARLIN_NAMESPACE_NAME::kernel_types {
@@ -78,6 +78,6 @@ class ScalarType<nv_bfloat16> {
 #endif
 };
 
-}  // namespace MARLIN_NAMESPACE_NAME
+}  // namespace MARLIN_NAMESPACE_NAME::kernel_types
 
 #endif
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin_template.h b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin_template.h
index 6de2cd83d9..3697407744 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin_template.h
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/marlin_template.h
@@ -20,7 +20,7 @@
  */
 
 #ifndef MARLIN_NAMESPACE_NAME
-  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
 #endif
 
 #include "moe/moe_wna16_marlin_utils/marlin.cuh"
@@ -38,7 +38,8 @@ namespace MARLIN_NAMESPACE_NAME {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
 
 template <typename scalar_t,  // compute dtype, half or nv_float16
-          const MARLIN_NAMESPACE_NAME::ScalarTypeId w_type_id,  // weight ScalarType id
+          const MARLIN_NAMESPACE_NAME::ScalarTypeId w_type_id,  // weight
+                                                                // ScalarType id
           const int threads,          // number of threads in a threadblock
           const int thread_m_blocks,  // number of 16x16 blocks in the m
                                       // dimension (batchsize) of the
@@ -86,9 +87,10 @@ __global__ void Marlin(
 // m16n8k16 tensor core mma instruction with fp16 inputs and fp32
 // output/accumulation.
 template <typename scalar_t>
-__device__ inline void mma(const typename kernel_types::ScalarType<scalar_t>::FragA& a_frag,
-                           const typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
-                           typename kernel_types::ScalarType<scalar_t>::FragC& frag_c) {
+__device__ inline void mma(
+    const typename kernel_types::ScalarType<scalar_t>::FragA& a_frag,
+    const typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
+    typename kernel_types::ScalarType<scalar_t>::FragC& frag_c) {
   const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
   const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
   float* c = reinterpret_cast<float*>(&frag_c);
@@ -97,15 +99,31 @@ __device__ inline void mma(const typename kernel_types::ScalarType<scalar_t>::Fr
         "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
         "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
         : "=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]));
+        : "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 (std::is_same<scalar_t, nv_bfloat16>::value) {
     asm volatile(
         "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
         "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
         : "=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]));
+        : "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 {
     STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
   }
@@ -126,15 +144,31 @@ __device__ inline void mma_trans(
         "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
         "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
         : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
-        : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
-          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+        : "r"(b[0]),
+          "r"(b2[0]),
+          "r"(b[1]),
+          "r"(b2[1]),
+          "r"(a[0]),
+          "r"(a[1]),
+          "f"(c[0]),
+          "f"(c[1]),
+          "f"(c[2]),
+          "f"(c[3]));
   } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
     asm volatile(
         "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
         "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
         : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
-        : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
-          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+        : "r"(b[0]),
+          "r"(b2[0]),
+          "r"(b[1]),
+          "r"(b2[1]),
+          "r"(a[0]),
+          "r"(a[1]),
+          "f"(c[0]),
+          "f"(c[1]),
+          "f"(c[2]),
+          "f"(c[3]));
   } else {
     STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
   }
@@ -143,8 +177,9 @@ __device__ inline void mma_trans(
 // Instruction for loading a full 16x16 matrix fragment of operand A from shared
 // memory, directly in tensor core layout.
 template <int count, typename scalar_t>
-__device__ inline void ldsm(typename kernel_types::ScalarType<scalar_t>::FragA& frag_a,
-                            const void* smem_ptr) {
+__device__ inline void ldsm(
+    typename kernel_types::ScalarType<scalar_t>::FragA& frag_a,
+    const void* smem_ptr) {
   uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
   uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
   if constexpr (count == 4) {
@@ -168,19 +203,22 @@ __device__ inline void ldsm(typename kernel_types::ScalarType<scalar_t>::FragA&
 // Multiply dequantized values by the corresponding quantization scale; used
 // only for grouped quantization.
 template <typename scalar_t>
-__device__ inline void scale(typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
-                             typename kernel_types::ScalarType<scalar_t>::FragS& frag_s,
-                             int i) {
+__device__ inline void scale(
+    typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
+    typename kernel_types::ScalarType<scalar_t>::FragS& frag_s,
+    int i) {
   using scalar_t2 = typename kernel_types::ScalarType<scalar_t>::scalar_t2;
-  scalar_t2 s =
-      kernel_types::ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_s)[i]);
+  scalar_t2 s = kernel_types::ScalarType<scalar_t>::num2num2(
+      reinterpret_cast<scalar_t*>(&frag_s)[i]);
   frag_b[0] = __hmul2(frag_b[0], s);
   frag_b[1] = __hmul2(frag_b[1], s);
 }
 
 template <typename scalar_t>
 __device__ inline void scale_and_sub(
-    typename kernel_types::ScalarType<scalar_t>::FragB& frag_b, scalar_t s, scalar_t zp) {
+    typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
+    scalar_t s,
+    scalar_t zp) {
   using scalar_t2 = typename kernel_types::ScalarType<scalar_t>::scalar_t2;
   scalar_t2 s2 = kernel_types::ScalarType<scalar_t>::num2num2(s);
   scalar_t2 zp2 = kernel_types::ScalarType<scalar_t>::num2num2(zp);
@@ -189,24 +227,26 @@ __device__ inline void scale_and_sub(
 }
 
 template <typename scalar_t>
-__device__ inline void sub_zp(typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
-                              typename kernel_types::ScalarType<scalar_t>::scalar_t2& frag_zp,
-                              int i) {
+__device__ inline void sub_zp(
+    typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
+    typename kernel_types::ScalarType<scalar_t>::scalar_t2& frag_zp,
+    int i) {
   using scalar_t2 = typename kernel_types::ScalarType<scalar_t>::scalar_t2;
-  scalar_t2 zp =
-      kernel_types::ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_zp)[i]);
+  scalar_t2 zp = kernel_types::ScalarType<scalar_t>::num2num2(
+      reinterpret_cast<scalar_t*>(&frag_zp)[i]);
   frag_b[0] = __hsub2(frag_b[0], zp);
   frag_b[1] = __hsub2(frag_b[1], zp);
 }
 
 // Same as above, but for act_order (each K is multiplied individually)
 template <typename scalar_t>
-__device__ inline void scale4(typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
-                              typename kernel_types::ScalarType<scalar_t>::FragS& frag_s_1,
-                              typename kernel_types::ScalarType<scalar_t>::FragS& frag_s_2,
-                              typename kernel_types::ScalarType<scalar_t>::FragS& frag_s_3,
-                              typename kernel_types::ScalarType<scalar_t>::FragS& frag_s_4,
-                              int i) {
+__device__ inline void scale4(
+    typename kernel_types::ScalarType<scalar_t>::FragB& frag_b,
+    typename kernel_types::ScalarType<scalar_t>::FragS& frag_s_1,
+    typename kernel_types::ScalarType<scalar_t>::FragS& frag_s_2,
+    typename kernel_types::ScalarType<scalar_t>::FragS& frag_s_3,
+    typename kernel_types::ScalarType<scalar_t>::FragS& frag_s_4,
+    int i) {
   using scalar_t2 = typename kernel_types::ScalarType<scalar_t>::scalar_t2;
   scalar_t2 s_val_1_2;
   s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
@@ -222,11 +262,13 @@ __device__ inline void scale4(typename kernel_types::ScalarType<scalar_t>::FragB
 
 // Given 2 floats multiply by 2 scales (halves)
 template <typename scalar_t>
-__device__ inline void scale_float(float* c,
-                                   typename kernel_types::ScalarType<scalar_t>::FragS& s) {
+__device__ inline void scale_float(
+    float* c, typename kernel_types::ScalarType<scalar_t>::FragS& s) {
   scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s);
-  c[0] = __fmul_rn(c[0], kernel_types::ScalarType<scalar_t>::num2float(s_ptr[0]));
-  c[1] = __fmul_rn(c[1], kernel_types::ScalarType<scalar_t>::num2float(s_ptr[1]));
+  c[0] =
+      __fmul_rn(c[0], kernel_types::ScalarType<scalar_t>::num2float(s_ptr[0]));
+  c[1] =
+      __fmul_rn(c[1], kernel_types::ScalarType<scalar_t>::num2float(s_ptr[1]));
 }
 
 // Wait until barrier reaches `count`, then lock for current threadblock.
@@ -279,7 +321,8 @@ __device__ inline void wait_negative_and_add(int* lock) {
 }
 
 template <typename scalar_t,  // compute dtype, half or nv_float16
-          const MARLIN_NAMESPACE_NAME::ScalarTypeId w_type_id,  // weight ScalarType id
+          const MARLIN_NAMESPACE_NAME::ScalarTypeId w_type_id,  // weight
+                                                                // ScalarType id
           const int threads,          // number of threads in a threadblock
           const int thread_m_blocks,  // number of 16x16 blocks in the m
                                       // dimension (batchsize) of the
@@ -341,10 +384,14 @@ __global__ void Marlin(
   using FragZP = typename kernel_types::ScalarType<scalar_t>::FragZP;
 
   extern __shared__ int4 sh[];
-  static constexpr auto w_type = MARLIN_NAMESPACE_NAME::ScalarType::from_id(w_type_id);
-  constexpr bool has_zp = w_type == MARLIN_NAMESPACE_NAME::kU4 || w_type == MARLIN_NAMESPACE_NAME::kU8;
-  constexpr bool is_int_type = w_type == MARLIN_NAMESPACE_NAME::kU4 || w_type == MARLIN_NAMESPACE_NAME::kU8 ||
-                               w_type == MARLIN_NAMESPACE_NAME::kU4B8 || w_type == MARLIN_NAMESPACE_NAME::kU8B128;
+  static constexpr auto w_type =
+      MARLIN_NAMESPACE_NAME::ScalarType::from_id(w_type_id);
+  constexpr bool has_zp = w_type == MARLIN_NAMESPACE_NAME::kU4 ||
+                          w_type == MARLIN_NAMESPACE_NAME::kU8;
+  constexpr bool is_int_type = w_type == MARLIN_NAMESPACE_NAME::kU4 ||
+                               w_type == MARLIN_NAMESPACE_NAME::kU8 ||
+                               w_type == MARLIN_NAMESPACE_NAME::kU4B8 ||
+                               w_type == MARLIN_NAMESPACE_NAME::kU8B128;
   // see comments of dequant.h for more details
   constexpr bool dequant_skip_flop =
       !is_int_type ||
@@ -361,7 +408,8 @@ __global__ void Marlin(
   const int group_size =
       (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups;
   const int scales_expert_stride =
-      prob_n * prob_k / group_size / (w_type == MARLIN_NAMESPACE_NAME::kFE2M1f ? 16 : 8);
+      prob_n * prob_k / group_size /
+      (w_type == MARLIN_NAMESPACE_NAME::kFE2M1f ? 16 : 8);
   const int zp_expert_stride =
       is_zp_float ? prob_n * prob_k / group_size / 8
                   : prob_n * prob_k / group_size / (pack_factor * 4);
@@ -444,12 +492,12 @@ __global__ void Marlin(
   // block_sorted_ids / block_num_valid_tokens / block_topk_weights
   auto read_moe_block_data = [&](int block_id) {
     block_num_valid_tokens = moe_block_size;
-  #pragma unroll
+#pragma unroll
     for (int i = 0; i < moe_block_size / 4; i++) {
       int4 sorted_token_ids_int4 = reinterpret_cast<const int4*>(
           sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i];
       int* sorted_token_ids = reinterpret_cast<int*>(&sorted_token_ids_int4);
-  #pragma unroll
+#pragma unroll
       for (int j = 0; j < 4; j++) {
         if (sorted_token_ids[j] >= prob_m * top_k) {
           block_num_valid_tokens = i * 4 + j;
@@ -465,20 +513,21 @@ __global__ void Marlin(
       sh_block_sorted_ids_int4[tid4] = reinterpret_cast<const int4*>(
           sorted_token_ids_ptr)[block_id * moe_block_size / 4 + tid4];
 
-  #pragma unroll
+#pragma unroll
       for (int i = 0; i < 4; i++)
         sh_rd_block_sorted_ids[tid4 * 4 + i] =
             sh_block_sorted_ids[tid4 * 4 + i] / top_k;
 
       if (mul_topk_weights) {
-  #pragma unroll
+#pragma unroll
         for (int i = 0; i < 4; i++) {
           int idx = tid4 * 4 + i;
           idx = idx < block_num_valid_tokens ? idx : 0;
           if constexpr (w_type == MARLIN_NAMESPACE_NAME::kFE2M1f) {
-            sh_block_topk_weights[idx] = __hmul2(
-                global_scale, Dtype::num2num2(Dtype::float2num(
-                                  topk_weights_ptr[sh_block_sorted_ids[idx]])));
+            sh_block_topk_weights[idx] =
+                __hmul2(global_scale,
+                        Dtype::num2num2(Dtype::float2num(
+                            topk_weights_ptr[sh_block_sorted_ids[idx]])));
           } else {
             sh_block_topk_weights[idx] = Dtype::num2num2(
                 Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]));
@@ -631,7 +680,8 @@ __global__ void Marlin(
   constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
   constexpr int s_tb_groups =
       !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
-          ? thread_k_blocks / group_blocks / (w_type == MARLIN_NAMESPACE_NAME::kFE2M1f ? 2 : 1)
+          ? thread_k_blocks / group_blocks /
+                (w_type == MARLIN_NAMESPACE_NAME::kFE2M1f ? 2 : 1)
           : 1;
   constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
   int s_gl_rd_delta = s_gl_stride;
@@ -714,7 +764,8 @@ __global__ void Marlin(
   // we scale a `half2` tile in column-major layout in the former and in
   // row-major in the latter case.
   int s_sh_rd;
-  if constexpr (group_blocks != -1 && w_type == MARLIN_NAMESPACE_NAME::kFE2M1f) {
+  if constexpr (group_blocks != -1 &&
+                w_type == MARLIN_NAMESPACE_NAME::kFE2M1f) {
     auto warp_id = threadIdx.x / 32;
     int n_warps = thread_n_blocks / 4;
     int warp_row = warp_id / n_warps;
@@ -767,13 +818,13 @@ __global__ void Marlin(
   // loop unrolls, all shared memory accesses are static, we simply precompute
   // both transformed reads and writes.
   int a_sh_wr_trans[a_sh_wr_iters];
-  #pragma unroll
+#pragma unroll
   for (int i = 0; i < a_sh_wr_iters; i++)
     a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
   int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];
-  #pragma unroll
+#pragma unroll
   for (int i = 0; i < b_sh_wr_iters; i++) {
-  #pragma unroll
+#pragma unroll
     for (int j = 0; j < thread_m_blocks; j++)
       a_sh_rd_trans[i][j] =
           transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
@@ -784,7 +835,7 @@ __global__ void Marlin(
   // maintining multiple pointers (we have enough registers), a tiny
   // optimization.
   const int4* B_ptr[b_sh_wr_iters];
-  #pragma unroll
+#pragma unroll
   for (int i = 0; i < b_sh_wr_iters; i++)
     B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
 
@@ -824,7 +875,7 @@ __global__ void Marlin(
 
   // Zero accumulators.
   auto zero_accums = [&]() {
-  #pragma unroll
+#pragma unroll
     for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
       reinterpret_cast<float*>(frag_c)[i] = 0;
   };
@@ -832,39 +883,39 @@ __global__ void Marlin(
   int sh_first_group_id = -1;
   int sh_num_groups = -1;
 
-  auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id,
-                                              int last_group_id) {
-    sh_first_group_id = first_group_id;
-    sh_num_groups = last_group_id - first_group_id + 1;
+  auto fetch_act_order_scales_to_shared =
+      [&](bool is_async, int first_group_id, int last_group_id) {
+        sh_first_group_id = first_group_id;
+        sh_num_groups = last_group_id - first_group_id + 1;
 
-    if (sh_num_groups > act_s_max_num_groups) {
-      sh_num_groups = act_s_max_num_groups;
-    }
-
-    if (sh_first_group_id + sh_num_groups > num_groups) {
-      sh_num_groups = num_groups - sh_first_group_id;
-    }
-
-    int row_offset = first_group_id * s_gl_stride;
-
-    if (is_async) {
-      for (int i = 0; i < sh_num_groups; i++) {
-        if (threadIdx.x < s_sh_stride) {
-          cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
-                         &scales_ptr[row_offset + (i * s_gl_stride) +
-                                     slice_n_offset + threadIdx.x]);
+        if (sh_num_groups > act_s_max_num_groups) {
+          sh_num_groups = act_s_max_num_groups;
         }
-      }
-    } else {
-      for (int i = 0; i < sh_num_groups; i++) {
-        if (threadIdx.x < s_sh_stride) {
-          sh_s[(i * s_sh_stride) + threadIdx.x] =
-              scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
-                         threadIdx.x];
+
+        if (sh_first_group_id + sh_num_groups > num_groups) {
+          sh_num_groups = num_groups - sh_first_group_id;
         }
-      }
-    }
-  };
+
+        int row_offset = first_group_id * s_gl_stride;
+
+        if (is_async) {
+          for (int i = 0; i < sh_num_groups; i++) {
+            if (threadIdx.x < s_sh_stride) {
+              cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
+                             &scales_ptr[row_offset + (i * s_gl_stride) +
+                                         slice_n_offset + threadIdx.x]);
+            }
+          }
+        } else {
+          for (int i = 0; i < sh_num_groups; i++) {
+            if (threadIdx.x < s_sh_stride) {
+              sh_s[(i * s_sh_stride) + threadIdx.x] =
+                  scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
+                             threadIdx.x];
+            }
+          }
+        }
+      };
 
   // Asynchronously fetch the next A, B and s tile from global to the next
   // shared memory pipeline location.
@@ -872,101 +923,102 @@ __global__ void Marlin(
   int max_num_stage_groups =
       ((sh_a_max_row - moe_block_size) / moe_block_size + 1) / stages;
   max_num_stage_groups = max(max_num_stage_groups, 1);
-  auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true,
-                             int pipe_a = 0) {
-    if (pred) {
-      if (should_load_a) {
-        int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a;
-  #pragma unroll
-        for (int i = 0; i < a_sh_wr_iters; i++) {
-          int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row;
-          int64_t sorted_row = 0;
-          if (!m_block_size_8 || row < 8)
-            sorted_row = sh_rd_block_sorted_ids[row];
-          int64_t true_idx =
-              sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off;
-          cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx],
-                         row < block_num_valid_tokens);
-        }
-      }
-
-      int4* sh_b_stage = sh_b + b_sh_stage * pipe;
-  #pragma unroll
-      for (int i = 0; i < b_sh_wr_iters; i++) {
-  #pragma unroll
-        for (int j = 0; j < b_thread_vecs; j++) {
-          cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j],
-                    B_ptr[i] + j + B_expert_off);
-        }
-
-        B_ptr[i] += b_gl_rd_delta_o;
-      }
-
-      if constexpr (has_act_order) {
-        // Fetch g_idx thread-block portion
-        int full_pipe = a_off;
-        int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
-        if (cur_k < prob_k && cur_k < slice_k_finish) {
-          int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
-
-          int4 const* cur_g_idx_stage_ptr =
-              reinterpret_cast<int4 const*>(&g_idx[cur_k]);
-
-          if (threadIdx.x < g_idx_stage) {
-            cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
-                           &cur_g_idx_stage_ptr[threadIdx.x]);
+  auto fetch_to_shared =
+      [&](int pipe, int a_off, bool pred = true, int pipe_a = 0) {
+        if (pred) {
+          if (should_load_a) {
+            int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a;
+#pragma unroll
+            for (int i = 0; i < a_sh_wr_iters; i++) {
+              int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row;
+              int64_t sorted_row = 0;
+              if (!m_block_size_8 || row < 8)
+                sorted_row = sh_rd_block_sorted_ids[row];
+              int64_t true_idx = sorted_row * a_gl_stride + a_gl_rd_col +
+                                 a_gl_rd_delta_o * a_off;
+              cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]],
+                             &A[true_idx],
+                             row < block_num_valid_tokens);
+            }
           }
-        }
-      } else {
-        if constexpr (group_blocks != -1) {
-          int4* sh_s_stage = sh_s + s_sh_stage * pipe;
 
-          if constexpr (group_blocks >= thread_k_blocks) {
-            // Only fetch scales if this tile starts a new group
-            if (pipe % (group_blocks / thread_k_blocks) == 0) {
-              if (s_sh_wr_pred) {
-                cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
+          int4* sh_b_stage = sh_b + b_sh_stage * pipe;
+#pragma unroll
+          for (int i = 0; i < b_sh_wr_iters; i++) {
+#pragma unroll
+            for (int j = 0; j < b_thread_vecs; j++) {
+              cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j],
+                        B_ptr[i] + j + B_expert_off);
+            }
+
+            B_ptr[i] += b_gl_rd_delta_o;
+          }
+
+          if constexpr (has_act_order) {
+            // Fetch g_idx thread-block portion
+            int full_pipe = a_off;
+            int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
+            if (cur_k < prob_k && cur_k < slice_k_finish) {
+              int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
+
+              int4 const* cur_g_idx_stage_ptr =
+                  reinterpret_cast<int4 const*>(&g_idx[cur_k]);
+
+              if (threadIdx.x < g_idx_stage) {
+                cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
+                               &cur_g_idx_stage_ptr[threadIdx.x]);
               }
-              s_gl_rd += s_gl_rd_delta;
             }
           } else {
-            for (int i = 0; i < s_tb_groups; i++) {
-              if (s_sh_wr_pred) {
-                cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr],
-                          &scales_ptr[s_gl_rd]);
+            if constexpr (group_blocks != -1) {
+              int4* sh_s_stage = sh_s + s_sh_stage * pipe;
+
+              if constexpr (group_blocks >= thread_k_blocks) {
+                // Only fetch scales if this tile starts a new group
+                if (pipe % (group_blocks / thread_k_blocks) == 0) {
+                  if (s_sh_wr_pred) {
+                    cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
+                  }
+                  s_gl_rd += s_gl_rd_delta;
+                }
+              } else {
+                for (int i = 0; i < s_tb_groups; i++) {
+                  if (s_sh_wr_pred) {
+                    cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr],
+                              &scales_ptr[s_gl_rd]);
+                  }
+                  s_gl_rd += s_gl_rd_delta;
+                }
+              }
+            }
+
+            if constexpr (has_zp && group_blocks != -1) {
+              int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
+
+              if constexpr (group_blocks >= thread_k_blocks) {
+                // Only fetch zero-points if this tile starts a new group
+                if (pipe % (group_blocks / thread_k_blocks) == 0) {
+                  if (zp_sh_wr_pred) {
+                    cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
+                  }
+                  zp_gl_rd += zp_gl_rd_delta;
+                }
+              } else {
+                for (int i = 0; i < zp_tb_groups; i++) {
+                  if (zp_sh_wr_pred) {
+                    cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr],
+                              &zp_ptr[zp_gl_rd]);
+                  }
+                  zp_gl_rd += zp_gl_rd_delta;
+                }
               }
-              s_gl_rd += s_gl_rd_delta;
             }
           }
         }
-
-        if constexpr (has_zp && group_blocks != -1) {
-          int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
-
-          if constexpr (group_blocks >= thread_k_blocks) {
-            // Only fetch zero-points if this tile starts a new group
-            if (pipe % (group_blocks / thread_k_blocks) == 0) {
-              if (zp_sh_wr_pred) {
-                cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
-              }
-              zp_gl_rd += zp_gl_rd_delta;
-            }
-          } else {
-            for (int i = 0; i < zp_tb_groups; i++) {
-              if (zp_sh_wr_pred) {
-                cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr],
-                          &zp_ptr[zp_gl_rd]);
-              }
-              zp_gl_rd += zp_gl_rd_delta;
-            }
-          }
-        }
-      }
-    }
-    // Insert a fence even when we are winding down the pipeline to ensure that
-    // waiting is also correct at this point.
-    cp_async_fence();
-  };
+        // Insert a fence even when we are winding down the pipeline to ensure
+        // that waiting is also correct at this point.
+        cp_async_fence();
+      };
 
   auto fetch_col_zp_to_shared = [&]() {
     if (zp_sh_wr_pred) {
@@ -994,13 +1046,13 @@ __global__ void Marlin(
   // into the current register buffer.
   auto fetch_to_registers = [&](int k, int pipe, int pipe_a = 0) {
     int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a;
-  #pragma unroll
+#pragma unroll
     for (int i = 0; i < thread_m_blocks; i++)
       ldsm<m_block_size_8 ? 2 : 4, scalar_t>(
           frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
     int4* sh_b_stage = sh_b + b_sh_stage * pipe;
 
-  #pragma unroll
+#pragma unroll
     for (int i = 0; i < b_thread_vecs; i++) {
       frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
           &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]);
@@ -1058,7 +1110,8 @@ __global__ void Marlin(
 
           int k_blocks = cur_k / 16;
           int cur_group_id =
-              k_blocks / (group_blocks * (w_type == MARLIN_NAMESPACE_NAME::kFE2M1f ? 2 : 1));
+              k_blocks / (group_blocks *
+                          (w_type == MARLIN_NAMESPACE_NAME::kFE2M1f ? 2 : 1));
 
           int4* sh_s_stage = sh_s + s_sh_stage * pipe;
 
@@ -1129,10 +1182,10 @@ __global__ void Marlin(
     int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
     int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
 
-    constexpr int k_frag_offsets[4] = {0, 1, 8,
-                                       9};  // Tensor core offsets per thread
+    constexpr int k_frag_offsets[4] = {
+        0, 1, 8, 9};  // Tensor core offsets per thread
 
-  #pragma unroll
+#pragma unroll
     for (int i = 0; i < 4; i++) {
       int actual_k = cur_k + k_frag_offsets[i];
 
@@ -1156,7 +1209,7 @@ __global__ void Marlin(
       if constexpr (group_blocks == -1) {
         // load only when starting a new slice
         if (k == 0 && full_pipe == 0) {
-  #pragma unroll
+#pragma unroll
           for (int i = 0; i < num_ints_per_thread; i++) {
             frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];
           }
@@ -1167,7 +1220,7 @@ __global__ void Marlin(
           int4* sh_zp_stage =
               sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
                                      (pipe / (group_blocks / thread_k_blocks)));
-  #pragma unroll
+#pragma unroll
           for (int i = 0; i < num_ints_per_thread; i++) {
             frag_qzp[k % 2][i] =
                 (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
@@ -1186,16 +1239,16 @@ __global__ void Marlin(
         int cur_group_id = 0;
 
         // Suppress bogus and persistent divide-by-zero warning
-  #pragma nv_diagnostic push
-  #pragma nv_diag_suppress divide_by_zero
+#pragma nv_diagnostic push
+#pragma nv_diag_suppress divide_by_zero
         cur_group_id = k_blocks / group_blocks;
-  #pragma nv_diagnostic pop
+#pragma nv_diagnostic pop
 
         int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
 
         sh_zp_stage += cur_group_id * zp_sh_stride;
 
-  #pragma unroll
+#pragma unroll
         for (int i = 0; i < num_ints_per_thread; i++) {
           frag_qzp[k % 2][i] =
               (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
@@ -1227,10 +1280,10 @@ __global__ void Marlin(
 
           int k_blocks = cur_k / 16;
           // Suppress bogus and persistent divide-by-zero warning
-  #pragma nv_diagnostic push
-  #pragma nv_diag_suppress divide_by_zero
+#pragma nv_diagnostic push
+#pragma nv_diag_suppress divide_by_zero
           int cur_group_id = k_blocks / group_blocks;
-  #pragma nv_diagnostic pop
+#pragma nv_diagnostic pop
 
           int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
 
@@ -1287,9 +1340,9 @@ __global__ void Marlin(
           s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
     }
 
-  // We have the m dimension as the inner loop in order to encourage overlapping
-  // dequantization and matmul operations.
-  #pragma unroll
+// We have the m dimension as the inner loop in order to encourage overlapping
+// dequantization and matmul operations.
+#pragma unroll
     for (int j = 0; j < 4; j++) {
       FragB frag_b0;
       FragB frag_b1;
@@ -1319,10 +1372,18 @@ __global__ void Marlin(
       // Apply scale to frag_b0
       if constexpr (has_act_order) {
         static_assert(group_blocks != -1);
-        scale4<scalar_t>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
-                         act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
-        scale4<scalar_t>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
-                         act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
+        scale4<scalar_t>(frag_b0,
+                         act_frag_s[k2][0][j],
+                         act_frag_s[k2][1][j],
+                         act_frag_s[k2][2][j],
+                         act_frag_s[k2][3][j],
+                         0);
+        scale4<scalar_t>(frag_b1,
+                         act_frag_s[k2][0][j],
+                         act_frag_s[k2][1][j],
+                         act_frag_s[k2][2][j],
+                         act_frag_s[k2][3][j],
+                         1);
       } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
                            group_blocks == -1) {
         int idx = (threadIdx.x / 4) % 2;
@@ -1343,7 +1404,7 @@ __global__ void Marlin(
         scale<scalar_t>(frag_b1, frag_s[k2][j], 1);
       }
 
-  #pragma unroll
+#pragma unroll
       for (int i = 0; i < thread_m_blocks; i++) {
         if constexpr (m_block_size_8) {
           mma_trans<scalar_t>(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]);
@@ -1372,12 +1433,12 @@ __global__ void Marlin(
       // unnecessary read or write iterations, e.g., for two warps we write only
       // once by warp 1 and read only once by warp 0.
 
-  #pragma unroll
+#pragma unroll
       for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
-  #pragma unroll
+#pragma unroll
         for (int i = red_off; i > 0; i /= 2) {
           if (i <= red_idx && red_idx < 2 * i) {
-  #pragma unroll
+#pragma unroll
             for (int j = 0; j < 4 * 2; j += (m_block_size_8 ? 2 : 1)) {
               int red_sh_wr =
                   red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
@@ -1385,7 +1446,7 @@ __global__ void Marlin(
                 float* c_rd = reinterpret_cast<float*>(
                     &sh_red[red_sh_delta * j + red_sh_rd]);
                 float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]);
-  #pragma unroll
+#pragma unroll
                 for (int k = 0; k < 4; k++)
                   reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
                       c_rd[k] + c_wr[k];
@@ -1397,11 +1458,11 @@ __global__ void Marlin(
           __syncthreads();
         }
         if (red_idx == 0) {
-  #pragma unroll
+#pragma unroll
           for (int i = 0; i < 4 * 2; i += (m_block_size_8 ? 2 : 1)) {
             float* c_rd =
                 reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]);
-  #pragma unroll
+#pragma unroll
             for (int j = 0; j < 4; j++)
               reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
                   c_rd[j];
@@ -1444,7 +1505,7 @@ __global__ void Marlin(
 
     if (!first) {
 
-  #pragma unroll
+#pragma unroll
       for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
         int c_idx;
         if constexpr (m_block_size_8)
@@ -1461,11 +1522,11 @@ __global__ void Marlin(
       }
     }
 
-  #pragma unroll
+#pragma unroll
     for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
       if (!first) {
         int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta];
-  #pragma unroll
+#pragma unroll
         for (int j = 0; j < 2 * 4; j++) {
           int delta = 0;
           if constexpr (m_block_size_8) {
@@ -1478,7 +1539,7 @@ __global__ void Marlin(
       }
       if (!last) {
         int4 c;
-  #pragma unroll
+#pragma unroll
         for (int j = 0; j < 2 * 4; j++) {
           int delta = 0;
           if constexpr (m_block_size_8) {
@@ -1527,7 +1588,7 @@ __global__ void Marlin(
 
     if (!first) {
       float* frag_c_ptr = reinterpret_cast<float*>(&frag_c);
-  #pragma unroll
+#pragma unroll
       for (int k = 0; k < th_size; k++) {
         if constexpr (m_block_size_8) {
           if (k % 2) continue;
@@ -1540,7 +1601,7 @@ __global__ void Marlin(
             C_tmp[c_cur_offset + active_threads * k + threadIdx.x];
 
         float* sh_c_ptr = reinterpret_cast<float*>(&sh_red[threadIdx.x]);
-  #pragma unroll
+#pragma unroll
         for (int f = 0; f < 4; f++) {
           frag_c_ptr[k * 4 + f] += sh_c_ptr[f];
         }
@@ -1549,7 +1610,7 @@ __global__ void Marlin(
 
     if (!last) {
       int4* frag_c_ptr = reinterpret_cast<int4*>(&frag_c);
-  #pragma unroll
+#pragma unroll
       for (int k = 0; k < th_size; k++) {
         if constexpr (m_block_size_8) {
           if (k % 2) continue;
@@ -1619,26 +1680,38 @@ __global__ void Marlin(
     };
 
     if (threadIdx.x / 32 < thread_n_blocks / 4) {
-  #pragma unroll
+#pragma unroll
       for (int i = 0; i < thread_m_blocks; i++) {
-  #pragma unroll
+#pragma unroll
         for (int j = 0; j < 4; j++) {
           if constexpr (m_block_size_8) {
             int wr = c_sh_wr + 16 * j;
-            write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
+            write(wr,
+                  frag_c[i][j][0][0],
+                  frag_c[i][j][0][1],
                   frag_s[j / 2][2 * (j % 2) + 0]);
-            write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3],
+            write(wr + 8,
+                  frag_c[i][j][0][2],
+                  frag_c[i][j][0][3],
                   frag_s[j / 2][2 * (j % 2) + 1]);
           } else {
             int wr = c_sh_wr + 8 * j;
-            write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
-                  frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]);
-            write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
-                  frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]);
-            write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
-                  frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]);
-            write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
-                  frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]);
+            write(wr + (4 * c_sh_stride) * 0 + 0,
+                  frag_c[i][j][0][0],
+                  frag_c[i][j][0][1],
+                  frag_s[j / 2][2 * (j % 2) + 0]);
+            write(wr + (4 * c_sh_stride) * 8 + 0,
+                  frag_c[i][j][0][2],
+                  frag_c[i][j][0][3],
+                  frag_s[j / 2][2 * (j % 2) + 0]);
+            write(wr + (4 * c_sh_stride) * 0 + 4,
+                  frag_c[i][j][1][0],
+                  frag_c[i][j][1][1],
+                  frag_s[j / 2][2 * (j % 2) + 1]);
+            write(wr + (4 * c_sh_stride) * 8 + 4,
+                  frag_c[i][j][1][2],
+                  frag_c[i][j][1][3],
+                  frag_s[j / 2][2 * (j % 2) + 1]);
           }
         }
         c_sh_wr += 16 * (4 * c_sh_stride);
@@ -1646,7 +1719,7 @@ __global__ void Marlin(
     }
     __syncthreads();
 
-  #pragma unroll
+#pragma unroll
     for (int i = 0;
          i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
          i++) {
@@ -1660,7 +1733,7 @@ __global__ void Marlin(
           scalar_t2* C_half2 = reinterpret_cast<scalar_t2*>(&C[true_idx]);
           scalar_t2* sh_red_half2 =
               reinterpret_cast<scalar_t2*>(&sh_red[c_sh_rd]);
-  #pragma unroll
+#pragma unroll
           for (int a = 0; a < 4; a++) {
             scalar_t2 res = sh_red_half2[a];
             if (mul_topk_weights) {
@@ -1686,15 +1759,15 @@ __global__ void Marlin(
   // Start global fetch and register load pipelines.
   auto start_pipes = [&]() {
 
-  #pragma unroll
+#pragma unroll
     for (int i = 0; i < stages - 1; i++) {
       if (has_act_order && i == 0) {
         int last_g_idx = slice_k_start + stages * tb_k * 2;
         if (last_g_idx >= prob_k) {
           last_g_idx = prob_k - 1;
         }
-        fetch_act_order_scales_to_shared(true, g_idx[slice_k_start],
-                                         g_idx[last_g_idx]);
+        fetch_act_order_scales_to_shared(
+            true, g_idx[slice_k_start], g_idx[last_g_idx]);
       }
 
       if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
@@ -1732,9 +1805,9 @@ __global__ void Marlin(
 
     for (int stage_group_id = 0; stage_group_id < max_num_stage_groups;
          stage_group_id++) {
-  #pragma unroll
+#pragma unroll
       for (int pipe = 0; pipe < stages;) {
-  #pragma unroll
+#pragma unroll
         for (int k = 0; k < b_sh_wr_iters; k++) {
           int idx =
               (pipe >= stages && stage_group_id == max_num_stage_groups - 1)
@@ -1747,8 +1820,8 @@ __global__ void Marlin(
             int idx = (pipe >= 1 && stage_group_id == max_num_stage_groups - 1)
                           ? (pipe - 1)
                           : (pipe + (stage_group_id + 1) * stages - 1);
-            fetch_to_shared((pipe + stages - 1) % stages, pipe,
-                            slice_iters >= stages, idx);
+            fetch_to_shared(
+                (pipe + stages - 1) % stages, pipe, slice_iters >= stages, idx);
             pipe++;
             wait_for_stage();
             init_same_group(pipe % stages);
@@ -1775,8 +1848,8 @@ __global__ void Marlin(
           }
           int last_group_id = g_idx[last_g_idx];
           if (last_group_id >= sh_first_group_id + sh_num_groups) {
-            fetch_act_order_scales_to_shared(false, first_group_id,
-                                             last_group_id);
+            fetch_act_order_scales_to_shared(
+                false, first_group_id, last_group_id);
             __syncthreads();
           }
         }
@@ -1816,7 +1889,7 @@ __global__ void Marlin(
             if constexpr (m_block_size_8) {
               int idx = (threadIdx.x / 4) % 2;
               scalar_t2* frag_s_half2 = reinterpret_cast<scalar_t2*>(frag_s);
-  #pragma unroll
+#pragma unroll
               for (int i = 0; i < 8; i++) {
                 frag_s_half2[i] = Dtype::num2num2(
                     reinterpret_cast<scalar_t*>(&frag_s_half2[i])[idx]);
@@ -1833,9 +1906,9 @@ __global__ void Marlin(
                     w_type.size_bits() == 8 &&
                     (has_zp && dequant_skip_flop || !has_zp)) {
         if (threadIdx.x / 32 < thread_n_blocks / 4) {
-  #pragma unroll
+#pragma unroll
           for (int i = 0; i < thread_m_blocks; i++) {
-  #pragma unroll
+#pragma unroll
             for (int j = 0; j < 4; j++) {
               scale_float<scalar_t>(
                   reinterpret_cast<float*>(&frag_c[i][j][0][0]),
@@ -1896,11 +1969,11 @@ __global__ void Marlin(
 
       if (slice_iters) {
         a_gl_rd_col = (threadIdx.x % a_gl_rd_delta_o);
-  #pragma unroll
+#pragma unroll
         for (int i = 0; i < b_sh_wr_iters; i++)
           B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
         if (slice_col == 0) {
-  #pragma unroll
+#pragma unroll
           for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
         }
 
diff --git a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/types.h b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/types.h
index d905133390..79c51b4c8a 100644
--- a/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/types.h
+++ b/custom_ops/gpu_ops/moe/moe_wna16_marlin_utils/types.h
@@ -49,7 +49,6 @@ struct Tensor {
 
   int64_t size(int64_t d) const { return raw_tensor_.dims().at(d); }
 
-
   template <typename T>
   T *data_ptr() const {
     return const_cast<T *>(raw_tensor_.data<T>());
@@ -74,7 +73,6 @@ struct Tensor {
     return raw_tensor_.place().GetType() == phi::AllocationType::GPU;
   }
 
-
   // MARLIN_NAMESPACE_NAME::ScalarType scalar_type() const {
   //   return raw_tensor_.dtype();
   // }
diff --git a/custom_ops/gpu_ops/moe/swigluoai.cu b/custom_ops/gpu_ops/moe/swigluoai.cu
index 24ffacce0b..a6cd97a7c6 100644
--- a/custom_ops/gpu_ops/moe/swigluoai.cu
+++ b/custom_ops/gpu_ops/moe/swigluoai.cu
@@ -18,168 +18,171 @@
 
 #pragma once
 
-
 // dim3 grid(256)
 // dim3 block(512)
 template <typename T, int VecSize>
 __global__ void swigluoai_interleave_kernel(T* act_out,
-                                 const T* input,
-                                 const float alpha,
-                                 const float limit,
-                                 const int64_t seq_len,
-                                 const int64_t hidden_dim) {
-    int64_t tid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
-    const int64_t num = seq_len * hidden_dim;
-    using LoadT = AlignedVector<T, VecSize>;
-    LoadT src_vec0, src_vec1;
-    LoadT res_vec;
+                                            const T* input,
+                                            const float alpha,
+                                            const float limit,
+                                            const int64_t seq_len,
+                                            const int64_t hidden_dim) {
+  int64_t tid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const int64_t num = seq_len * hidden_dim;
+  using LoadT = AlignedVector<T, VecSize>;
+  LoadT src_vec0, src_vec1;
+  LoadT res_vec;
 
-    int64_t vec_num = hidden_dim / VecSize * seq_len;
-    int64_t col_size = hidden_dim / VecSize;
-    int64_t times = (vec_num - 1) / (gridDim.x * blockDim.x) + 1;
+  int64_t vec_num = hidden_dim / VecSize * seq_len;
+  int64_t col_size = hidden_dim / VecSize;
+  int64_t times = (vec_num - 1) / (gridDim.x * blockDim.x) + 1;
 
-    for(int i = 0; i < times; i++)
-    {
-        int64_t index = tid + i * gridDim.x * blockDim.x ;
-        int64_t row = index / col_size;
-        int64_t col = index % col_size;
+  for (int i = 0; i < times; i++) {
+    int64_t index = tid + i * gridDim.x * blockDim.x;
+    int64_t row = index / col_size;
+    int64_t col = index % col_size;
 
-        if(row < seq_len && col < col_size)
-        {
-            Load<T, VecSize>(&input[row*hidden_dim*2 + col*VecSize*2], &src_vec0);
-            Load<T, VecSize>(&input[row*hidden_dim*2 + col*VecSize*2 + VecSize], &src_vec1);
+    if (row < seq_len && col < col_size) {
+      Load<T, VecSize>(&input[row * hidden_dim * 2 + col * VecSize * 2],
+                       &src_vec0);
+      Load<T, VecSize>(
+          &input[row * hidden_dim * 2 + col * VecSize * 2 + VecSize],
+          &src_vec1);
 
-            for (int j = 0; j < VecSize/2; ++j) {
-                float a = static_cast<float>(src_vec0[2*j]);
-                float b = static_cast<float>(src_vec0[2*j + 1]);
-                a = fminf(a, limit);
-                b = fminf(fmaxf(b,-limit), limit);
-                float res = (b + 1) * a / (1.f + expf(-a * alpha));
-                res_vec[j] = static_cast<T>(res);
-            }
-            for (int j = 0; j < VecSize/2; ++j) {
-                float a = static_cast<float>(src_vec1[2*j]);
-                float b = static_cast<float>(src_vec1[2*j + 1]);
-                a = fminf(a, limit);
-                b = fminf(fmaxf(b,-limit), limit);
-                float res = (b + 1) * a / (1.f + expf(-a * alpha));
-                res_vec[j + VecSize/2] = static_cast<T>(res);
-            }
+      for (int j = 0; j < VecSize / 2; ++j) {
+        float a = static_cast<float>(src_vec0[2 * j]);
+        float b = static_cast<float>(src_vec0[2 * j + 1]);
+        a = fminf(a, limit);
+        b = fminf(fmaxf(b, -limit), limit);
+        float res = (b + 1) * a / (1.f + expf(-a * alpha));
+        res_vec[j] = static_cast<T>(res);
+      }
+      for (int j = 0; j < VecSize / 2; ++j) {
+        float a = static_cast<float>(src_vec1[2 * j]);
+        float b = static_cast<float>(src_vec1[2 * j + 1]);
+        a = fminf(a, limit);
+        b = fminf(fmaxf(b, -limit), limit);
+        float res = (b + 1) * a / (1.f + expf(-a * alpha));
+        res_vec[j + VecSize / 2] = static_cast<T>(res);
+      }
 
-            Store<T, VecSize>(res_vec, &act_out[row*hidden_dim + col*VecSize]);
-        }
+      Store<T, VecSize>(res_vec, &act_out[row * hidden_dim + col * VecSize]);
     }
+  }
 }
 
-
 // dim3 grid(256)
 // dim3 block(512)
 template <typename T, int VecSize>
 __global__ void swigluoai_norm_kernel(T* act_out,
-                                 const T* input,
-                                 const float alpha,
-                                 const float limit,
-                                 const int64_t seq_len,
-                                 const int64_t hidden_dim) {
-    int64_t tid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
-    const int64_t num = seq_len * hidden_dim;
-    using LoadT = AlignedVector<T, VecSize>;
-    LoadT src_vec0, src_vec1;
-    LoadT res_vec;
+                                      const T* input,
+                                      const float alpha,
+                                      const float limit,
+                                      const int64_t seq_len,
+                                      const int64_t hidden_dim) {
+  int64_t tid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const int64_t num = seq_len * hidden_dim;
+  using LoadT = AlignedVector<T, VecSize>;
+  LoadT src_vec0, src_vec1;
+  LoadT res_vec;
 
-    int64_t vec_num = hidden_dim / VecSize * seq_len;
-    int64_t col_size = hidden_dim / VecSize;
-    int64_t times = (vec_num - 1) / (gridDim.x * blockDim.x) + 1;
+  int64_t vec_num = hidden_dim / VecSize * seq_len;
+  int64_t col_size = hidden_dim / VecSize;
+  int64_t times = (vec_num - 1) / (gridDim.x * blockDim.x) + 1;
 
-    for(int i = 0; i < times; i++)
-    {
-        int64_t index = tid + i * gridDim.x * blockDim.x ;
-        int64_t row = index / col_size;
-        int64_t col = index % col_size;
+  for (int i = 0; i < times; i++) {
+    int64_t index = tid + i * gridDim.x * blockDim.x;
+    int64_t row = index / col_size;
+    int64_t col = index % col_size;
 
-        if(row < seq_len && col < col_size)
-        {
-            Load<T, VecSize>(&input[row*hidden_dim*2 + col*VecSize], &src_vec0);
-            Load<T, VecSize>(&input[row*hidden_dim*2 + hidden_dim + col*VecSize], &src_vec1);
+    if (row < seq_len && col < col_size) {
+      Load<T, VecSize>(&input[row * hidden_dim * 2 + col * VecSize], &src_vec0);
+      Load<T, VecSize>(
+          &input[row * hidden_dim * 2 + hidden_dim + col * VecSize], &src_vec1);
 
-            for (int j = 0; j < VecSize; ++j) {
-                float a = static_cast<float>(src_vec0[j]);
-                float b = static_cast<float>(src_vec1[j]);
-                float z = fminf(fmaxf(a * alpha, -limit), limit);
-                float res = b * a / (1.f + expf(-z));
-                res_vec[j] = static_cast<T>(res);
-            }
+      for (int j = 0; j < VecSize; ++j) {
+        float a = static_cast<float>(src_vec0[j]);
+        float b = static_cast<float>(src_vec1[j]);
+        float z = fminf(fmaxf(a * alpha, -limit), limit);
+        float res = b * a / (1.f + expf(-z));
+        res_vec[j] = static_cast<T>(res);
+      }
 
-            Store<T, VecSize>(res_vec, &act_out[row*hidden_dim + col*VecSize]);
-        }
+      Store<T, VecSize>(res_vec, &act_out[row * hidden_dim + col * VecSize]);
     }
+  }
 }
 
-paddle::Tensor SwigluOAI(const paddle::Tensor &fc1_out_tensor, const float alpha, const float limit, const std::string& type)
-{
-    // const int64_t group_size = fc1_out_tensor.shape()[1];
-    const int64_t seq_len = fc1_out_tensor.shape()[0];
-    const int64_t hidden_dim = fc1_out_tensor.shape()[1] / 2;
-    auto act_out_tensor = GetEmptyTensor({seq_len, hidden_dim}, fc1_out_tensor.dtype(), fc1_out_tensor.place());
+paddle::Tensor SwigluOAI(const paddle::Tensor& fc1_out_tensor,
+                         const float alpha,
+                         const float limit,
+                         const std::string& type) {
+  // const int64_t group_size = fc1_out_tensor.shape()[1];
+  const int64_t seq_len = fc1_out_tensor.shape()[0];
+  const int64_t hidden_dim = fc1_out_tensor.shape()[1] / 2;
+  auto act_out_tensor = GetEmptyTensor(
+      {seq_len, hidden_dim}, fc1_out_tensor.dtype(), fc1_out_tensor.place());
 
-    constexpr int VecSize = 8;
-    PD_CHECK(fc1_out_tensor.dtype() == paddle::DataType::BFLOAT16);
-    PD_CHECK(hidden_dim % VecSize == 0);
+  constexpr int VecSize = 8;
+  PD_CHECK(fc1_out_tensor.dtype() == paddle::DataType::BFLOAT16);
+  PD_CHECK(hidden_dim % VecSize == 0);
 
-    constexpr paddle::DataType D = paddle::DataType::BFLOAT16;
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+  constexpr paddle::DataType D = paddle::DataType::BFLOAT16;
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    const int block_size = 512;
-    const int grid_size = 256;
+  const int block_size = 512;
+  const int grid_size = 256;
 
-    #define dispatch_norm() do {\
-    swigluoai_norm_kernel<DataType_, VecSize><<<grid_size, block_size, 0, fc1_out_tensor.stream()>>>(\
-        reinterpret_cast<DataType_*>(const_cast<data_t*>(act_out_tensor.data<data_t>())),\
-        reinterpret_cast<const DataType_*>(fc1_out_tensor.data<data_t>()),\
-        alpha,\
-        limit,\
-        seq_len,\
-        hidden_dim\
-    );} while(0)
+#define dispatch_norm()                                                        \
+  do {                                                                         \
+    swigluoai_norm_kernel<DataType_, VecSize>                                  \
+        <<<grid_size, block_size, 0, fc1_out_tensor.stream()>>>(               \
+            reinterpret_cast<DataType_*>(                                      \
+                const_cast<data_t*>(act_out_tensor.data<data_t>())),           \
+            reinterpret_cast<const DataType_*>(fc1_out_tensor.data<data_t>()), \
+            alpha,                                                             \
+            limit,                                                             \
+            seq_len,                                                           \
+            hidden_dim);                                                       \
+  } while (0)
 
-    #define dispatch_interleave() do {\
-    swigluoai_interleave_kernel<DataType_, VecSize><<<grid_size, block_size, 0, fc1_out_tensor.stream()>>>(\
-        reinterpret_cast<DataType_*>(const_cast<data_t*>(act_out_tensor.data<data_t>())),\
-        reinterpret_cast<const DataType_*>(fc1_out_tensor.data<data_t>()),\
-        alpha,\
-        limit,\
-        seq_len,\
-        hidden_dim\
-    );} while(0)
+#define dispatch_interleave()                                                  \
+  do {                                                                         \
+    swigluoai_interleave_kernel<DataType_, VecSize>                            \
+        <<<grid_size, block_size, 0, fc1_out_tensor.stream()>>>(               \
+            reinterpret_cast<DataType_*>(                                      \
+                const_cast<data_t*>(act_out_tensor.data<data_t>())),           \
+            reinterpret_cast<const DataType_*>(fc1_out_tensor.data<data_t>()), \
+            alpha,                                                             \
+            limit,                                                             \
+            seq_len,                                                           \
+            hidden_dim);                                                       \
+  } while (0)
 
-    if(type == "interleave")
-    {
-        dispatch_interleave();
-    }
-    else
-    {
-        dispatch_norm();
-    }
-    // if (token_nums_per_expert.dtype() == paddle::DataType::INT64) {
-    //     dispatch_by_index(int64_t);
-    // } else if(token_nums_per_expert.dtype() == paddle::DataType::INT32) {
-    //     dispatch_by_index(int32_t);
-    // } else {
-    //     PD_THROW("Unsupported token_nums_per_expert's data dtype.");
-    // }
+  if (type == "interleave") {
+    dispatch_interleave();
+  } else {
+    dispatch_norm();
+  }
+  // if (token_nums_per_expert.dtype() == paddle::DataType::INT64) {
+  //     dispatch_by_index(int64_t);
+  // } else if(token_nums_per_expert.dtype() == paddle::DataType::INT32) {
+  //     dispatch_by_index(int32_t);
+  // } else {
+  //     PD_THROW("Unsupported token_nums_per_expert's data dtype.");
+  // }
 
-    return act_out_tensor;
+  return act_out_tensor;
 }
 
-
 std::vector<paddle::Tensor> SwigluOAIWrapper(
     const paddle::Tensor& fc1_out_tensor,
     const float alpha,
     const float limit,
     const std::string& type) {
-     return {SwigluOAI(fc1_out_tensor, alpha, limit, type)};
+  return {SwigluOAI(fc1_out_tensor, alpha, limit, type)};
 }
 
 PD_BUILD_STATIC_OP(swigluoai)
diff --git a/custom_ops/gpu_ops/moe/swigluoai.h b/custom_ops/gpu_ops/moe/swigluoai.h
index 63e8c850e3..45f1b18978 100644
--- a/custom_ops/gpu_ops/moe/swigluoai.h
+++ b/custom_ops/gpu_ops/moe/swigluoai.h
@@ -15,5 +15,7 @@
 #pragma once
 #include "helper.h"
 
-paddle::Tensor
-SwigluOAI(const paddle::Tensor &fc1_out_tensor, const float alpha, const float limit, const std::string& type);
+paddle::Tensor SwigluOAI(const paddle::Tensor& fc1_out_tensor,
+                         const float alpha,
+                         const float limit,
+                         const std::string& type);
diff --git a/custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu b/custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu
index 1faf8b4cd5..071e0a9b41 100644
--- a/custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu
+++ b/custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu
@@ -15,13 +15,14 @@
 #include "helper.h"
 #include "paddle/extension.h"
 
-#define CEILDIV(a,b) (((a+b-1)/b))
+#define CEILDIV(a, b) (((a + b - 1) / b))
 
 template <typename scalar_t>
-__global__ void count_and_sort_expert_tokens_kernel(const scalar_t* __restrict__ topk_ids,
-                                                    int32_t* __restrict__ sorted_token_ids,
-                                                    int32_t* __restrict__ cumsum_buffer,
-                                                    size_t numel) {
+__global__ void count_and_sort_expert_tokens_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ cumsum_buffer,
+    size_t numel) {
   const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
   const size_t stride = blockDim.x * gridDim.x;
 
@@ -33,16 +34,17 @@ __global__ void count_and_sort_expert_tokens_kernel(const scalar_t* __restrict__
 }
 
 template <typename scalar_t, int num_experts>
-__global__ void moe_align_block_size_kernel(const scalar_t* __restrict__ topk_ids,
-                                            int32_t* __restrict__ expert_ids,
-                                            int32_t* __restrict__ total_tokens_post_pad,
-                                            int32_t GEMM_BLOCK_SIZE_M,
-                                            size_t numel,
-                                            int32_t* __restrict__ cumsum_buffer) {
+__global__ void moe_align_block_size_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t GEMM_BLOCK_SIZE_M,
+    size_t numel,
+    int32_t* __restrict__ cumsum_buffer) {
   __shared__ int32_t tokens_per_ep[num_experts];
 
   for (int i = threadIdx.x; i < num_experts; i += blockDim.x) {
-      tokens_per_ep[i] = 0;
+    tokens_per_ep[i] = 0;
   }
 
   __syncthreads();
@@ -57,8 +59,10 @@ __global__ void moe_align_block_size_kernel(const scalar_t* __restrict__ topk_id
   if (threadIdx.x == 0) {
     cumsum_buffer[0] = 0;
     for (int i = 1; i <= num_experts; ++i) {
-      int expert_count = tokens_per_ep[i-1];
-      cumsum_buffer[i] = cumsum_buffer[i - 1] + CEILDIV(expert_count, GEMM_BLOCK_SIZE_M) * GEMM_BLOCK_SIZE_M;
+      int expert_count = tokens_per_ep[i - 1];
+      cumsum_buffer[i] =
+          cumsum_buffer[i - 1] +
+          CEILDIV(expert_count, GEMM_BLOCK_SIZE_M) * GEMM_BLOCK_SIZE_M;
     }
     *total_tokens_post_pad = cumsum_buffer[num_experts];
   }
@@ -66,33 +70,39 @@ __global__ void moe_align_block_size_kernel(const scalar_t* __restrict__ topk_id
   __syncthreads();
 
   if (threadIdx.x < num_experts) {
-    for (int i = cumsum_buffer[threadIdx.x]; i < cumsum_buffer[threadIdx.x + 1]; i += GEMM_BLOCK_SIZE_M) {
+    for (int i = cumsum_buffer[threadIdx.x]; i < cumsum_buffer[threadIdx.x + 1];
+         i += GEMM_BLOCK_SIZE_M) {
       expert_ids[i / GEMM_BLOCK_SIZE_M] = threadIdx.x;
     }
   }
 }
 
+std::vector<std::vector<int64_t>> tritonmoe_preprocessInferShape(
+    const std::vector<int64_t>& topk_ids,
+    int64_t num_experts,
+    int64_t GEMM_BLOCK_SIZE_M) {
+  int topk_ids_numel = topk_ids[0] * topk_ids[1];
+  int max_num_tokens_padded =
+      topk_ids_numel + num_experts * (GEMM_BLOCK_SIZE_M - 1);
 
-std::vector<std::vector<int64_t>> tritonmoe_preprocessInferShape(const std::vector<int64_t>& topk_ids, int64_t num_experts, int64_t GEMM_BLOCK_SIZE_M) {
+  std::vector<int64_t> sorted_ids = {max_num_tokens_padded};
 
+  int max_num_m_blocks = max_num_tokens_padded / GEMM_BLOCK_SIZE_M;
+  std::vector<int64_t> expert_ids = {max_num_m_blocks};
+  std::vector<int64_t> num_tokens_post_pad = {1};
 
-    int topk_ids_numel = topk_ids[0] * topk_ids[1];
-    int max_num_tokens_padded = topk_ids_numel + num_experts * (GEMM_BLOCK_SIZE_M - 1);
-
-    std::vector<int64_t> sorted_ids = {max_num_tokens_padded};
-
-    int max_num_m_blocks = max_num_tokens_padded / GEMM_BLOCK_SIZE_M;
-    std::vector<int64_t> expert_ids = {max_num_m_blocks};
-    std::vector<int64_t> num_tokens_post_pad = {1};
-
-    return {sorted_ids, expert_ids, num_tokens_post_pad};
+  return {sorted_ids, expert_ids, num_tokens_post_pad};
 }
 
-std::vector<paddle::DataType> tritonmoe_preprocessIferDtype(const paddle::DataType& topk_ids, int64_t num_experts, int64_t GEMM_BLOCK_SIZE_M) {
-    return {paddle::DataType::INT32, paddle::DataType::INT32, paddle::DataType::INT32};
+std::vector<paddle::DataType> tritonmoe_preprocessIferDtype(
+    const paddle::DataType& topk_ids,
+    int64_t num_experts,
+    int64_t GEMM_BLOCK_SIZE_M) {
+  return {paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32};
 }
 
-
 /*
 supporse num_experts = 8, GEMM_BLOCK_SIZE_M = 4,
 topk_ids.shape = [4,4], means=topk=4
@@ -113,85 +123,74 @@ Then return value `sorted_ids` is
 0,16,16,16
 */
 
+std::vector<paddle::Tensor> tritonmoe_preprocess_kernel(
+    const paddle::Tensor& topk_ids,
+    int64_t num_experts,
+    int64_t GEMM_BLOCK_SIZE_M) {
+  int topk_ids_numel = topk_ids.shape()[0] * topk_ids.shape()[1];
+  int max_num_tokens_padded =
+      topk_ids_numel + num_experts * (GEMM_BLOCK_SIZE_M - 1);
 
-std::vector<paddle::Tensor> tritonmoe_preprocess_kernel(const paddle::Tensor& topk_ids, int64_t num_experts, int64_t GEMM_BLOCK_SIZE_M) {
+  auto sorted_ids = paddle::full({max_num_tokens_padded},
+                                 topk_ids_numel,
+                                 paddle::DataType::INT32,
+                                 topk_ids.place());
 
-    int topk_ids_numel = topk_ids.shape()[0] * topk_ids.shape()[1];
-    int max_num_tokens_padded = topk_ids_numel + num_experts * (GEMM_BLOCK_SIZE_M - 1);
+  int max_num_m_blocks = max_num_tokens_padded / GEMM_BLOCK_SIZE_M;
 
-    auto sorted_ids = paddle::full(
-        {max_num_tokens_padded},
-        topk_ids_numel,
-        paddle::DataType::INT32,
-        topk_ids.place()
-    );
+  auto expert_ids = paddle::empty(
+      {max_num_m_blocks}, paddle::DataType::INT32, topk_ids.place());
 
-    int max_num_m_blocks = max_num_tokens_padded / GEMM_BLOCK_SIZE_M;
+  auto num_tokens_post_pad =
+      paddle::empty({1}, paddle::DataType::INT32, topk_ids.place());
 
-    auto expert_ids = paddle::empty(
-        {max_num_m_blocks}, paddle::DataType::INT32,
-        topk_ids.place()
-    );
+  auto cumsum_buffer = paddle::empty(
+      {num_experts + 1}, paddle::DataType::INT32, topk_ids.place());
 
-    auto num_tokens_post_pad = paddle::empty(
-        {1},
-        paddle::DataType::INT32,
-         topk_ids.place()
-    );
+  auto stream = topk_ids.stream();
+  using scalar_t = int64_t;
 
-    auto cumsum_buffer = paddle::empty(
-        {num_experts + 1},
-        paddle::DataType::INT32,
-        topk_ids.place()
-    );
+#define run_align_kernel(num_experts)                                       \
+  auto align_kernel = moe_align_block_size_kernel<scalar_t, num_experts>;   \
+  align_kernel<<<1, 1024, 0, stream>>>(topk_ids.data<scalar_t>(),           \
+                                       expert_ids.data<int32_t>(),          \
+                                       num_tokens_post_pad.data<int32_t>(), \
+                                       GEMM_BLOCK_SIZE_M,                   \
+                                       topk_ids_numel,                      \
+                                       cumsum_buffer.data<int32_t>());
 
-    auto stream = topk_ids.stream();
-    using scalar_t = int64_t;
+  if (num_experts == 8) {
+    run_align_kernel(8);
+  } else if (num_experts == 256) {
+    run_align_kernel(256);
+  } else if (num_experts == 2) {
+    run_align_kernel(2);
+  } else if (num_experts == 64) {
+    run_align_kernel(64);
+  } else if (num_experts == 128) {
+    run_align_kernel(128);
+  } else if (num_experts == 160) {
+    run_align_kernel(160);
+  } else if (num_experts == 32) {
+    run_align_kernel(32);
+  } else {
+    PD_THROW("Not support num_experts: %d", num_experts);
+  }
 
-    # define run_align_kernel(num_experts) \
-    auto align_kernel = moe_align_block_size_kernel<scalar_t, num_experts>; \
-    align_kernel<<<1, 1024, 0, stream>>>( \
-    topk_ids.data<scalar_t>(),  \
-    expert_ids.data<int32_t>(), \
-    num_tokens_post_pad.data<int32_t>(), \
-    GEMM_BLOCK_SIZE_M,  \
-    topk_ids_numel, \
-    cumsum_buffer.data<int32_t>());
+  const int block_threads = 256;
+  const int num_blocks = CEILDIV(topk_ids_numel, block_threads);
+  const int max_blocks = 65535;
+  const int actual_blocks = std::min(num_blocks, max_blocks);
 
-    if (num_experts == 8) {
-      run_align_kernel(8);
-    } else if (num_experts == 256) {
-      run_align_kernel(256);
-    } else if (num_experts == 2) {
-      run_align_kernel(2);
-    } else if (num_experts == 64) {
-      run_align_kernel(64);
-    } else if (num_experts == 128) {
-      run_align_kernel(128);
-    } else if (num_experts == 160) {
-      run_align_kernel(160);
-    } else if (num_experts == 32) {
-      run_align_kernel(32);
-    }
-    else {
-      PD_THROW("Not support num_experts: %d", num_experts);
-    }
+  auto sort_kernel = count_and_sort_expert_tokens_kernel<scalar_t>;
 
-    const int block_threads = 256;
-    const int num_blocks = CEILDIV(topk_ids_numel, block_threads);
-    const int max_blocks = 65535;
-    const int actual_blocks = std::min(num_blocks, max_blocks);
+  sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(
+      topk_ids.data<scalar_t>(),
+      sorted_ids.data<int32_t>(),
+      cumsum_buffer.data<int32_t>(),
+      topk_ids_numel);
 
-    auto sort_kernel = count_and_sort_expert_tokens_kernel<scalar_t>;
-
-    sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(topk_ids.data<scalar_t>(),
-                                                              sorted_ids.data<int32_t>(),
-                                                              cumsum_buffer.data<int32_t>(),
-                                                              topk_ids_numel);
-
-
-
-    return {sorted_ids, expert_ids, num_tokens_post_pad};
+  return {sorted_ids, expert_ids, num_tokens_post_pad};
 }
 
 PD_BUILD_STATIC_OP(tritonmoe_preprocess)
diff --git a/custom_ops/gpu_ops/moe/winx_unzip.cu b/custom_ops/gpu_ops/moe/winx_unzip.cu
index 22bebf1ef7..2c96e47979 100644
--- a/custom_ops/gpu_ops/moe/winx_unzip.cu
+++ b/custom_ops/gpu_ops/moe/winx_unzip.cu
@@ -17,12 +17,17 @@
 
 template <typename T, int TileRows, int TileColumns, int NumThreads>
 __global__ void Wint25UnzipKernel(const uint16_t *zipped_weight_ptr,
-                                  const T *super_scale_ptr, T *weight_ptr,
-                                  const int64_t batch, const int64_t num_rows,
+                                  const T *super_scale_ptr,
+                                  T *weight_ptr,
+                                  const int64_t batch,
+                                  const int64_t num_rows,
                                   const int64_t num_columns) {
-  using UnzipFunctor =
-      cutlass::gemm::threadblock::UnzipAndDequantFunctor<T, cutlass::WintQuantMethod::kWeightOnlyInt25, TileRows,
-                             TileColumns, NumThreads>;
+  using UnzipFunctor = cutlass::gemm::threadblock::UnzipAndDequantFunctor<
+      T,
+      cutlass::WintQuantMethod::kWeightOnlyInt25,
+      TileRows,
+      TileColumns,
+      NumThreads>;
 
   __shared__ T smem[TileRows * TileColumns];
 
@@ -41,7 +46,8 @@ __global__ void Wint25UnzipKernel(const uint16_t *zipped_weight_ptr,
 
   // unzip to shared memory
   UnzipFunctor unzip_functor;
-  unzip_functor(block_zipped_weight_ptr, block_super_scale_ptr, smem, num_columns);
+  unzip_functor(
+      block_zipped_weight_ptr, block_super_scale_ptr, smem, num_columns);
 
   // write back to global memory
   for (int row = 0; row < TileRows; ++row) {
@@ -55,19 +61,26 @@ __global__ void Wint25UnzipKernel(const uint16_t *zipped_weight_ptr,
 }
 
 template <typename T, int64_t TileRows, int64_t TileColumns, int NumThreads>
-__global__ void
-Wint2UnzipKernel(const uint8_t *zipped_weight_ptr,
-                 const uint8_t *local_scale_ptr, const float *code_scale_ptr,
-                 const float *code_zp_ptr, const T *super_scale_ptr,
-                 T *weight_ptr, const int64_t batch, const int64_t num_rows,
-                 const int64_t num_columns) {
-  using UnzipFunctor =
-      cutlass::gemm::threadblock::UnzipAndDequantFunctor<T, cutlass::WintQuantMethod::kWeightOnlyInt2, TileRows,
-                             TileColumns, NumThreads>;
+__global__ void Wint2UnzipKernel(const uint8_t *zipped_weight_ptr,
+                                 const uint8_t *local_scale_ptr,
+                                 const float *code_scale_ptr,
+                                 const float *code_zp_ptr,
+                                 const T *super_scale_ptr,
+                                 T *weight_ptr,
+                                 const int64_t batch,
+                                 const int64_t num_rows,
+                                 const int64_t num_columns) {
+  using UnzipFunctor = cutlass::gemm::threadblock::UnzipAndDequantFunctor<
+      T,
+      cutlass::WintQuantMethod::kWeightOnlyInt2,
+      TileRows,
+      TileColumns,
+      NumThreads>;
 
   constexpr bool kUseAsyncLoad = true;
 
-  __shared__ uint8_t zipped_smem[UnzipFunctor::kZippedSmemBytes + UnzipFunctor::kColumnWiseSmemBytes];
+  __shared__ uint8_t zipped_smem[UnzipFunctor::kZippedSmemBytes +
+                                 UnzipFunctor::kColumnWiseSmemBytes];
   __shared__ T smem[TileRows * TileColumns];
 
   int64_t block_start_column = blockIdx.x * TileColumns;
@@ -95,15 +108,21 @@ Wint2UnzipKernel(const uint8_t *zipped_weight_ptr,
           ? super_scale_ptr + blockIdx.z * num_columns + block_start_column
           : nullptr;
 
-  typename UnzipFunctor::Arguments args(zipped_smem, zipped_smem + UnzipFunctor::kZippedSmemBytes);
+  typename UnzipFunctor::Arguments args(
+      zipped_smem, zipped_smem + UnzipFunctor::kZippedSmemBytes);
 
   // unzip to shared memory
   UnzipFunctor functor;
 
   if (kUseAsyncLoad) {
-    functor.LoadAsync(block_zipped_weight_ptr, block_local_scale_ptr,
-                      block_code_scale_ptr, block_code_zp_ptr, block_super_scale_ptr,
-                      &args, num_columns, true);
+    functor.LoadAsync(block_zipped_weight_ptr,
+                      block_local_scale_ptr,
+                      block_code_scale_ptr,
+                      block_code_zp_ptr,
+                      block_super_scale_ptr,
+                      &args,
+                      num_columns,
+                      true);
 
     // 发起 cp.async 的收束
     cutlass::arch::cp_async_fence();
@@ -112,9 +131,14 @@ Wint2UnzipKernel(const uint8_t *zipped_weight_ptr,
     cutlass::arch::cp_async_wait<0>();
     __syncthreads();
   } else {
-    functor.Load(block_zipped_weight_ptr, block_local_scale_ptr,
-                 block_code_scale_ptr, block_code_zp_ptr, block_super_scale_ptr,
-                 &args, num_columns, true);
+    functor.Load(block_zipped_weight_ptr,
+                 block_local_scale_ptr,
+                 block_code_scale_ptr,
+                 block_code_zp_ptr,
+                 block_super_scale_ptr,
+                 &args,
+                 num_columns,
+                 true);
   }
 
   functor.Compute(args, smem, block_start_row);
@@ -132,8 +156,10 @@ Wint2UnzipKernel(const uint8_t *zipped_weight_ptr,
 
 template <typename T>
 void Wint25UnzipKernelLauncher(const uint16_t *zipped_weight,
-                               const T *supper_scale, T *weight,
-                               const int64_t batch, const int64_t num_rows,
+                               const T *supper_scale,
+                               T *weight,
+                               const int64_t batch,
+                               const int64_t num_rows,
                                const int64_t num_columns) {
   constexpr int kTileRows = 64;
   constexpr int kTileColumns = 128;
@@ -146,16 +172,19 @@ void Wint25UnzipKernelLauncher(const uint16_t *zipped_weight,
   dim3 grid_dim(block_dim_x, block_dim_y, batch);
 
   Wint25UnzipKernel<T, kTileRows, kTileColumns, kNumThreads>
-      <<<grid_dim, block_dim>>>(zipped_weight, supper_scale, weight, batch,
-                                num_rows, num_columns);
+      <<<grid_dim, block_dim>>>(
+          zipped_weight, supper_scale, weight, batch, num_rows, num_columns);
 }
 
 template <typename T>
 void Wint2UnzipKernelLauncher(const uint8_t *zipped_weight,
                               const uint8_t *local_scale,
-                              const float *code_scale, const float *code_zp,
-                              const T *supper_scale, T *weight,
-                              const int64_t batch, const int64_t num_rows,
+                              const float *code_scale,
+                              const float *code_zp,
+                              const T *supper_scale,
+                              T *weight,
+                              const int64_t batch,
+                              const int64_t num_rows,
                               const int64_t num_columns) {
   constexpr int kTileRows = 64;
   constexpr int kTileColumns = 256;
@@ -168,8 +197,14 @@ void Wint2UnzipKernelLauncher(const uint8_t *zipped_weight,
   dim3 grid_dim(block_dim_x, block_dim_y, batch);
 
   Wint2UnzipKernel<T, kTileRows, kTileColumns, kNumThreads>
-      <<<grid_dim, block_dim>>>(zipped_weight, local_scale, code_scale, code_zp,
-                                supper_scale, weight, batch, num_rows,
+      <<<grid_dim, block_dim>>>(zipped_weight,
+                                local_scale,
+                                code_scale,
+                                code_zp,
+                                supper_scale,
+                                weight,
+                                batch,
+                                num_rows,
                                 num_columns);
 }
 
@@ -179,7 +214,8 @@ void WintxUnzipKernel(const paddle::Tensor &zipped_weight,
                       const paddle::optional<paddle::Tensor> &code_scale,
                       const paddle::optional<paddle::Tensor> &code_zp,
                       const paddle::optional<paddle::Tensor> &super_scale,
-                      paddle::Tensor &weight, const std::string &quant_method) {
+                      paddle::Tensor &weight,
+                      const std::string &quant_method) {
   using data_t = typename PDTraits<T>::data_t;
   using NvType = typename PDTraits<T>::DataType;
 
@@ -199,7 +235,10 @@ void WintxUnzipKernel(const paddle::Tensor &zipped_weight,
     Wint25UnzipKernelLauncher<NvType>(
         reinterpret_cast<const uint16_t *>(zipped_weight_ptr),
         reinterpret_cast<const NvType *>(super_scale_ptr),
-        reinterpret_cast<NvType *>(weight_ptr), batch, num_rows, num_columns);
+        reinterpret_cast<NvType *>(weight_ptr),
+        batch,
+        num_rows,
+        num_columns);
   } else if (quant_method == "weight_only_int2") {
     paddle::Tensor *local_scale_tensor =
         const_cast<paddle::Tensor *>(local_scale.get_ptr());
@@ -209,22 +248,27 @@ void WintxUnzipKernel(const paddle::Tensor &zipped_weight,
         const_cast<paddle::Tensor *>(code_zp.get_ptr());
 
     Wint2UnzipKernelLauncher<NvType>(
-        zipped_weight.data<uint8_t>(), local_scale_tensor->data<uint8_t>(),
-        code_scale_tensor->data<float>(), code_zp_tensor->data<float>(),
+        zipped_weight.data<uint8_t>(),
+        local_scale_tensor->data<uint8_t>(),
+        code_scale_tensor->data<float>(),
+        code_zp_tensor->data<float>(),
         reinterpret_cast<const NvType *>(super_scale_ptr),
-        reinterpret_cast<NvType *>(weight_ptr), batch, num_rows, num_columns);
+        reinterpret_cast<NvType *>(weight_ptr),
+        batch,
+        num_rows,
+        num_columns);
   } else {
     PD_THROW("Unsupported quant_method for WintxUnzip.");
   }
 }
 
-std::vector<paddle::Tensor>
-WintXUnzip(const paddle::Tensor &zipped_weight,
-           const paddle::optional<paddle::Tensor> &local_scale,
-           const paddle::optional<paddle::Tensor> &code_scale,
-           const paddle::optional<paddle::Tensor> &code_zp,
-           const paddle::optional<paddle::Tensor> &super_scale,
-           const std::string &quant_method) {
+std::vector<paddle::Tensor> WintXUnzip(
+    const paddle::Tensor &zipped_weight,
+    const paddle::optional<paddle::Tensor> &local_scale,
+    const paddle::optional<paddle::Tensor> &code_scale,
+    const paddle::optional<paddle::Tensor> &code_zp,
+    const paddle::optional<paddle::Tensor> &super_scale,
+    const std::string &quant_method) {
   paddle::Tensor *local_scale_tensor =
       const_cast<paddle::Tensor *>(local_scale.get_ptr());
   paddle::Tensor *super_scale_tensor =
@@ -251,18 +295,26 @@ WintXUnzip(const paddle::Tensor &zipped_weight,
   auto output_tensor = GetEmptyTensor(output_dims, dtype, place);
 
   switch (dtype) {
-  case paddle::DataType::BFLOAT16:
-    WintxUnzipKernel<paddle::DataType::BFLOAT16>(
-        zipped_weight, local_scale, code_scale, code_zp, super_scale,
-        output_tensor, quant_method);
-    break;
-  case paddle::DataType::FLOAT16:
-    WintxUnzipKernel<paddle::DataType::FLOAT16>(
-        zipped_weight, local_scale, code_scale, code_zp, super_scale,
-        output_tensor, quant_method);
-    break;
-  default:
-    PD_THROW("Unsupported data type for WintxUnzip");
+    case paddle::DataType::BFLOAT16:
+      WintxUnzipKernel<paddle::DataType::BFLOAT16>(zipped_weight,
+                                                   local_scale,
+                                                   code_scale,
+                                                   code_zp,
+                                                   super_scale,
+                                                   output_tensor,
+                                                   quant_method);
+      break;
+    case paddle::DataType::FLOAT16:
+      WintxUnzipKernel<paddle::DataType::FLOAT16>(zipped_weight,
+                                                  local_scale,
+                                                  code_scale,
+                                                  code_zp,
+                                                  super_scale,
+                                                  output_tensor,
+                                                  quant_method);
+      break;
+    default:
+      PD_THROW("Unsupported data type for WintxUnzip");
   }
   return {output_tensor};
 }
@@ -306,8 +358,10 @@ std::vector<paddle::DataType> WintXUnzipInferDtype(
 }
 
 PD_BUILD_STATIC_OP(winx_unzip)
-    .Inputs({"zipped_weight", paddle::Optional("local_scale"),
-             paddle::Optional("code_scale"), paddle::Optional("code_zp"),
+    .Inputs({"zipped_weight",
+             paddle::Optional("local_scale"),
+             paddle::Optional("code_scale"),
+             paddle::Optional("code_zp"),
              paddle::Optional("super_scale")})
     .Outputs({"weight"})
     .Attrs({"quant_method:std::string"})
diff --git a/custom_ops/gpu_ops/msg_utils.h b/custom_ops/gpu_ops/msg_utils.h
index ff46ccb004..9976ee1be4 100644
--- a/custom_ops/gpu_ops/msg_utils.h
+++ b/custom_ops/gpu_ops/msg_utils.h
@@ -29,11 +29,11 @@
 #define MAX_BSZ 512
 
 struct msgdata {
-    long mtype;
-    int mtext[MAX_BSZ + 2];  // stop_flag, bsz, tokens
+  long mtype;
+  int mtext[MAX_BSZ + 2];  // stop_flag, bsz, tokens
 };
 
 struct msgdatakv {
-    long mtype;
-    int mtext[MAX_BSZ * 3 + 2];  // encoder_count, layer_id, bid- pair
+  long mtype;
+  int mtext[MAX_BSZ * 3 + 2];  // encoder_count, layer_id, bid- pair
 };
diff --git a/custom_ops/gpu_ops/ngram_mask.cu b/custom_ops/gpu_ops/ngram_mask.cu
index 7260244c5f..d712aad6f4 100644
--- a/custom_ops/gpu_ops/ngram_mask.cu
+++ b/custom_ops/gpu_ops/ngram_mask.cu
@@ -40,178 +40,178 @@ class PDTraits;
 
 template <>
 class PDTraits<paddle::DataType::FLOAT32> {
-    public:
-    typedef float DataType;
-    typedef float data_t;
+ public:
+  typedef float DataType;
+  typedef float data_t;
 };
 
 template <>
 class PDTraits<paddle::DataType::FLOAT16> {
-    public:
-    typedef half DataType;
-    typedef paddle::float16 data_t;
+ public:
+  typedef half DataType;
+  typedef paddle::float16 data_t;
 };
 
 class TreeNode {
-    public:
-    int token_id_;
-    TreeNode *children_;
-    int children_node_len_;
+ public:
+  int token_id_;
+  TreeNode *children_;
+  int children_node_len_;
 
-    TreeNode() : token_id_(-1), children_node_len_(0) { children_ = nullptr; }
+  TreeNode() : token_id_(-1), children_node_len_(0) { children_ = nullptr; }
 
-    __host__ __device__ bool is_in_tree(TreeNode *node,
-                                        int node_num,
-                                        int token,
-                                        int *idx) {
-        for (int i = 0; i < node_num; i++) {
-            if (node->children_) {
-                TreeNode *tmp = &node->children_[i];
-                if (tmp->token_id_ == token) {
-                    *idx = i;
-                    return true;
-                } else if (tmp->token_id_ != -1) {
-                    continue;
-                } else {
-                    *idx = i;
-                    break;
-                }
-            } else {
-                *idx = i;
-                break;
-            }
+  __host__ __device__ bool is_in_tree(TreeNode *node,
+                                      int node_num,
+                                      int token,
+                                      int *idx) {
+    for (int i = 0; i < node_num; i++) {
+      if (node->children_) {
+        TreeNode *tmp = &node->children_[i];
+        if (tmp->token_id_ == token) {
+          *idx = i;
+          return true;
+        } else if (tmp->token_id_ != -1) {
+          continue;
+        } else {
+          *idx = i;
+          break;
         }
-        return false;
+      } else {
+        *idx = i;
+        break;
+      }
     }
+    return false;
+  }
 
-    __host__ __device__ void insert(const int *stop_list,
-                                    const int stop_list_len) {
-        TreeNode *node = this;
-        for (int i = 0; i < stop_list_len; i++) {
-            int token = stop_list[i];
-            if (token == -1) break;
-            int idx;
-            bool in_tree = is_in_tree(node, STOP_LIST_BS, token, &idx);
-            if (!in_tree) {
-                node->children_node_len_++;
-                if (!node->children_) {
-                    node->children_ = new TreeNode[STOP_LIST_BS];
-                }
-                node->children_[idx].token_id_ = token;
-            }
-            node = &(node->children_[idx]);
+  __host__ __device__ void insert(const int *stop_list,
+                                  const int stop_list_len) {
+    TreeNode *node = this;
+    for (int i = 0; i < stop_list_len; i++) {
+      int token = stop_list[i];
+      if (token == -1) break;
+      int idx;
+      bool in_tree = is_in_tree(node, STOP_LIST_BS, token, &idx);
+      if (!in_tree) {
+        node->children_node_len_++;
+        if (!node->children_) {
+          node->children_ = new TreeNode[STOP_LIST_BS];
         }
+        node->children_[idx].token_id_ = token;
+      }
+      node = &(node->children_[idx]);
     }
+  }
 
-    __host__ __device__ void search(const int *prefix_token,
-                                    int token_len,
-                                    int *res,
-                                    int *res_len) {
-        TreeNode *node = this;
-        int idx;
-        for (int i = 0; i < token_len; i++) {
-            if (node) {
-                bool in_tree = is_in_tree(
-                    node, node->children_node_len_, prefix_token[i], &idx);
-                if (in_tree) {
-                    node = &(node->children_[idx]);
-                } else {
-                    *res_len = 0;
-                    return;
-                }
-            } else {
-                break;
-            }
-        }
-        if (node) {
-            int id = 0;
-            for (int i = 0; i < node->children_node_len_; i++) {
-                if (node->children_[i].token_id_ != -1) {
-                    res[id++] = node->children_[i].token_id_;
-                } else {
-                    break;
-                }
-            }
-            *res_len = id;
+  __host__ __device__ void search(const int *prefix_token,
+                                  int token_len,
+                                  int *res,
+                                  int *res_len) {
+    TreeNode *node = this;
+    int idx;
+    for (int i = 0; i < token_len; i++) {
+      if (node) {
+        bool in_tree =
+            is_in_tree(node, node->children_node_len_, prefix_token[i], &idx);
+        if (in_tree) {
+          node = &(node->children_[idx]);
+        } else {
+          *res_len = 0;
+          return;
         }
+      } else {
+        break;
+      }
     }
+    if (node) {
+      int id = 0;
+      for (int i = 0; i < node->children_node_len_; i++) {
+        if (node->children_[i].token_id_ != -1) {
+          res[id++] = node->children_[i].token_id_;
+        } else {
+          break;
+        }
+      }
+      *res_len = id;
+    }
+  }
 
-    void destroy_node(TreeNode *node) {
-        if (node) {
-            delete[] node;
-        }
+  void destroy_node(TreeNode *node) {
+    if (node) {
+      delete[] node;
     }
+  }
 
-    void destroy(TreeNode *head) {
-        if (head->children_node_len_ == 0) {
-            // last layer
-            return;
-        }
-        for (int i = 0; i < head->children_node_len_; i++) {
-            destroy(&head->children_[i]);
-        }
-        destroy_node(head->children_);
+  void destroy(TreeNode *head) {
+    if (head->children_node_len_ == 0) {
+      // last layer
+      return;
     }
+    for (int i = 0; i < head->children_node_len_; i++) {
+      destroy(&head->children_[i]);
+    }
+    destroy_node(head->children_);
+  }
 };
 
 class TreeNodeGPU {
-    public:
-    int token_id_;
-    int children_node_len_;
-    int next;
+ public:
+  int token_id_;
+  int children_node_len_;
+  int next;
 
-    __device__ void is_in_tree(TreeNodeGPU *node,
-                               TreeNodeGPU *head,
-                               int token,
-                               int *idx) {
-        for (int i = 0; i < node->children_node_len_; i++) {
-            TreeNodeGPU *children = head + node->next + i;
-            if (children->token_id_ == token) {
-                *idx = node->next + i;
-                return;
-            }
-        }
-        *idx = -1;
+  __device__ void is_in_tree(TreeNodeGPU *node,
+                             TreeNodeGPU *head,
+                             int token,
+                             int *idx) {
+    for (int i = 0; i < node->children_node_len_; i++) {
+      TreeNodeGPU *children = head + node->next + i;
+      if (children->token_id_ == token) {
+        *idx = node->next + i;
+        return;
+      }
     }
+    *idx = -1;
+  }
 
-    __device__ void search(const int *prefix_token,
-                           int token_len,
-                           int *res,
-                           int *res_len) {
-        TreeNodeGPU *node = this;
-        int idx;
-        for (int i = 0; i < token_len; i++) {
-            is_in_tree(node, this, prefix_token[i], &idx);
-            if (idx != -1) {
-                node = this + idx;
-            } else {
-                *res_len = 0;
-                return;
-            }
-        }
-        for (int i = 0; i < node->children_node_len_; i++) {
-            res[i] = (this + node->next + i)->token_id_;
-        }
-        res_len[0] = node->children_node_len_;
+  __device__ void search(const int *prefix_token,
+                         int token_len,
+                         int *res,
+                         int *res_len) {
+    TreeNodeGPU *node = this;
+    int idx;
+    for (int i = 0; i < token_len; i++) {
+      is_in_tree(node, this, prefix_token[i], &idx);
+      if (idx != -1) {
+        node = this + idx;
+      } else {
+        *res_len = 0;
+        return;
+      }
     }
+    for (int i = 0; i < node->children_node_len_; i++) {
+      res[i] = (this + node->next + i)->token_id_;
+    }
+    res_len[0] = node->children_node_len_;
+  }
 
-    template <typename T>
-    __device__ void search(const int *prefix_token, int token_len, T *logits) {
-        TreeNodeGPU *node = this;
-        int idx;
-        for (int i = 0; i < token_len; i++) {
-            is_in_tree(node, this, prefix_token[i], &idx);
-            if (idx != -1) {
-                node = this + idx;
-            } else {
-                return;
-            }
-        }
-        for (int i = 0; i < node->children_node_len_; i++) {
-            // printf("child token: %d\n", (this + node->next + i)->token_id_);
-            logits[(this + node->next + i)->token_id_] = -10000.;
-        }
+  template <typename T>
+  __device__ void search(const int *prefix_token, int token_len, T *logits) {
+    TreeNodeGPU *node = this;
+    int idx;
+    for (int i = 0; i < token_len; i++) {
+      is_in_tree(node, this, prefix_token[i], &idx);
+      if (idx != -1) {
+        node = this + idx;
+      } else {
+        return;
+      }
     }
+    for (int i = 0; i < node->children_node_len_; i++) {
+      // printf("child token: %d\n", (this + node->next + i)->token_id_);
+      logits[(this + node->next + i)->token_id_] = -10000.;
+    }
+  }
 };
 
 template <typename T>
@@ -220,15 +220,15 @@ __global__ void search_on_gpu(TreeNodeGPU *head,
                               T *logits,
                               const int logits_len,
                               const int max_input_len) {
-    int bi = blockIdx.x;
-    int ti = threadIdx.x;
-    if (ti < max_input_len) {
-        int seq_offset = bi * max_input_len;
-        const int *seq_this_thread =
-            input_sequences + seq_offset + max_input_len - ti - 1;
-        T *logits_this_thread = logits + bi * logits_len;
-        head->search(seq_this_thread, ti + 1, logits_this_thread);
-    }
+  int bi = blockIdx.x;
+  int ti = threadIdx.x;
+  if (ti < max_input_len) {
+    int seq_offset = bi * max_input_len;
+    const int *seq_this_thread =
+        input_sequences + seq_offset + max_input_len - ti - 1;
+    T *logits_this_thread = logits + bi * logits_len;
+    head->search(seq_this_thread, ti + 1, logits_this_thread);
+  }
 }
 
 __global__ void search_on_gpu(TreeNodeGPU *head,
@@ -236,19 +236,18 @@ __global__ void search_on_gpu(TreeNodeGPU *head,
                               int *res,
                               int *res_len,
                               const int max_input_len) {
-    int bi = blockIdx.x;
-    int ti = threadIdx.x;
-    if (ti < max_input_len) {
-        int seq_offset = bi * max_input_len;
-        int res_offset = bi * max_input_len * STOP_LIST_BS + ti * STOP_LIST_BS;
-        int res_len_offset = bi * max_input_len + ti;
-        const int *seq_this_thread =
-            input_sequences + seq_offset + max_input_len - ti - 1;
-        int *res_this_thread = res + res_offset;
-        int *res_len_this_thread = res_len + res_len_offset;
-        head->search(
-            seq_this_thread, ti + 1, res_this_thread, res_len_this_thread);
-    }
+  int bi = blockIdx.x;
+  int ti = threadIdx.x;
+  if (ti < max_input_len) {
+    int seq_offset = bi * max_input_len;
+    int res_offset = bi * max_input_len * STOP_LIST_BS + ti * STOP_LIST_BS;
+    int res_len_offset = bi * max_input_len + ti;
+    const int *seq_this_thread =
+        input_sequences + seq_offset + max_input_len - ti - 1;
+    int *res_this_thread = res + res_offset;
+    int *res_len_this_thread = res_len + res_len_offset;
+    head->search(seq_this_thread, ti + 1, res_this_thread, res_len_this_thread);
+  }
 }
 
 template <typename T>
@@ -257,24 +256,24 @@ __global__ void set_value_reverse(const int *res,
                                   T *logits,
                                   const int max_input_len,
                                   const int logits_len) {
-    int bi = blockIdx.x;
-    T *logits_now = logits + bi * logits_len;
-    for (int i = threadIdx.x; i < logits_len; i += blockDim.x) {
-        bool set_flag = true;
-        const int *res_len_now = res_len + bi * max_input_len;
-        for (int j = 0; j < max_input_len; j++) {
-            const int *res_now =
-                res + bi * max_input_len * STOP_LIST_BS + j * STOP_LIST_BS;
-            for (int k = 0; k < res_len_now[j]; k++) {
-                if (i == res_now[k]) {
-                    set_flag = false;
-                    break;
-                }
-            }
-            if (!set_flag) break;
+  int bi = blockIdx.x;
+  T *logits_now = logits + bi * logits_len;
+  for (int i = threadIdx.x; i < logits_len; i += blockDim.x) {
+    bool set_flag = true;
+    const int *res_len_now = res_len + bi * max_input_len;
+    for (int j = 0; j < max_input_len; j++) {
+      const int *res_now =
+          res + bi * max_input_len * STOP_LIST_BS + j * STOP_LIST_BS;
+      for (int k = 0; k < res_len_now[j]; k++) {
+        if (i == res_now[k]) {
+          set_flag = false;
+          break;
         }
-        if (set_flag) logits_now[i] = -10000.;
+      }
+      if (!set_flag) break;
     }
+    if (set_flag) logits_now[i] = -10000.;
+  }
 }
 
 __global__ void setup_gpu_node(TreeNodeGPU *d_head,
@@ -282,30 +281,30 @@ __global__ void setup_gpu_node(TreeNodeGPU *d_head,
                                const int childre_node_num,
                                const int now_id,
                                const int next_id) {
-    d_head[now_id].token_id_ = token_id;
-    d_head[now_id].children_node_len_ = childre_node_num;
-    d_head[now_id].next = next_id;
+  d_head[now_id].token_id_ = token_id;
+  d_head[now_id].children_node_len_ = childre_node_num;
+  d_head[now_id].next = next_id;
 }
 
 __global__ void print_kernel(TreeNodeGPU *d_node, int now_id) {
-    printf("NodeGPU token_id: %d, next_id: %d, children_num: %d\n",
-           d_node[now_id].token_id_,
-           d_node[now_id].next,
-           d_node[now_id].children_node_len_);
+  printf("NodeGPU token_id: %d, next_id: %d, children_num: %d\n",
+         d_node[now_id].token_id_,
+         d_node[now_id].next,
+         d_node[now_id].children_node_len_);
 }
 
 void get_nodes_num(TreeNode *node, int &res) {
-    res += node->children_node_len_;
-    for (int i = 0; i < node->children_node_len_; i++) {
-        get_nodes_num(&(node->children_[i]), res);
-    }
+  res += node->children_node_len_;
+  for (int i = 0; i < node->children_node_len_; i++) {
+    get_nodes_num(&(node->children_[i]), res);
+  }
 }
 
 void setup_tree_cpu(TreeNode *head, const int *stop_list, int stop_list_len) {
-    for (int i = 0; i < STOP_LIST_BS; i++) {
-        int offset = i * stop_list_len;
-        head->insert(stop_list + offset, stop_list_len);
-    }
+  for (int i = 0; i < STOP_LIST_BS; i++) {
+    int offset = i * stop_list_len;
+    head->insert(stop_list + offset, stop_list_len);
+  }
 }
 
 void copy_tree(TreeNode *head,
@@ -314,24 +313,24 @@ void copy_tree(TreeNode *head,
                int now_id,
                int &next_id,
                gpuStream_t stream) {
-    if (!head) {
-        return;
-    }
-    int tmp_next_id = next_id;
-    if (head->children_node_len_ == 0) {
-        tmp_next_id = -1;
-    }
-    setup_gpu_node<<<1, 1, 0, stream>>>(
-        d_head, head->token_id_, head->children_node_len_, now_id, tmp_next_id);
-    GPU(DeviceSynchronize)();
-    depth++;
-    int next_id_this_arrays = next_id;
-    next_id += head->children_node_len_;
-    for (int i = 0; i < head->children_node_len_; i++) {
-        int tmp_now_id = next_id_this_arrays + i;
-        copy_tree(
-            &(head->children_[i]), d_head, depth, tmp_now_id, next_id, stream);
-    }
+  if (!head) {
+    return;
+  }
+  int tmp_next_id = next_id;
+  if (head->children_node_len_ == 0) {
+    tmp_next_id = -1;
+  }
+  setup_gpu_node<<<1, 1, 0, stream>>>(
+      d_head, head->token_id_, head->children_node_len_, now_id, tmp_next_id);
+  GPU(DeviceSynchronize)();
+  depth++;
+  int next_id_this_arrays = next_id;
+  next_id += head->children_node_len_;
+  for (int i = 0; i < head->children_node_len_; i++) {
+    int tmp_now_id = next_id_this_arrays + i;
+    copy_tree(
+        &(head->children_[i]), d_head, depth, tmp_now_id, next_id, stream);
+  }
 }
 
 void search_on_cpu(TreeNode *head,
@@ -340,20 +339,18 @@ void search_on_cpu(TreeNode *head,
                    int *res_len,
                    const int bs,
                    const int max_input_len) {
-    for (int i = 0; i < bs; i++) {
-        for (int j = 0; j < max_input_len; j++) {
-            int seq_offset = i * max_input_len;
-            int res_offset =
-                i * max_input_len * STOP_LIST_BS + j * STOP_LIST_BS;
-            int res_len_offset = i * max_input_len + j;
-            const int *seq_this_time =
-                input_sequences + seq_offset + max_input_len - j - 1;
-            int *res_this_time = res + res_offset;
-            int *res_len_this_time = res_len + res_len_offset;
-            head->search(
-                seq_this_time, j + 1, res_this_time, res_len_this_time);
-        }
+  for (int i = 0; i < bs; i++) {
+    for (int j = 0; j < max_input_len; j++) {
+      int seq_offset = i * max_input_len;
+      int res_offset = i * max_input_len * STOP_LIST_BS + j * STOP_LIST_BS;
+      int res_len_offset = i * max_input_len + j;
+      const int *seq_this_time =
+          input_sequences + seq_offset + max_input_len - j - 1;
+      int *res_this_time = res + res_offset;
+      int *res_len_this_time = res_len + res_len_offset;
+      head->search(seq_this_time, j + 1, res_this_time, res_len_this_time);
     }
+  }
 }
 
 template <paddle::DataType D>
@@ -362,112 +359,111 @@ std::vector<paddle::Tensor> NgramMaskKernel(
     const paddle::Tensor &input_sequences,
     const paddle::Tensor &logits,
     int reverse) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
-    std::vector<int64_t> stop_list_shape = stop_list_tensor.shape();
-    std::vector<int64_t> input_shape = input_sequences.shape();
-    std::vector<int64_t> logits_shape = logits.shape();
-    int logits_len = logits_shape[1];
-    auto logits_out = logits.copy_to(logits.place(), false);
-    auto cu_stream = input_sequences.stream();
-    int bs = input_shape[0];
-    int stop_list_len = stop_list_shape[1];
-    int max_input_len = input_shape[1];
-    static int run_flag = 0;
-    static TreeNode *head = new TreeNode();
-    static TreeNodeGPU *d_head;
-    if (!run_flag) {
-        setup_tree_cpu(head, stop_list_tensor.data<int>(), stop_list_len);
-        int node_num = 0;
-        get_nodes_num(head, node_num);
-        node_num++;
-        printf("node_num: %d\n", node_num);
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+  std::vector<int64_t> stop_list_shape = stop_list_tensor.shape();
+  std::vector<int64_t> input_shape = input_sequences.shape();
+  std::vector<int64_t> logits_shape = logits.shape();
+  int logits_len = logits_shape[1];
+  auto logits_out = logits.copy_to(logits.place(), false);
+  auto cu_stream = input_sequences.stream();
+  int bs = input_shape[0];
+  int stop_list_len = stop_list_shape[1];
+  int max_input_len = input_shape[1];
+  static int run_flag = 0;
+  static TreeNode *head = new TreeNode();
+  static TreeNodeGPU *d_head;
+  if (!run_flag) {
+    setup_tree_cpu(head, stop_list_tensor.data<int>(), stop_list_len);
+    int node_num = 0;
+    get_nodes_num(head, node_num);
+    node_num++;
+    printf("node_num: %d\n", node_num);
 
-        run_flag++;
+    run_flag++;
 
-        GPU(Malloc)(&d_head, node_num * sizeof(TreeNodeGPU));
-        GPU(DeviceSynchronize)();
-        int now_id = 0;
-        int next_id = 1;
-        copy_tree(head, d_head, 0, now_id, next_id, cu_stream);
-        GPU(DeviceSynchronize)();
+    GPU(Malloc)(&d_head, node_num * sizeof(TreeNodeGPU));
+    GPU(DeviceSynchronize)();
+    int now_id = 0;
+    int next_id = 1;
+    copy_tree(head, d_head, 0, now_id, next_id, cu_stream);
+    GPU(DeviceSynchronize)();
 
-        head->destroy(head);
-    }
+    head->destroy(head);
+  }
 
+  int grid_size = bs;
+  int block_size = max_input_len;
+  if (reverse) {
+    auto out_ids = paddle::empty({bs, max_input_len, STOP_LIST_BS},
+                                 paddle::DataType::INT32,
+                                 input_sequences.place());
+    auto out_lens = paddle::empty(
+        {bs, max_input_len}, paddle::DataType::INT32, input_sequences.place());
     int grid_size = bs;
     int block_size = max_input_len;
-    if (reverse) {
-        auto out_ids = paddle::empty({bs, max_input_len, STOP_LIST_BS},
-                                     paddle::DataType::INT32,
-                                     input_sequences.place());
-        auto out_lens = paddle::empty({bs, max_input_len},
-                                      paddle::DataType::INT32,
-                                      input_sequences.place());
-        int grid_size = bs;
-        int block_size = max_input_len;
-        search_on_gpu<<<grid_size, block_size, 0, cu_stream>>>(
-            d_head,
-            input_sequences.data<int>(),
-            out_ids.data<int>(),
-            out_lens.data<int>(),
-            max_input_len);
+    search_on_gpu<<<grid_size, block_size, 0, cu_stream>>>(
+        d_head,
+        input_sequences.data<int>(),
+        out_ids.data<int>(),
+        out_lens.data<int>(),
+        max_input_len);
 
-        set_value_reverse<DataType_><<<grid_size, 256, 0, cu_stream>>>(
-            out_ids.data<int>(),
-            out_lens.data<int>(),
-            reinterpret_cast<DataType_ *>(
-                const_cast<data_t *>(logits_out.data<data_t>())),
-            max_input_len,
-            logits_len);
-    } else {
-        search_on_gpu<DataType_><<<grid_size, block_size, 0, cu_stream>>>(
-            d_head,
-            input_sequences.data<int>(),
-            reinterpret_cast<DataType_ *>(
-                const_cast<data_t *>(logits_out.data<data_t>())),
-            logits_len,
-            max_input_len);
-    }
+    set_value_reverse<DataType_><<<grid_size, 256, 0, cu_stream>>>(
+        out_ids.data<int>(),
+        out_lens.data<int>(),
+        reinterpret_cast<DataType_ *>(
+            const_cast<data_t *>(logits_out.data<data_t>())),
+        max_input_len,
+        logits_len);
+  } else {
+    search_on_gpu<DataType_><<<grid_size, block_size, 0, cu_stream>>>(
+        d_head,
+        input_sequences.data<int>(),
+        reinterpret_cast<DataType_ *>(
+            const_cast<data_t *>(logits_out.data<data_t>())),
+        logits_len,
+        max_input_len);
+  }
 
-    return {logits_out};
+  return {logits_out};
 }
 
 std::vector<paddle::Tensor> NgramMask(const paddle::Tensor &stop_list_tensor,
                                       const paddle::Tensor &input_sequences,
                                       const paddle::Tensor &logits,
                                       int reverse) {
-    switch (logits.type()) {
-        case paddle::DataType::FLOAT16: {
-            return NgramMaskKernel<paddle::DataType::FLOAT16>(
-                stop_list_tensor, input_sequences, logits, reverse);
-        }
-        case paddle::DataType::FLOAT32: {
-            return NgramMaskKernel<paddle::DataType::FLOAT32>(
-                stop_list_tensor, input_sequences, logits, reverse);
-        }
-        default: {
-            PD_THROW(
-                "NOT supported data type. "
-                "Only float16 and float32 are supported. ");
-            break;
-        }
+  switch (logits.type()) {
+    case paddle::DataType::FLOAT16: {
+      return NgramMaskKernel<paddle::DataType::FLOAT16>(
+          stop_list_tensor, input_sequences, logits, reverse);
     }
+    case paddle::DataType::FLOAT32: {
+      return NgramMaskKernel<paddle::DataType::FLOAT32>(
+          stop_list_tensor, input_sequences, logits, reverse);
+    }
+    default: {
+      PD_THROW(
+          "NOT supported data type. "
+          "Only float16 and float32 are supported. ");
+      break;
+    }
+  }
 }
 
 std::vector<std::vector<int64_t>> NgramMaskInferShape(
     const std::vector<int64_t> &stop_list_tensor_shape,
     const std::vector<int64_t> &input_sequences_shape,
     const std::vector<int64_t> &logits_shape) {
-    return {logits_shape};
+  return {logits_shape};
 }
 
 std::vector<paddle::DataType> NgramMaskInferDtype(
     const paddle::DataType &stop_list_tensor_dtype,
     const paddle::DataType &input_sequences_dtype,
     const paddle::DataType &logits_dtype) {
-    return {logits_dtype};
+  return {logits_dtype};
 }
 
 PD_BUILD_STATIC_OP(ngram_mask)
diff --git a/custom_ops/gpu_ops/noaux_tc.cu b/custom_ops/gpu_ops/noaux_tc.cu
index 7b239b8cf4..b9b2d9d32b 100644
--- a/custom_ops/gpu_ops/noaux_tc.cu
+++ b/custom_ops/gpu_ops/noaux_tc.cu
@@ -36,22 +36,24 @@ std::vector<paddle::Tensor> NoauxTc(paddle::Tensor& scores,
   auto place = scores_with_bias.place();
   auto group_scores = paddle::empty({num_tokens, n_group}, input_type, place);
   auto topk_values = paddle::empty({num_tokens, topk}, input_type, place);
-  auto topk_indices = paddle::empty({num_tokens, topk}, paddle::DataType::INT64, place);
+  auto topk_indices =
+      paddle::empty({num_tokens, topk}, paddle::DataType::INT64, place);
   auto stream = scores_with_bias.stream();
 
-  invokeNoAuxTc<float, int64_t>(reinterpret_cast<float*>(scores.data<float>()),
-                       reinterpret_cast<float*>(group_scores.data<float>()),
-                       reinterpret_cast<float*>(topk_values.data<float>()),
-                       reinterpret_cast<int64_t*>(topk_indices.data<int64_t>()),
-                       reinterpret_cast<float*>(scores_with_bias.data<float>()),
-                       num_tokens,
-                       num_experts,
-                       n_group,
-                       topk_group,
-                       topk,
-                       renormalize,
-                       routed_scaling_factor,
-                       stream);
+  invokeNoAuxTc<float, int64_t>(
+      reinterpret_cast<float*>(scores.data<float>()),
+      reinterpret_cast<float*>(group_scores.data<float>()),
+      reinterpret_cast<float*>(topk_values.data<float>()),
+      reinterpret_cast<int64_t*>(topk_indices.data<int64_t>()),
+      reinterpret_cast<float*>(scores_with_bias.data<float>()),
+      num_tokens,
+      num_experts,
+      n_group,
+      topk_group,
+      topk,
+      renormalize,
+      routed_scaling_factor,
+      stream);
 
   return {scores, topk_values, topk_indices};
 }
@@ -64,9 +66,9 @@ std::vector<paddle::DataType> NoauxTcInferDtype(
 
 std::vector<std::vector<int64_t>> NoauxTcInferShape(
     const std::vector<int64_t>& scores_shape,
-    const std::vector<int64_t>& ,
+    const std::vector<int64_t>&,
     const int topk) {
-    auto num_tokens = scores_shape[0];
+  auto num_tokens = scores_shape[0];
   auto topk_values_shape = std::vector<int64_t>{num_tokens, topk};
   auto topk_indices_shape = std::vector<int64_t>{num_tokens, topk};
   return {scores_shape, topk_values_shape, topk_indices_shape};
diff --git a/custom_ops/gpu_ops/open_shm_and_get_meta_signal.cc b/custom_ops/gpu_ops/open_shm_and_get_meta_signal.cc
index 64723eadda..e06868a417 100644
--- a/custom_ops/gpu_ops/open_shm_and_get_meta_signal.cc
+++ b/custom_ops/gpu_ops/open_shm_and_get_meta_signal.cc
@@ -21,22 +21,32 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-using cache_write_complete_signal_type = RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data;
+using cache_write_complete_signal_type =
+    RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data;
 
-paddle::Tensor OpenShmAndGetMetaSignalFunc(const int rank, const int device_id, const bool keep_pd_step_flag) {
-    cache_write_complete_signal_type kv_signal_metadata;
-    const char* fmt_write_cache_completed_signal_str = std::getenv("FLAGS_fmt_write_cache_completed_signal");
-    if (fmt_write_cache_completed_signal_str &&
-        (std::strcmp(fmt_write_cache_completed_signal_str, "true") == 0 ||
-         std::strcmp(fmt_write_cache_completed_signal_str, "1") == 0)) {
-        kv_signal_metadata = RemoteCacheKvIpc::open_shm_and_get_complete_signal_meta_data(rank, device_id, keep_pd_step_flag);
-    }
+paddle::Tensor OpenShmAndGetMetaSignalFunc(const int rank,
+                                           const int device_id,
+                                           const bool keep_pd_step_flag) {
+  cache_write_complete_signal_type kv_signal_metadata;
+  const char *fmt_write_cache_completed_signal_str =
+      std::getenv("FLAGS_fmt_write_cache_completed_signal");
+  if (fmt_write_cache_completed_signal_str &&
+      (std::strcmp(fmt_write_cache_completed_signal_str, "true") == 0 ||
+       std::strcmp(fmt_write_cache_completed_signal_str, "1") == 0)) {
+    kv_signal_metadata =
+        RemoteCacheKvIpc::open_shm_and_get_complete_signal_meta_data(
+            rank, device_id, keep_pd_step_flag);
+  }
 
-    auto kv_signal_metadata_out = paddle::full({3}, -1, paddle::DataType::INT64, paddle::CPUPlace());
-    kv_signal_metadata_out.data<int64_t>()[0] = static_cast<int64_t>(kv_signal_metadata.layer_id);
-    kv_signal_metadata_out.data<int64_t>()[1] = reinterpret_cast<int64_t>(kv_signal_metadata.shm_ptr);
-    kv_signal_metadata_out.data<int64_t>()[2] = static_cast<int64_t>(kv_signal_metadata.shm_fd);
-    return kv_signal_metadata_out;
+  auto kv_signal_metadata_out =
+      paddle::full({3}, -1, paddle::DataType::INT64, paddle::CPUPlace());
+  kv_signal_metadata_out.data<int64_t>()[0] =
+      static_cast<int64_t>(kv_signal_metadata.layer_id);
+  kv_signal_metadata_out.data<int64_t>()[1] =
+      reinterpret_cast<int64_t>(kv_signal_metadata.shm_ptr);
+  kv_signal_metadata_out.data<int64_t>()[2] =
+      static_cast<int64_t>(kv_signal_metadata.shm_fd);
+  return kv_signal_metadata_out;
 }
 
 void InitKVSignalPerQuery(const paddle::Tensor &seq_lens_encoder_tensor,
@@ -44,32 +54,39 @@ void InitKVSignalPerQuery(const paddle::Tensor &seq_lens_encoder_tensor,
                           const paddle::Tensor &seq_lens_decoder_tensor,
                           const int rank,
                           const int num_layers) {
-    const char* fmt_write_cache_completed_signal_str = std::getenv("FLAGS_fmt_write_cache_completed_signal");
-    if (fmt_write_cache_completed_signal_str &&
-        (std::strcmp(fmt_write_cache_completed_signal_str, "true") == 0 ||
-         std::strcmp(fmt_write_cache_completed_signal_str, "1") == 0)) {
-        int real_bsz = seq_lens_this_time_tensor.dims()[0];
-        // GPU init, cp to cpu?
-        auto seq_lens_encoder_cpu = seq_lens_encoder_tensor.copy_to(paddle::CPUPlace(), false);
-        auto seq_lens_decoder_cpu = seq_lens_decoder_tensor.copy_to(paddle::CPUPlace(), false);
-        RemoteCacheKvIpc::kv_complete_signal_meta_data_per_query.init(
-            seq_lens_encoder_cpu.data<int>(),
-            seq_lens_decoder_cpu.data<int>(),
-            rank, num_layers, real_bsz);
-    }
+  const char *fmt_write_cache_completed_signal_str =
+      std::getenv("FLAGS_fmt_write_cache_completed_signal");
+  if (fmt_write_cache_completed_signal_str &&
+      (std::strcmp(fmt_write_cache_completed_signal_str, "true") == 0 ||
+       std::strcmp(fmt_write_cache_completed_signal_str, "1") == 0)) {
+    int real_bsz = seq_lens_this_time_tensor.dims()[0];
+    // GPU init, cp to cpu?
+    auto seq_lens_encoder_cpu =
+        seq_lens_encoder_tensor.copy_to(paddle::CPUPlace(), false);
+    auto seq_lens_decoder_cpu =
+        seq_lens_decoder_tensor.copy_to(paddle::CPUPlace(), false);
+    RemoteCacheKvIpc::kv_complete_signal_meta_data_per_query.init(
+        seq_lens_encoder_cpu.data<int>(),
+        seq_lens_decoder_cpu.data<int>(),
+        rank,
+        num_layers,
+        real_bsz);
+  }
 }
 
-std::vector<paddle::Tensor> OpenShmAndGetMetaSignal(const int rank, const int device_id, const bool keep_pd_step_flag) {
-    return {OpenShmAndGetMetaSignalFunc(rank, device_id, keep_pd_step_flag)};
+std::vector<paddle::Tensor> OpenShmAndGetMetaSignal(
+    const int rank, const int device_id, const bool keep_pd_step_flag) {
+  return {OpenShmAndGetMetaSignalFunc(rank, device_id, keep_pd_step_flag)};
 }
 
-
-std::vector<std::vector<int64_t>> OpenShmAndGetMetaSignalShape(const int rank, const int device_id, const bool keep_pd_step_flag) {
-    return {{3}};
+std::vector<std::vector<int64_t>> OpenShmAndGetMetaSignalShape(
+    const int rank, const int device_id, const bool keep_pd_step_flag) {
+  return {{3}};
 }
 
-std::vector<paddle::DataType> OpenShmAndGetMetaSignalDtype(const int rank, const int device_id, const bool keep_pd_step_flag) {
-    return {paddle::DataType::INT64};
+std::vector<paddle::DataType> OpenShmAndGetMetaSignalDtype(
+    const int rank, const int device_id, const bool keep_pd_step_flag) {
+  return {paddle::DataType::INT64};
 }
 
 PD_BUILD_STATIC_OP(open_shm_and_get_meta_signal)
diff --git a/custom_ops/gpu_ops/quantization/common.cu b/custom_ops/gpu_ops/quantization/common.cu
index c0e8f48ee2..461ca53167 100644
--- a/custom_ops/gpu_ops/quantization/common.cu
+++ b/custom_ops/gpu_ops/quantization/common.cu
@@ -1,35 +1,39 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
-
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
 
 #include "quantization/common.cuh"
 
-// adapted from: https://github.com/sgl-project/sglang/blob/v0.5.2rc2/sgl-kernel/csrc/gemm/per_token_quant_fp8.cu
+// adapted from:
+// https://github.com/sgl-project/sglang/blob/v0.5.2rc2/sgl-kernel/csrc/gemm/per_token_quant_fp8.cu
 
 // ---------------------------------------------------------------------------
 // 1. Warp‑local, no shared memory
 //    • One warp handles one token.
 //    • Eight tokens per 256‑thread CTA.
 // ---------------------------------------------------------------------------
-template <typename T, typename DST_DTYPE, int kTokensPerCTA = 8, int kVecSize = 16>
-__global__ void per_token_quant_fp8_kernel(
-    const T* __restrict__ input,
-    DST_DTYPE* __restrict__ output_q,
-    float* __restrict__ output_s,
-    const float scale_ub,
-    const int64_t hidden_size,
-    const int64_t num_tokens) {
+template <typename T,
+          typename DST_DTYPE,
+          int kTokensPerCTA = 8,
+          int kVecSize = 16>
+__global__ void per_token_quant_fp8_kernel(const T *__restrict__ input,
+                                           DST_DTYPE *__restrict__ output_q,
+                                           float *__restrict__ output_s,
+                                           const float scale_ub,
+                                           const int64_t hidden_size,
+                                           const int64_t num_tokens) {
   const int warp_id = threadIdx.x / WARP_SIZE;        // 0‑7  (8 warps)
   const int lane_id = threadIdx.x & (WARP_SIZE - 1);  // 0‑31
   const int token_id = blockIdx.x * kTokensPerCTA + warp_id;
   if (token_id >= num_tokens) return;
 
   // Global tensors for this token
-  const T* token_input = input + token_id * hidden_size;
-  DST_DTYPE* token_output = output_q + token_id * hidden_size;
-  float* token_scale = output_s + token_id;
+  const T *token_input = input + token_id * hidden_size;
+  DST_DTYPE *token_output = output_q + token_id * hidden_size;
+  float *token_scale = output_s + token_id;
 
   //
-  // Pass-1: Perform a warp reduce to find the max_value of a token's hidden_size
+  // Pass-1: Perform a warp reduce to find the max_value of a token's
+  // hidden_size
   //
   float max_value = 0.f;
   using vec_t = AlignedVector<T, kVecSize>;
@@ -46,7 +50,7 @@ __global__ void per_token_quant_fp8_kernel(
   }
 
   float warp_max = warpReduceMax(max_value);
-  if (scale_ub > 0){
+  if (scale_ub > 0) {
     warp_max = fminf(warp_max, scale_ub);
   }
   float scale;
@@ -71,7 +75,7 @@ __global__ void per_token_quant_fp8_kernel(
       output_arr[j] = static_cast<DST_DTYPE>(val);
     }
     if constexpr (kVecSize == 16) {
-      *(uint4*)(token_output + i * kVecSize) = *(uint4*)output_arr;
+      *(uint4 *)(token_output + i * kVecSize) = *(uint4 *)output_arr;
     } else {
       // Use element-wise copy for vector size 8 to ensure correctness
       for (int k = 0; k < kVecSize; ++k) {
@@ -86,9 +90,9 @@ __global__ void per_token_quant_fp8_kernel(
 // ---------------------------------------------------------------------------
 template <typename T, typename DST_DTYPE, int kVecSize = 16>
 __global__ void per_token_quant_fp8_small_batch_kernel(
-    const T* __restrict__ input,
-    DST_DTYPE* __restrict__ output_q,
-    float* __restrict__ output_s,
+    const T *__restrict__ input,
+    DST_DTYPE *__restrict__ output_q,
+    float *__restrict__ output_s,
     const float scale_ub,
     const int64_t hidden_size,
     const int64_t num_tokens) {
@@ -98,8 +102,8 @@ __global__ void per_token_quant_fp8_small_batch_kernel(
   const int tid = threadIdx.x;
   const int block_dim = blockDim.x;
 
-  const T* token_input = input + token_idx * hidden_size;
-  DST_DTYPE* token_output = output_q + token_idx * hidden_size;
+  const T *token_input = input + token_idx * hidden_size;
+  DST_DTYPE *token_output = output_q + token_idx * hidden_size;
 
   float max_value = 0.0f;
 
@@ -120,7 +124,7 @@ __global__ void per_token_quant_fp8_small_batch_kernel(
   }
 
   max_value = blockReduceMax(max_value);
-  if (scale_ub > 0){
+  if (scale_ub > 0) {
     max_value = fminf(max_value, scale_ub);
   }
   __shared__ float scale;
@@ -140,12 +144,14 @@ __global__ void per_token_quant_fp8_small_batch_kernel(
     DST_DTYPE output_arr[kVecSize];
 #pragma unroll
     for (uint32_t j = 0; j < kVecSize; ++j) {
-      float val = fmaxf(fminf(static_cast<float>(input_vec[j]) * scale_inv, FP8_E4M3_MAX), -FP8_E4M3_MAX);
+      float val = fmaxf(
+          fminf(static_cast<float>(input_vec[j]) * scale_inv, FP8_E4M3_MAX),
+          -FP8_E4M3_MAX);
       output_arr[j] = static_cast<DST_DTYPE>(val);
     }
 
     if constexpr (kVecSize == 16) {
-      *(uint4*)(token_output + i * kVecSize) = *(uint4*)output_arr;
+      *(uint4 *)(token_output + i * kVecSize) = *(uint4 *)output_arr;
     } else {
       // Use element-wise copy for vector size 8 to ensure correctness
       for (int k = 0; k < kVecSize; ++k) {
@@ -173,8 +179,11 @@ __global__ void scaled_fp8_quant_kernel(fp8_type *__restrict__ out,
 
 template <typename scalar_t, typename fp8_type>
 __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
-    fp8_type *__restrict__ out, float *__restrict__ scale,
-    scalar_t const *__restrict__ input, float scale_ub, const int hidden_size) {
+    fp8_type *__restrict__ out,
+    float *__restrict__ scale,
+    scalar_t const *__restrict__ input,
+    float scale_ub,
+    const int hidden_size) {
   int const tid = threadIdx.x;
   int const token_idx = blockIdx.x;
 
@@ -228,11 +237,11 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy
 
-void StaticScaledFp8Quant(paddle::Tensor &out,         // [..., d]
-                          paddle::Tensor const &input, // [..., d]
-                          paddle::Tensor const &scale) // [1]
+void StaticScaledFp8Quant(paddle::Tensor &out,          // [..., d]
+                          paddle::Tensor const &input,  // [..., d]
+                          paddle::Tensor const &scale)  // [1]
 {
   PD_CHECK(out.dtype() == paddle::DataType::FLOAT8_E4M3FN);
   using fp8_t = phi::dtype::float8_e4m3fn;
@@ -245,35 +254,41 @@ void StaticScaledFp8Quant(paddle::Tensor &out,         // [..., d]
   cudaStream_t stream = input.stream();
 
   switch (input.dtype()) {
-  case paddle::DataType::FLOAT32: {
-    using scalar_t = float;
-    fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(out.data<fp8_t>(), input.data<scalar_t>(),
-                                     scale.data<float>(), num_elems);
-    break;
-  }
-  case paddle::DataType::FLOAT16: {
-    using scalar_t = phi::dtype::float16;
-    fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(out.data<fp8_t>(), input.data<scalar_t>(),
-                                     scale.data<float>(), num_elems);
-    break;
-  }
-  case paddle::DataType::BFLOAT16: {
-    using scalar_t = phi::dtype::bfloat16;
-    fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(out.data<fp8_t>(), input.data<scalar_t>(),
-                                     scale.data<float>(), num_elems);
-    break;
-  }
-  default:
-    PD_THROW("Only supported attr of input type in [fp32, fp16, bf16].");
+    case paddle::DataType::FLOAT32: {
+      using scalar_t = float;
+      fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(out.data<fp8_t>(),
+                                       input.data<scalar_t>(),
+                                       scale.data<float>(),
+                                       num_elems);
+      break;
+    }
+    case paddle::DataType::FLOAT16: {
+      using scalar_t = phi::dtype::float16;
+      fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(out.data<fp8_t>(),
+                                       input.data<scalar_t>(),
+                                       scale.data<float>(),
+                                       num_elems);
+      break;
+    }
+    case paddle::DataType::BFLOAT16: {
+      using scalar_t = phi::dtype::bfloat16;
+      fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(out.data<fp8_t>(),
+                                       input.data<scalar_t>(),
+                                       scale.data<float>(),
+                                       num_elems);
+      break;
+    }
+    default:
+      PD_THROW("Only supported attr of input type in [fp32, fp16, bf16].");
   }
 }
 
-void DynamicScaledFp8Quant(paddle::Tensor &out,         // [..., d]
-                           paddle::Tensor const &input, // [..., d]
-                           paddle::Tensor &scale)       // [1]
+void DynamicScaledFp8Quant(paddle::Tensor &out,          // [..., d]
+                           paddle::Tensor const &input,  // [..., d]
+                           paddle::Tensor &scale)        // [1]
 {
   PD_CHECK(out.dtype() == paddle::DataType::FLOAT8_E4M3FN);
   using fp8_t = phi::dtype::float8_e4m3fn;
@@ -286,44 +301,51 @@ void DynamicScaledFp8Quant(paddle::Tensor &out,         // [..., d]
   cudaStream_t stream = input.stream();
 
   switch (input.dtype()) {
-  case paddle::DataType::FLOAT32: {
-    using scalar_t = float;
-    fastdeploy::segmented_max_reduction<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(scale.data<float>(),
-                                     input.data<scalar_t>(), num_elems);
-    fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(out.data<fp8_t>(), input.data<scalar_t>(),
-                                     scale.data<float>(), num_elems);
-    break;
-  }
-  case paddle::DataType::FLOAT16: {
-    using scalar_t = phi::dtype::float16;
-    fastdeploy::segmented_max_reduction<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(scale.data<float>(),
-                                     input.data<scalar_t>(), num_elems);
-    fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(out.data<fp8_t>(), input.data<scalar_t>(),
-                                     scale.data<float>(), num_elems);
-    break;
-  }
-  case paddle::DataType::BFLOAT16: {
-    using scalar_t = phi::dtype::bfloat16;
-    fastdeploy::segmented_max_reduction<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(scale.data<float>(),
-                                     input.data<scalar_t>(), num_elems);
-    fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(out.data<fp8_t>(), input.data<scalar_t>(),
-                                     scale.data<float>(), num_elems);
-    break;
-  }
-  default:
-    PD_THROW("Only supported attr of input type in [fp32, fp16, bf16].");
+    case paddle::DataType::FLOAT32: {
+      using scalar_t = float;
+      fastdeploy::segmented_max_reduction<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(
+              scale.data<float>(), input.data<scalar_t>(), num_elems);
+      fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(out.data<fp8_t>(),
+                                       input.data<scalar_t>(),
+                                       scale.data<float>(),
+                                       num_elems);
+      break;
+    }
+    case paddle::DataType::FLOAT16: {
+      using scalar_t = phi::dtype::float16;
+      fastdeploy::segmented_max_reduction<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(
+              scale.data<float>(), input.data<scalar_t>(), num_elems);
+      fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(out.data<fp8_t>(),
+                                       input.data<scalar_t>(),
+                                       scale.data<float>(),
+                                       num_elems);
+      break;
+    }
+    case paddle::DataType::BFLOAT16: {
+      using scalar_t = phi::dtype::bfloat16;
+      fastdeploy::segmented_max_reduction<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(
+              scale.data<float>(), input.data<scalar_t>(), num_elems);
+      fastdeploy::scaled_fp8_quant_kernel<scalar_t, fp8_t>
+          <<<grid, block, 0, stream>>>(out.data<fp8_t>(),
+                                       input.data<scalar_t>(),
+                                       scale.data<float>(),
+                                       num_elems);
+      break;
+    }
+    default:
+      PD_THROW("Only supported attr of input type in [fp32, fp16, bf16].");
   }
 }
 
-void DynamicPerTokenScaledFp8Quant(paddle::Tensor &out,         // [..., d]
-                                   paddle::Tensor const &input, // [..., d]
-                                   paddle::Tensor &scales, float scale_ub) {
+void DynamicPerTokenScaledFp8Quant(paddle::Tensor &out,          // [..., d]
+                                   paddle::Tensor const &input,  // [..., d]
+                                   paddle::Tensor &scales,
+                                   float scale_ub) {
   PD_CHECK(input.is_contiguous());
   PD_CHECK(out.is_contiguous());
   PD_CHECK(out.dtype() == paddle::DataType::FLOAT8_E4M3FN);
@@ -333,7 +355,7 @@ void DynamicPerTokenScaledFp8Quant(paddle::Tensor &out,         // [..., d]
   int const num_tokens = input.numel() / hidden_size;
   cudaStream_t stream = input.stream();
 
-  if (hidden_size % 8 == 0){
+  if (hidden_size % 8 == 0) {
     int device = 0;
     cudaGetDevice(&device);
     int sm_count = 0;
@@ -343,50 +365,58 @@ void DynamicPerTokenScaledFp8Quant(paddle::Tensor &out,         // [..., d]
     const bool use_vec16 = (hidden_size % 16 == 0);
     DISPATCH_FLOAT_FP6_DTYPE(input.dtype(), scalar_t, {
       if (use_warp_kernel) {
-        // -------- warp‑local ---------------------------------------------------
+        // -------- warp‑local
+        // ---------------------------------------------------
         constexpr int THREADS = TOKENS_PER_CTA * WARP_SIZE;  // 256
         dim3 grid((num_tokens + TOKENS_PER_CTA - 1) / TOKENS_PER_CTA);
         dim3 block(THREADS);
 
         if (use_vec16) {
-          per_token_quant_fp8_kernel<scalar_t, __nv_fp8_e4m3, TOKENS_PER_CTA, 16><<<grid, block, 0, stream>>>(
-              reinterpret_cast<const scalar_t*>(input.data<scalar_t>()),
-              reinterpret_cast<__nv_fp8_e4m3*>(out.data<fp8_t>()),
-              reinterpret_cast<float*>(scales.data<float>()),
+          per_token_quant_fp8_kernel<scalar_t,
+                                     __nv_fp8_e4m3,
+                                     TOKENS_PER_CTA,
+                                     16><<<grid, block, 0, stream>>>(
+              reinterpret_cast<const scalar_t *>(input.data<scalar_t>()),
+              reinterpret_cast<__nv_fp8_e4m3 *>(out.data<fp8_t>()),
+              reinterpret_cast<float *>(scales.data<float>()),
               scale_ub,
               hidden_size,
               num_tokens);
         } else {
-          per_token_quant_fp8_kernel<scalar_t, __nv_fp8_e4m3, TOKENS_PER_CTA, 8><<<grid, block, 0, stream>>>(
-              reinterpret_cast<const scalar_t*>(input.data<scalar_t>()),
-              reinterpret_cast<__nv_fp8_e4m3*>(out.data<fp8_t>()),
-              reinterpret_cast<float*>(scales.data<float>()),
-              scale_ub,
-              hidden_size,
-              num_tokens);
+          per_token_quant_fp8_kernel<scalar_t, __nv_fp8_e4m3, TOKENS_PER_CTA, 8>
+              <<<grid, block, 0, stream>>>(
+                  reinterpret_cast<const scalar_t *>(input.data<scalar_t>()),
+                  reinterpret_cast<__nv_fp8_e4m3 *>(out.data<fp8_t>()),
+                  reinterpret_cast<float *>(scales.data<float>()),
+                  scale_ub,
+                  hidden_size,
+                  num_tokens);
         }
       } else {
-        // -------- baseline -----------------------------------------------------
+        // -------- baseline
+        // -----------------------------------------------------
         constexpr int THREADS = 256;
         dim3 grid(num_tokens);
         dim3 block(THREADS);
 
         if (use_vec16) {
-          per_token_quant_fp8_small_batch_kernel<scalar_t, __nv_fp8_e4m3, 16><<<grid, block, 0, stream>>>(
-              reinterpret_cast<const scalar_t*>(input.data<scalar_t>()),
-              reinterpret_cast<__nv_fp8_e4m3*>(out.data<fp8_t>()),
-              reinterpret_cast<float*>(scales.data<float>()),
-              scale_ub,
-              hidden_size,
-              num_tokens);
+          per_token_quant_fp8_small_batch_kernel<scalar_t, __nv_fp8_e4m3, 16>
+              <<<grid, block, 0, stream>>>(
+                  reinterpret_cast<const scalar_t *>(input.data<scalar_t>()),
+                  reinterpret_cast<__nv_fp8_e4m3 *>(out.data<fp8_t>()),
+                  reinterpret_cast<float *>(scales.data<float>()),
+                  scale_ub,
+                  hidden_size,
+                  num_tokens);
         } else {
-          per_token_quant_fp8_small_batch_kernel<scalar_t, __nv_fp8_e4m3, 8><<<grid, block, 0, stream>>>(
-              reinterpret_cast<const scalar_t*>(input.data<scalar_t>()),
-              reinterpret_cast<__nv_fp8_e4m3*>(out.data<fp8_t>()),
-              reinterpret_cast<float*>(scales.data<float>()),
-              scale_ub,
-              hidden_size,
-              num_tokens);
+          per_token_quant_fp8_small_batch_kernel<scalar_t, __nv_fp8_e4m3, 8>
+              <<<grid, block, 0, stream>>>(
+                  reinterpret_cast<const scalar_t *>(input.data<scalar_t>()),
+                  reinterpret_cast<__nv_fp8_e4m3 *>(out.data<fp8_t>()),
+                  reinterpret_cast<float *>(scales.data<float>()),
+                  scale_ub,
+                  hidden_size,
+                  num_tokens);
         }
       }
     });
@@ -398,11 +428,12 @@ void DynamicPerTokenScaledFp8Quant(paddle::Tensor &out,         // [..., d]
 
   DISPATCH_FLOAT_FP6_DTYPE(input.dtype(), scalar_t, {
     fastdeploy::dynamic_per_token_scaled_fp8_quant_kernel<scalar_t, fp8_t>
-        <<<grid, block, 0, stream>>>(out.data<fp8_t>(), scales.data<float>(),
-                                     input.data<scalar_t>(), scale_ub,
+        <<<grid, block, 0, stream>>>(out.data<fp8_t>(),
+                                     scales.data<float>(),
+                                     input.data<scalar_t>(),
+                                     scale_ub,
                                      hidden_size);
   });
-
 }
 
 PD_BUILD_STATIC_OP(static_scaled_fp8_quant)
@@ -414,8 +445,7 @@ PD_BUILD_STATIC_OP(static_scaled_fp8_quant)
 PD_BUILD_STATIC_OP(dynamic_scaled_fp8_quant)
     .Inputs({"out", "input", "scale"})
     .Outputs({"out_q", "out_scale"})
-    .SetInplaceMap({{"out", "out_q"},
-                    {"scale", "out_scale"}})
+    .SetInplaceMap({{"out", "out_q"}, {"scale", "out_scale"}})
     .SetKernelFn(PD_KERNEL(DynamicScaledFp8Quant));
 
 PD_BUILD_STATIC_OP(dynamic_per_token_scaled_fp8_quant)
diff --git a/custom_ops/gpu_ops/quantization/common.cuh b/custom_ops/gpu_ops/quantization/common.cuh
index 6581e67a7c..25dd5ee8fe 100644
--- a/custom_ops/gpu_ops/quantization/common.cuh
+++ b/custom_ops/gpu_ops/quantization/common.cuh
@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
 
 #pragma once
 
@@ -11,14 +12,16 @@
 namespace fastdeploy {
 
 // Vectorization containers
-template <typename scalar_t> struct __align__(8) vec4_t {
+template <typename scalar_t>
+struct __align__(8) vec4_t {
   scalar_t x;
   scalar_t y;
   scalar_t z;
   scalar_t w;
 };
 
-template <typename quant_type_t> struct __align__(4) q8x4_t {
+template <typename quant_type_t>
+struct __align__(4) q8x4_t {
   static_assert(std::is_same_v<quant_type_t, int8_t> ||
                 std::is_same_v<quant_type_t, phi::dtype::float8_e4m3fn>);
   quant_type_t x;
@@ -94,7 +97,8 @@ __global__ void segmented_max_reduction(float *__restrict__ scale,
 
 template <typename scalar_t>
 __device__ float thread_max_vec(scalar_t const *__restrict__ input,
-                                int64_t const num_elems, int const tid,
+                                int64_t const num_elems,
+                                int const tid,
                                 int const step) {
   // Vectorized input/output to better utilize memory bandwidth.
   vec4_t<scalar_t> const *vectorized_in =
@@ -125,7 +129,8 @@ __device__ void scaled_fp8_conversion_vec(fp8_type *__restrict__ out,
                                           scalar_t const *__restrict__ input,
                                           float const scale,
                                           int64_t const num_elems,
-                                          int const tid, int const step) {
+                                          int const tid,
+                                          int const step) {
   using float8x4_t = q8x4_t<fp8_type>;
   // Vectorized input/output to better utilize memory bandwidth.
   auto const *vectorized_in = reinterpret_cast<vec4_t<scalar_t> const *>(input);
@@ -156,4 +161,4 @@ __device__ void scaled_fp8_conversion_vec(fp8_type *__restrict__ out,
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy
diff --git a/custom_ops/gpu_ops/read_data_ipc.cu b/custom_ops/gpu_ops/read_data_ipc.cu
index 4de0d4a5c2..fffe29c5d1 100644
--- a/custom_ops/gpu_ops/read_data_ipc.cu
+++ b/custom_ops/gpu_ops/read_data_ipc.cu
@@ -30,44 +30,43 @@ namespace cub = hipcub;
 
 template <typename T>
 __global__ void set_data(T *input, int n) {
-    if (threadIdx.x == 0) {
-        for (int i = 0; i < n; ++i) {
-            *(input + i) = static_cast<T>(i);
-            printf("set[%d]: %f\n", i, *(input + i));
-        }
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < n; ++i) {
+      *(input + i) = static_cast<T>(i);
+      printf("set[%d]: %f\n", i, *(input + i));
     }
+  }
 }
 
 template <typename T>
 __global__ void print_data(T *input, int n) {
-    if (threadIdx.x == 0) {
-        for (int i = 0; i < n; ++i) {
-            printf("input[%d]: %f\n", i, input[i]);
-        }
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < n; ++i) {
+      printf("input[%d]: %f\n", i, input[i]);
     }
+  }
 }
 
 void ReadDataIpc(const paddle::Tensor &tmp_input,
                  int64_t data_ptr,
                  const std::string &shm_name) {
-    volatile shmStruct *shm = NULL;
-    sharedMemoryInfo info;
-    if (sharedMemoryOpen(shm_name.c_str(), sizeof(shmStruct), &info) != 0) {
-        printf("Failed to create shared memory slab\n");
-        printf("Func ReadDataIpc. Shm_name: %s\n", shm_name.c_str());
-        exit(EXIT_FAILURE);
-    }
-    shm = (volatile shmStruct *)info.addr;
-    void *ptr = nullptr;
-    checkCudaErrors(
-        GPU(IpcOpenMemHandle)(&ptr,
-                              *(GPU(IpcMemHandle_t) *)&shm->memHandle,
-                              GPU(IpcMemLazyEnablePeerAccess)));
-    printf("ptr: %p\n", ptr);
-    print_data<float><<<1, 1>>>(reinterpret_cast<float *>(ptr), 10);
-    GPU(DeviceSynchronize)();
-    checkCudaErrors(GPU(IpcCloseMemHandle)(ptr));
-    sharedMemoryClose(&info);
+  volatile shmStruct *shm = NULL;
+  sharedMemoryInfo info;
+  if (sharedMemoryOpen(shm_name.c_str(), sizeof(shmStruct), &info) != 0) {
+    printf("Failed to create shared memory slab\n");
+    printf("Func ReadDataIpc. Shm_name: %s\n", shm_name.c_str());
+    exit(EXIT_FAILURE);
+  }
+  shm = (volatile shmStruct *)info.addr;
+  void *ptr = nullptr;
+  checkCudaErrors(GPU(IpcOpenMemHandle)(&ptr,
+                                        *(GPU(IpcMemHandle_t) *)&shm->memHandle,
+                                        GPU(IpcMemLazyEnablePeerAccess)));
+  printf("ptr: %p\n", ptr);
+  print_data<float><<<1, 1>>>(reinterpret_cast<float *>(ptr), 10);
+  GPU(DeviceSynchronize)();
+  checkCudaErrors(GPU(IpcCloseMemHandle)(ptr));
+  sharedMemoryClose(&info);
 }
 
 PD_BUILD_STATIC_OP(read_data_ipc)
diff --git a/custom_ops/gpu_ops/recover_decode_task.cu b/custom_ops/gpu_ops/recover_decode_task.cu
index ae4e77ad6d..edb1164c56 100644
--- a/custom_ops/gpu_ops/recover_decode_task.cu
+++ b/custom_ops/gpu_ops/recover_decode_task.cu
@@ -24,115 +24,119 @@ __global__ void recover_decode_task(bool *stop_flags,
                                     const int bsz,
                                     const int block_num_per_seq,
                                     const int block_size) {
-    int thread_idx = threadIdx.x;
-    if (thread_idx < bsz) {
-        if(is_block_step[thread_idx] == true) {
-            int *block_table_now = block_tables + thread_idx * block_num_per_seq;
-            if (block_table_now[step_seq_lens_decoder[thread_idx] / block_size] != -1) {
-                // can be recovered for decoding
-                is_block_step[thread_idx] = false;
-                seq_lens_this_time[thread_idx]= 1;
-                stop_flags[thread_idx] = false;
-                seq_lens_encoder[thread_idx] = 0;
-                seq_lens_decoder[thread_idx] = step_seq_lens_decoder[thread_idx];
-
-            }
-        }
+  int thread_idx = threadIdx.x;
+  if (thread_idx < bsz) {
+    if (is_block_step[thread_idx] == true) {
+      int *block_table_now = block_tables + thread_idx * block_num_per_seq;
+      if (block_table_now[step_seq_lens_decoder[thread_idx] / block_size] !=
+          -1) {
+        // can be recovered for decoding
+        is_block_step[thread_idx] = false;
+        seq_lens_this_time[thread_idx] = 1;
+        stop_flags[thread_idx] = false;
+        seq_lens_encoder[thread_idx] = 0;
+        seq_lens_decoder[thread_idx] = step_seq_lens_decoder[thread_idx];
+      }
     }
+  }
 }
 
 __global__ void recover_spec_decode_task(bool *stop_flags,
-                                        int *seq_lens_this_time,
-                                        int *seq_lens_encoder,
-                                        int *seq_lens_decoder,
-                                        int *step_seq_lens_decoder,
-                                        int *block_tables,
-                                        bool *is_block_step,
-                                        int64_t *draft_tokens,
-                                        const int64_t *step_draft_tokens,
-                                        const int *step_seq_lens_this_time,
-                                        const int bsz,
-                                        const int block_num_per_seq,
-                                        const int block_size,
-                                        const int draft_tokens_len,
-                                        const int num_extra_tokens) {
-    int thread_idx = threadIdx.x;
-    if (thread_idx < bsz) {
-        if(is_block_step[thread_idx] == true) {
-            int *block_table_now = block_tables + thread_idx * block_num_per_seq;
-            int max_possible_block_idx = (step_seq_lens_decoder[thread_idx] + num_extra_tokens) / block_size;
-            max_possible_block_idx = min(max_possible_block_idx, block_num_per_seq);
-            if (block_table_now[max_possible_block_idx] != -1) {
-                // can be recovered for decoding
-                int64_t *draft_tokens_now = draft_tokens + thread_idx * draft_tokens_len;
-                const int64_t *step_draft_tokens_now = step_draft_tokens + thread_idx * draft_tokens_len;
-                is_block_step[thread_idx] = false;
-                seq_lens_this_time[thread_idx] = step_seq_lens_this_time[thread_idx];
-                stop_flags[thread_idx] = false;
-                seq_lens_encoder[thread_idx] = 0;
-                seq_lens_decoder[thread_idx] = step_seq_lens_decoder[thread_idx];
-                for (int i = 0; i < seq_lens_this_time[thread_idx]; i++) {
-                    draft_tokens_now[i] = step_draft_tokens_now[i];
-                }
-
-            }
+                                         int *seq_lens_this_time,
+                                         int *seq_lens_encoder,
+                                         int *seq_lens_decoder,
+                                         int *step_seq_lens_decoder,
+                                         int *block_tables,
+                                         bool *is_block_step,
+                                         int64_t *draft_tokens,
+                                         const int64_t *step_draft_tokens,
+                                         const int *step_seq_lens_this_time,
+                                         const int bsz,
+                                         const int block_num_per_seq,
+                                         const int block_size,
+                                         const int draft_tokens_len,
+                                         const int num_extra_tokens) {
+  int thread_idx = threadIdx.x;
+  if (thread_idx < bsz) {
+    if (is_block_step[thread_idx] == true) {
+      int *block_table_now = block_tables + thread_idx * block_num_per_seq;
+      int max_possible_block_idx =
+          (step_seq_lens_decoder[thread_idx] + num_extra_tokens) / block_size;
+      max_possible_block_idx = min(max_possible_block_idx, block_num_per_seq);
+      if (block_table_now[max_possible_block_idx] != -1) {
+        // can be recovered for decoding
+        int64_t *draft_tokens_now =
+            draft_tokens + thread_idx * draft_tokens_len;
+        const int64_t *step_draft_tokens_now =
+            step_draft_tokens + thread_idx * draft_tokens_len;
+        is_block_step[thread_idx] = false;
+        seq_lens_this_time[thread_idx] = step_seq_lens_this_time[thread_idx];
+        stop_flags[thread_idx] = false;
+        seq_lens_encoder[thread_idx] = 0;
+        seq_lens_decoder[thread_idx] = step_seq_lens_decoder[thread_idx];
+        for (int i = 0; i < seq_lens_this_time[thread_idx]; i++) {
+          draft_tokens_now[i] = step_draft_tokens_now[i];
         }
+      }
     }
+  }
 }
 
-
-void RecoverDecodeTask(const paddle::Tensor &stop_flags,
-                   const paddle::Tensor &seq_lens_this_time,
-                   const paddle::Tensor &seq_lens_encoder,
-                   const paddle::Tensor &seq_lens_decoder,
-                   const paddle::Tensor &step_seq_lens_decoder,
-                   const paddle::Tensor &block_tables,
-                   const paddle::Tensor &is_block_step,
-                   const paddle::optional<paddle::Tensor> &draft_tokens,
-                   const paddle::optional<paddle::Tensor> &step_draft_tokens,
-                   const paddle::optional<paddle::Tensor> &step_seq_lens_this_time,
-                   const int block_size,
-                   const int max_draft_tokens) {
+void RecoverDecodeTask(
+    const paddle::Tensor &stop_flags,
+    const paddle::Tensor &seq_lens_this_time,
+    const paddle::Tensor &seq_lens_encoder,
+    const paddle::Tensor &seq_lens_decoder,
+    const paddle::Tensor &step_seq_lens_decoder,
+    const paddle::Tensor &block_tables,
+    const paddle::Tensor &is_block_step,
+    const paddle::optional<paddle::Tensor> &draft_tokens,
+    const paddle::optional<paddle::Tensor> &step_draft_tokens,
+    const paddle::optional<paddle::Tensor> &step_seq_lens_this_time,
+    const int block_size,
+    const int max_draft_tokens) {
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
-    auto dev_ctx = static_cast<const phi::CustomContext*>(paddle::experimental::DeviceContextPool::Instance().Get(seq_lens_this_time.place()));
-    auto cu_stream = dev_ctx->stream();
+  auto dev_ctx = static_cast<const phi::CustomContext *>(
+      paddle::experimental::DeviceContextPool::Instance().Get(
+          seq_lens_this_time.place()));
+  auto cu_stream = dev_ctx->stream();
 #else
-    auto cu_stream = seq_lens_this_time.stream();
+  auto cu_stream = seq_lens_this_time.stream();
 #endif
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    if (draft_tokens) {
-        const int draft_tokens_len = draft_tokens.get_ptr()->shape()[1];
-        recover_spec_decode_task<<<1, 1024, 0, cu_stream>>>(
-            const_cast<bool *>(stop_flags.data<bool>()),
-            const_cast<int *>(seq_lens_this_time.data<int>()),
-            const_cast<int *>(seq_lens_encoder.data<int>()),
-            const_cast<int *>(seq_lens_decoder.data<int>()),
-            const_cast<int *>(step_seq_lens_decoder.data<int>()),
-            const_cast<int *>(block_tables.data<int>()),
-            const_cast<bool *>(is_block_step.data<bool>()),
-            const_cast<int64_t *>(draft_tokens.get_ptr()->data<int64_t>()),
-            step_draft_tokens.get_ptr()->data<int64_t>(),
-            step_seq_lens_this_time.get_ptr()->data<int>(),
-            bsz,
-            block_num_per_seq,
-            block_size,
-            draft_tokens_len,
-            max_draft_tokens * 2 + 1);
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  if (draft_tokens) {
+    const int draft_tokens_len = draft_tokens.get_ptr()->shape()[1];
+    recover_spec_decode_task<<<1, 1024, 0, cu_stream>>>(
+        const_cast<bool *>(stop_flags.data<bool>()),
+        const_cast<int *>(seq_lens_this_time.data<int>()),
+        const_cast<int *>(seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_decoder.data<int>()),
+        const_cast<int *>(step_seq_lens_decoder.data<int>()),
+        const_cast<int *>(block_tables.data<int>()),
+        const_cast<bool *>(is_block_step.data<bool>()),
+        const_cast<int64_t *>(draft_tokens.get_ptr()->data<int64_t>()),
+        step_draft_tokens.get_ptr()->data<int64_t>(),
+        step_seq_lens_this_time.get_ptr()->data<int>(),
+        bsz,
+        block_num_per_seq,
+        block_size,
+        draft_tokens_len,
+        max_draft_tokens * 2 + 1);
 
-    } else {
-        recover_decode_task<<<1, 1024, 0, cu_stream>>>(
-            const_cast<bool *>(stop_flags.data<bool>()),
-            const_cast<int *>(seq_lens_this_time.data<int>()),
-            const_cast<int *>(seq_lens_encoder.data<int>()),
-            const_cast<int *>(seq_lens_decoder.data<int>()),
-            const_cast<int *>(step_seq_lens_decoder.data<int>()),
-            const_cast<int *>(block_tables.data<int>()),
-            const_cast<bool *>(is_block_step.data<bool>()),
-            bsz,
-            block_num_per_seq,
-            block_size);
-    }
+  } else {
+    recover_decode_task<<<1, 1024, 0, cu_stream>>>(
+        const_cast<bool *>(stop_flags.data<bool>()),
+        const_cast<int *>(seq_lens_this_time.data<int>()),
+        const_cast<int *>(seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_decoder.data<int>()),
+        const_cast<int *>(step_seq_lens_decoder.data<int>()),
+        const_cast<int *>(block_tables.data<int>()),
+        const_cast<bool *>(is_block_step.data<bool>()),
+        bsz,
+        block_num_per_seq,
+        block_size);
+  }
 }
 
 PD_BUILD_STATIC_OP(recover_decode_task)
diff --git a/custom_ops/gpu_ops/remote_cache_kv_ipc.cc b/custom_ops/gpu_ops/remote_cache_kv_ipc.cc
index f1f53513b8..2c34d49b74 100644
--- a/custom_ops/gpu_ops/remote_cache_kv_ipc.cc
+++ b/custom_ops/gpu_ops/remote_cache_kv_ipc.cc
@@ -14,107 +14,111 @@
 
 #include "remote_cache_kv_ipc.h"
 
-RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data RemoteCacheKvIpc::kv_complete_signal_meta_data;
+RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data
+    RemoteCacheKvIpc::kv_complete_signal_meta_data;
 RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data_per_query
     RemoteCacheKvIpc::kv_complete_signal_meta_data_per_query;
 void* RemoteCacheKvIpc::kv_complete_signal_identity_ptr = nullptr;
 bool RemoteCacheKvIpc::kv_complete_signal_shmem_opened = false;
 
 RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data
-                            RemoteCacheKvIpc::open_shm_and_get_complete_signal_meta_data(
-                                                            const int rank_id,
-                                                            const int device_id,
-                                                            const bool keep_pd_step_flag) {
-    if (RemoteCacheKvIpc::kv_complete_signal_shmem_opened){
-        if (keep_pd_step_flag) {
-            return RemoteCacheKvIpc::kv_complete_signal_meta_data;
-        }
-        RemoteCacheKvIpc::kv_complete_signal_meta_data.layer_id = -1;
-        int32_t* layer_complete_ptr = reinterpret_cast<int32_t*>(kv_complete_signal_meta_data.shm_ptr);
-        *layer_complete_ptr = -1;
-        int32_t current_identity = (*reinterpret_cast<int32_t*>(RemoteCacheKvIpc::kv_complete_signal_identity_ptr));
-        int32_t* write_ptr = reinterpret_cast<int32_t*>(RemoteCacheKvIpc::kv_complete_signal_identity_ptr);
-        *write_ptr = (current_identity + 1) % 100003;
-        return RemoteCacheKvIpc::kv_complete_signal_meta_data;
+RemoteCacheKvIpc::open_shm_and_get_complete_signal_meta_data(
+    const int rank_id, const int device_id, const bool keep_pd_step_flag) {
+  if (RemoteCacheKvIpc::kv_complete_signal_shmem_opened) {
+    if (keep_pd_step_flag) {
+      return RemoteCacheKvIpc::kv_complete_signal_meta_data;
     }
-
-    std::string flags_server_uuid;
-    if (const char* iflags_server_uuid_env_p = std::getenv("SHM_UUID")){
-        std::string iflags_server_uuid_env_str(iflags_server_uuid_env_p);
-        flags_server_uuid = iflags_server_uuid_env_str;
-    }
-
-    std::string step_shm_name = ("splitwise_complete_prefilled_step_"
-                    + std::to_string(rank_id) + "." + std::to_string(device_id));
-    std::string layer_shm_name = ("splitwise_complete_prefilled_layer_"
-                    + std::to_string(rank_id) + "." + std::to_string(device_id));
-    if (const char* use_ep = std::getenv("ENABLE_EP_DP")){
-        if(std::strcmp(use_ep, "1") == 0){
-        step_shm_name = "splitwise_complete_prefilled_step_tprank0_dprank"
-                + std::to_string(rank_id) + "_" + flags_server_uuid;
-        layer_shm_name = "splitwise_complete_prefilled_layer_tprank0_dprank"
-                + std::to_string(rank_id) + "_" + flags_server_uuid;
-        }
-    }
-
-    int signal_shm_fd = shm_open(layer_shm_name.c_str(), O_CREAT | O_RDWR, 0666);
-
-    PADDLE_ENFORCE_NE(signal_shm_fd,
-                        -1,
-                        phi::errors::InvalidArgument(
-                            "can not open shm for cache_kv_complete_signal."));
-    int signal_shm_ftruncate = ftruncate(signal_shm_fd, 4);
-    void* signal_ptr = mmap(0, 4, PROT_WRITE, MAP_SHARED, signal_shm_fd, 0);
-
-    PADDLE_ENFORCE_NE(
-        signal_ptr,
-        MAP_FAILED,
-        phi::errors::InvalidArgument(
-                            "MAP_FAILED for cache_kv_complete_signal."));
-    int32_t* write_signal_ptr = reinterpret_cast<int32_t*>(signal_ptr);
-    *write_signal_ptr = -1;
-    using type_meta_data = RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data;
-
-    // std::printf("#### open_shm_and_get_complete_signal_meta_data layer idx:%d, to ptx:%p \n",
-    //             -1, signal_ptr);
-
-    type_meta_data meta_data(
-        -1,
-        signal_ptr,
-        signal_shm_fd
-    );
-    RemoteCacheKvIpc::kv_complete_signal_meta_data = meta_data;
-    int identity_shm_fd = shm_open(step_shm_name.c_str(), O_CREAT | O_RDWR, 0666);
-    PADDLE_ENFORCE_NE(identity_shm_fd,
-                        -1,
-                        phi::errors::InvalidArgument(
-                            "can not open shm for cache_kv_complete_identity."));
-
-    int identity_shm_ftruncate = ftruncate(identity_shm_fd, 4);
-    void* identity_ptr = mmap(0, 4, PROT_WRITE, MAP_SHARED, identity_shm_fd, 0);
-    PADDLE_ENFORCE_NE(
-        identity_ptr,
-        MAP_FAILED,
-        phi::errors::InvalidArgument(
-                            "MAP_FAILED for prefill_identity."));
-    int32_t current_identity = (*reinterpret_cast<int32_t*>(identity_ptr));
-    int32_t* write_ptr = reinterpret_cast<int32_t*>(identity_ptr);
+    RemoteCacheKvIpc::kv_complete_signal_meta_data.layer_id = -1;
+    int32_t* layer_complete_ptr =
+        reinterpret_cast<int32_t*>(kv_complete_signal_meta_data.shm_ptr);
+    *layer_complete_ptr = -1;
+    int32_t current_identity = (*reinterpret_cast<int32_t*>(
+        RemoteCacheKvIpc::kv_complete_signal_identity_ptr));
+    int32_t* write_ptr = reinterpret_cast<int32_t*>(
+        RemoteCacheKvIpc::kv_complete_signal_identity_ptr);
     *write_ptr = (current_identity + 1) % 100003;
-    RemoteCacheKvIpc::kv_complete_signal_identity_ptr = identity_ptr;
-    RemoteCacheKvIpc::kv_complete_signal_shmem_opened = true;
-    return meta_data;
+    return RemoteCacheKvIpc::kv_complete_signal_meta_data;
+  }
+
+  std::string flags_server_uuid;
+  if (const char* iflags_server_uuid_env_p = std::getenv("SHM_UUID")) {
+    std::string iflags_server_uuid_env_str(iflags_server_uuid_env_p);
+    flags_server_uuid = iflags_server_uuid_env_str;
+  }
+
+  std::string step_shm_name =
+      ("splitwise_complete_prefilled_step_" + std::to_string(rank_id) + "." +
+       std::to_string(device_id));
+  std::string layer_shm_name =
+      ("splitwise_complete_prefilled_layer_" + std::to_string(rank_id) + "." +
+       std::to_string(device_id));
+  if (const char* use_ep = std::getenv("ENABLE_EP_DP")) {
+    if (std::strcmp(use_ep, "1") == 0) {
+      step_shm_name = "splitwise_complete_prefilled_step_tprank0_dprank" +
+                      std::to_string(rank_id) + "_" + flags_server_uuid;
+      layer_shm_name = "splitwise_complete_prefilled_layer_tprank0_dprank" +
+                       std::to_string(rank_id) + "_" + flags_server_uuid;
+    }
+  }
+
+  int signal_shm_fd = shm_open(layer_shm_name.c_str(), O_CREAT | O_RDWR, 0666);
+
+  PADDLE_ENFORCE_NE(signal_shm_fd,
+                    -1,
+                    phi::errors::InvalidArgument(
+                        "can not open shm for cache_kv_complete_signal."));
+  int signal_shm_ftruncate = ftruncate(signal_shm_fd, 4);
+  void* signal_ptr = mmap(0, 4, PROT_WRITE, MAP_SHARED, signal_shm_fd, 0);
+
+  PADDLE_ENFORCE_NE(
+      signal_ptr,
+      MAP_FAILED,
+      phi::errors::InvalidArgument("MAP_FAILED for cache_kv_complete_signal."));
+  int32_t* write_signal_ptr = reinterpret_cast<int32_t*>(signal_ptr);
+  *write_signal_ptr = -1;
+  using type_meta_data =
+      RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data;
+
+  // std::printf("#### open_shm_and_get_complete_signal_meta_data layer idx:%d,
+  // to ptx:%p \n",
+  //             -1, signal_ptr);
+
+  type_meta_data meta_data(-1, signal_ptr, signal_shm_fd);
+  RemoteCacheKvIpc::kv_complete_signal_meta_data = meta_data;
+  int identity_shm_fd = shm_open(step_shm_name.c_str(), O_CREAT | O_RDWR, 0666);
+  PADDLE_ENFORCE_NE(identity_shm_fd,
+                    -1,
+                    phi::errors::InvalidArgument(
+                        "can not open shm for cache_kv_complete_identity."));
+
+  int identity_shm_ftruncate = ftruncate(identity_shm_fd, 4);
+  void* identity_ptr = mmap(0, 4, PROT_WRITE, MAP_SHARED, identity_shm_fd, 0);
+  PADDLE_ENFORCE_NE(
+      identity_ptr,
+      MAP_FAILED,
+      phi::errors::InvalidArgument("MAP_FAILED for prefill_identity."));
+  int32_t current_identity = (*reinterpret_cast<int32_t*>(identity_ptr));
+  int32_t* write_ptr = reinterpret_cast<int32_t*>(identity_ptr);
+  *write_ptr = (current_identity + 1) % 100003;
+  RemoteCacheKvIpc::kv_complete_signal_identity_ptr = identity_ptr;
+  RemoteCacheKvIpc::kv_complete_signal_shmem_opened = true;
+  return meta_data;
 }
 
-void CUDART_CB RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise(void* meta_data) {
-    int64_t* meta_data_ptr = reinterpret_cast<int64_t*>(meta_data);
-    int32_t layer_id = meta_data_ptr[0];
-    int32_t* ptr = reinterpret_cast<int32_t*>(meta_data_ptr[1]);
-    *ptr = layer_id;
-    // std::printf("#### save_cache_kv_complete_signal_layerwise layer idx:%d, to ptx:%p \n",
-    //             *ptr, meta_data_ptr[1]);
+void CUDART_CB
+RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise(void* meta_data) {
+  int64_t* meta_data_ptr = reinterpret_cast<int64_t*>(meta_data);
+  int32_t layer_id = meta_data_ptr[0];
+  int32_t* ptr = reinterpret_cast<int32_t*>(meta_data_ptr[1]);
+  *ptr = layer_id;
+  // std::printf("#### save_cache_kv_complete_signal_layerwise layer idx:%d, to
+  // ptx:%p \n",
+  //             *ptr, meta_data_ptr[1]);
 }
 
-void CUDART_CB RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_per_query(void* meta_data) {
-    RemoteCacheKvIpc::kv_complete_signal_meta_data_per_query.send_signal();
-    // std::printf("#### save_cache_kv_complete_signal_layerwise_per_query);
+void CUDART_CB
+RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_per_query(
+    void* meta_data) {
+  RemoteCacheKvIpc::kv_complete_signal_meta_data_per_query.send_signal();
+  // std::printf("#### save_cache_kv_complete_signal_layerwise_per_query);
 }
diff --git a/custom_ops/gpu_ops/sample_kernels/min_p_sampling_from_probs.cu b/custom_ops/gpu_ops/sample_kernels/min_p_sampling_from_probs.cu
index c44c16b430..02d6d8c201 100644
--- a/custom_ops/gpu_ops/sample_kernels/min_p_sampling_from_probs.cu
+++ b/custom_ops/gpu_ops/sample_kernels/min_p_sampling_from_probs.cu
@@ -17,46 +17,45 @@
 #include "sample_kernels/sampling.cuh"
 
 std::vector<paddle::Tensor> MinPSamplingFromProbs(const paddle::Tensor &probs,
-                                               const paddle::Tensor &min_p) {
-    std::vector<int64_t> probs_shape = probs.shape();
-    unsigned int batch_size = probs_shape[0];
-    unsigned int vocab_size = probs_shape[1];
-    auto cu_stream = probs.stream();
+                                                  const paddle::Tensor &min_p) {
+  std::vector<int64_t> probs_shape = probs.shape();
+  unsigned int batch_size = probs_shape[0];
+  unsigned int vocab_size = probs_shape[1];
+  auto cu_stream = probs.stream();
 
-    auto renorm_probs =
+  auto renorm_probs =
       GetEmptyTensor(probs.dims(), paddle::DataType::FLOAT32, probs.place());
 
-    cudaError_t status;
+  cudaError_t status;
 
-    status = sampling::MinPSamplingFromProb<float, int>(
-        const_cast<float *>(probs.data<float>()),
-        const_cast<float *>(min_p.data<float>()),
-        renorm_probs.data<float>(),
-        batch_size,
-        vocab_size,
-        true,  // deterministic
-        cu_stream);
+  status = sampling::MinPSamplingFromProb<float, int>(
+      const_cast<float *>(probs.data<float>()),
+      const_cast<float *>(min_p.data<float>()),
+      renorm_probs.data<float>(),
+      batch_size,
+      vocab_size,
+      true,  // deterministic
+      cu_stream);
 
-
-  PD_CHECK(status == cudaSuccess, "SamplingFromProbs failed with error code " +
-                                      std::string(cudaGetErrorString(status)));
+  PD_CHECK(status == cudaSuccess,
+           "SamplingFromProbs failed with error code " +
+               std::string(cudaGetErrorString(status)));
 
   return {renorm_probs};
 }
 
-std::vector<std::vector<int64_t>>
-MinPSamplingFromProbsInferShape(const std::vector<int64_t> &probs_shape,
-                             const paddle::optional<std::vector<int64_t>> &min_p_shape) {
+std::vector<std::vector<int64_t>> MinPSamplingFromProbsInferShape(
+    const std::vector<int64_t> &probs_shape,
+    const paddle::optional<std::vector<int64_t>> &min_p_shape) {
   return {probs_shape};
 }
 
-std::vector<paddle::DataType>
-MinPSamplingFromProbsInferDtype(const paddle::DataType &probs_dtype,
-                             const paddle::optional<paddle::DataType> &min_p_dtype) {
+std::vector<paddle::DataType> MinPSamplingFromProbsInferDtype(
+    const paddle::DataType &probs_dtype,
+    const paddle::optional<paddle::DataType> &min_p_dtype) {
   return {probs_dtype};
 }
 
-
 PD_BUILD_STATIC_OP(min_p_sampling)
     .Inputs({"probs", "min_p"})
     .Outputs({"renorm_probs"})
diff --git a/custom_ops/gpu_ops/sample_kernels/rejection_top_p_sampling.cu b/custom_ops/gpu_ops/sample_kernels/rejection_top_p_sampling.cu
index 99c87d36f4..1716d67cb3 100644
--- a/custom_ops/gpu_ops/sample_kernels/rejection_top_p_sampling.cu
+++ b/custom_ops/gpu_ops/sample_kernels/rejection_top_p_sampling.cu
@@ -16,10 +16,11 @@
 #include "paddle/phi/backends/context_pool.h"
 #include "sample_kernels/sampling.cuh"
 
-std::vector<paddle::Tensor> TopPSamplingReject(const paddle::Tensor &probs,
-                                               const paddle::Tensor &top_p,
-                                               const paddle::optional<paddle::Tensor> &top_k,
-                                               int64_t seed) {
+std::vector<paddle::Tensor> TopPSamplingReject(
+    const paddle::Tensor &probs,
+    const paddle::Tensor &top_p,
+    const paddle::optional<paddle::Tensor> &top_k,
+    int64_t seed) {
   std::vector<int64_t> probs_shape = probs.shape();
   unsigned int batch_size = probs_shape[0];
   unsigned int vocab_size = probs_shape[1];
@@ -30,9 +31,11 @@ std::vector<paddle::Tensor> TopPSamplingReject(const paddle::Tensor &probs,
   // need_batch_random
   if (seed == -1) {
 #ifdef PADDLE_WITH_COREX
-    auto dev_ctx = static_cast<const phi::CustomContext*>(paddle::experimental::DeviceContextPool::Instance().Get(probs.place()));
+    auto dev_ctx = static_cast<const phi::CustomContext *>(
+        paddle::experimental::DeviceContextPool::Instance().Get(probs.place()));
 #else
-    phi::GPUContext* dev_ctx = static_cast<phi::GPUContext*>(phi::DeviceContextPool::Instance().Get(probs.place()));
+    phi::GPUContext *dev_ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(probs.place()));
 #endif
     auto gen_cuda = dev_ctx->GetGenerator();
     auto seed_offset = gen_cuda->IncrementOffset(32 * batch_size);
@@ -47,35 +50,48 @@ std::vector<paddle::Tensor> TopPSamplingReject(const paddle::Tensor &probs,
 
   if (top_k) {
     status = sampling::TopKTopPSamplingFromProb<float, int64_t>(
-        const_cast<float *>(probs.data<float>()), samples.data<int64_t>(),
-        batch_size, top_p.data<float>(), top_k.get().data<int64_t>(), vocab_size,
-        true, philox_seed, philox_offset, cu_stream);
-  }
-  else {
+        const_cast<float *>(probs.data<float>()),
+        samples.data<int64_t>(),
+        batch_size,
+        top_p.data<float>(),
+        top_k.get().data<int64_t>(),
+        vocab_size,
+        true,
+        philox_seed,
+        philox_offset,
+        cu_stream);
+  } else {
     status = sampling::TopPSamplingFromProb<float, int64_t>(
-        const_cast<float *>(probs.data<float>()), samples.data<int64_t>(),
-        batch_size, top_p.data<float>(), vocab_size,
-        true, philox_seed, philox_offset, cu_stream);
+        const_cast<float *>(probs.data<float>()),
+        samples.data<int64_t>(),
+        batch_size,
+        top_p.data<float>(),
+        vocab_size,
+        true,
+        philox_seed,
+        philox_offset,
+        cu_stream);
   }
 
-  PD_CHECK(status == cudaSuccess, "SamplingFromProbs failed with error code " +
-                                      std::string(cudaGetErrorString(status)));
+  PD_CHECK(status == cudaSuccess,
+           "SamplingFromProbs failed with error code " +
+               std::string(cudaGetErrorString(status)));
 
   return {samples};
 }
 
-std::vector<std::vector<int64_t>>
-TopPSamplingRejectInferShape(const std::vector<int64_t> &probs_shape,
-                             const std::vector<int64_t> &top_p_shape,
-                             const paddle::optional<std::vector<int64_t>> &top_k_shape) {
+std::vector<std::vector<int64_t>> TopPSamplingRejectInferShape(
+    const std::vector<int64_t> &probs_shape,
+    const std::vector<int64_t> &top_p_shape,
+    const paddle::optional<std::vector<int64_t>> &top_k_shape) {
   int64_t bs = probs_shape[0];
   return {{bs, 1}};
 }
 
-std::vector<paddle::DataType>
-TopPSamplingRejectInferDtype(const paddle::DataType &probs_dtype,
-                             const paddle::DataType &top_p_dtype,
-                             const paddle::optional<paddle::DataType> &top_k_dtype) {
+std::vector<paddle::DataType> TopPSamplingRejectInferDtype(
+    const paddle::DataType &probs_dtype,
+    const paddle::DataType &top_p_dtype,
+    const paddle::optional<paddle::DataType> &top_k_dtype) {
   return {paddle::DataType::INT64};
 }
 
diff --git a/custom_ops/gpu_ops/sample_kernels/top_k_renorm_probs.cu b/custom_ops/gpu_ops/sample_kernels/top_k_renorm_probs.cu
index ea4ab0dbb6..6d7945f0a4 100644
--- a/custom_ops/gpu_ops/sample_kernels/top_k_renorm_probs.cu
+++ b/custom_ops/gpu_ops/sample_kernels/top_k_renorm_probs.cu
@@ -28,28 +28,29 @@ std::vector<paddle::Tensor> TopKRenorm(const paddle::Tensor &probs,
 
   cudaError_t status;
 
-
   status = sampling::TopKRenormProb<float>(
-    const_cast<float *>(probs.data<float>()),
-    renorm_probs.data<float>(),
-    const_cast<int64_t *>(top_k.data<int64_t>()),
-    batch_size, vocab_size, cu_stream);
+      const_cast<float *>(probs.data<float>()),
+      renorm_probs.data<float>(),
+      const_cast<int64_t *>(top_k.data<int64_t>()),
+      batch_size,
+      vocab_size,
+      cu_stream);
 
-  PD_CHECK(status == cudaSuccess, "TopKRenormProb failed with error code " +
-                                      std::string(cudaGetErrorString(status)));
+  PD_CHECK(status == cudaSuccess,
+           "TopKRenormProb failed with error code " +
+               std::string(cudaGetErrorString(status)));
 
   return {renorm_probs};
 }
 
-std::vector<std::vector<int64_t>>
-TopKRenormInferShape(const std::vector<int64_t> &probs_shape,
-                    const std::vector<int64_t> &top_k_shape) {
+std::vector<std::vector<int64_t>> TopKRenormInferShape(
+    const std::vector<int64_t> &probs_shape,
+    const std::vector<int64_t> &top_k_shape) {
   return {probs_shape};
 }
 
-std::vector<paddle::DataType>
-TopKRenormInferDtype(const paddle::DataType &probs_dtype,
-                    const paddle::DataType &top_k_shape) {
+std::vector<paddle::DataType> TopKRenormInferDtype(
+    const paddle::DataType &probs_dtype, const paddle::DataType &top_k_shape) {
   return {probs_dtype};
 }
 
diff --git a/custom_ops/gpu_ops/save_with_output.cc b/custom_ops/gpu_ops/save_with_output.cc
index 598dc3549d..7d87056f42 100644
--- a/custom_ops/gpu_ops/save_with_output.cc
+++ b/custom_ops/gpu_ops/save_with_output.cc
@@ -32,33 +32,33 @@
 constexpr char kSEP = '/';
 
 std::string DirName(const std::string& filepath) {
-    auto pos = filepath.rfind(kSEP);
-    if (pos == std::string::npos) {
-        return "";
-    }
-    return filepath.substr(0, pos);
+  auto pos = filepath.rfind(kSEP);
+  if (pos == std::string::npos) {
+    return "";
+  }
+  return filepath.substr(0, pos);
 }
 
 bool FileExists(const std::string& filepath) {
-    struct stat buffer;
-    return (stat(filepath.c_str(), &buffer) == 0);
+  struct stat buffer;
+  return (stat(filepath.c_str(), &buffer) == 0);
 }
 
 void MkDir(const char* path) {
-    std::string path_error(path);
-    path_error += " mkdir failed!";
-    if (mkdir(path, 0755)) {
-        if (errno != EEXIST) {
-            throw std::runtime_error(path_error);
-        }
+  std::string path_error(path);
+  path_error += " mkdir failed!";
+  if (mkdir(path, 0755)) {
+    if (errno != EEXIST) {
+      throw std::runtime_error(path_error);
     }
+  }
 }
 
 void MkDirRecursively(const char* fullpath) {
-    if (*fullpath == '\0') return;  // empty string
-    if (FileExists(fullpath)) return;
-    MkDirRecursively(DirName(fullpath).c_str());
-    MkDir(fullpath);
+  if (*fullpath == '\0') return;  // empty string
+  if (FileExists(fullpath)) return;
+  MkDirRecursively(DirName(fullpath).c_str());
+  MkDir(fullpath);
 }
 
 template <typename data_t>
@@ -67,12 +67,12 @@ void saveToFile(std::ostream& os,
                 std::vector<int64_t> shape,
                 int64_t x_numel,
                 const char type_id) {
-    // 1.type
-    os.write(reinterpret_cast<const char*>(&type_id), sizeof(type_id));
-    // 2.data
-    uint64_t size = x_numel * sizeof(data_t);
-    os.write(static_cast<const char*>(x_data),
-             static_cast<std::streamsize>(size));
+  // 1.type
+  os.write(reinterpret_cast<const char*>(&type_id), sizeof(type_id));
+  // 2.data
+  uint64_t size = x_numel * sizeof(data_t);
+  os.write(static_cast<const char*>(x_data),
+           static_cast<std::streamsize>(size));
 }
 
 template <typename data_t>
@@ -82,52 +82,51 @@ void save_with_output_kernel(const paddle::Tensor& x,
                              std::string file_path,
                              int64_t rank_id,
                              char type_id) {
-    std::vector<int64_t> x_shape = x.shape();
+  std::vector<int64_t> x_shape = x.shape();
 
-    if (rank_id >= 0) {
-        file_path += "_rank_" + std::to_string(rank_id);
-    }
+  if (rank_id >= 0) {
+    file_path += "_rank_" + std::to_string(rank_id);
+  }
 
-    int batch_idx_data = -1, step_idx_data = -1;
+  int batch_idx_data = -1, step_idx_data = -1;
 
-    if (batch_idx.is_gpu()) {
-        paddle::Tensor batch_idx_cpu =
-            batch_idx.copy_to<int32_t>(paddle::CPUPlace());
-        batch_idx_data = batch_idx_cpu.data<int32_t>()[0];
-    } else {
-        batch_idx_data = batch_idx.data<int32_t>()[0];
-    }
-    if (step_idx.is_gpu()) {
-        paddle::Tensor step_idx_cpu =
-            step_idx.copy_to<int64_t>(paddle::CPUPlace());
-        step_idx_data = step_idx_cpu.data<int64_t>()[0];
-    } else {
-        step_idx_data = step_idx.data<int64_t>()[0];
-    }
-    auto x_data = x.data<data_t>();
+  if (batch_idx.is_gpu()) {
+    paddle::Tensor batch_idx_cpu =
+        batch_idx.copy_to<int32_t>(paddle::CPUPlace());
+    batch_idx_data = batch_idx_cpu.data<int32_t>()[0];
+  } else {
+    batch_idx_data = batch_idx.data<int32_t>()[0];
+  }
+  if (step_idx.is_gpu()) {
+    paddle::Tensor step_idx_cpu = step_idx.copy_to<int64_t>(paddle::CPUPlace());
+    step_idx_data = step_idx_cpu.data<int64_t>()[0];
+  } else {
+    step_idx_data = step_idx.data<int64_t>()[0];
+  }
+  auto x_data = x.data<data_t>();
 
-    if (batch_idx_data >= 0) {
-        file_path += "_batch_" + std::to_string(batch_idx_data);
-    }
-    if (step_idx_data >= 0) {
-        file_path += "_step_" + std::to_string(step_idx_data);
-    }
-    MkDirRecursively(DirName(file_path).c_str());
-    std::ofstream fout(file_path, std::ios::binary);
-    fout.write("0", 1);
-    saveToFile<data_t>(fout, x_data, x_shape, x.numel(), type_id);
-    fout.seekp(std::ios::beg);
-    fout.write("1", 1);
-    fout.close();
+  if (batch_idx_data >= 0) {
+    file_path += "_batch_" + std::to_string(batch_idx_data);
+  }
+  if (step_idx_data >= 0) {
+    file_path += "_step_" + std::to_string(step_idx_data);
+  }
+  MkDirRecursively(DirName(file_path).c_str());
+  std::ofstream fout(file_path, std::ios::binary);
+  fout.write("0", 1);
+  saveToFile<data_t>(fout, x_data, x_shape, x.numel(), type_id);
+  fout.seekp(std::ios::beg);
+  fout.write("1", 1);
+  fout.close();
 }
 
 void print_shape(const paddle::Tensor& tmp, char* tmp_str) {
-    std::vector<int64_t> shape = tmp.shape();
-    printf("%s's shape: \n", tmp_str);
-    for (int i = 0; i < shape.size(); i++) {
-        printf("%d ", (int)shape[i]);
-    }
-    printf("\n");
+  std::vector<int64_t> shape = tmp.shape();
+  printf("%s's shape: \n", tmp_str);
+  for (int i = 0; i < shape.size(); i++) {
+    printf("%d ", (int)shape[i]);
+  }
+  printf("\n");
 }
 
 std::vector<paddle::Tensor> SaveWithOutputForward(
@@ -136,40 +135,40 @@ std::vector<paddle::Tensor> SaveWithOutputForward(
     const paddle::Tensor& step_idx,
     std::string file_path,
     int64_t rank_id) {
-    auto out = x.copy_to(paddle::CPUPlace(), false);
-    switch (x.type()) {
-        case paddle::DataType::FLOAT32:
-            save_with_output_kernel<float>(
-                out, batch_idx, step_idx, file_path, rank_id, '0');
-            break;
-        case paddle::DataType::INT64:
-            save_with_output_kernel<int64_t>(
-                out, batch_idx, step_idx, file_path, rank_id, '1');
-            break;
-        case paddle::DataType::INT32:
-            save_with_output_kernel<int32_t>(
-                out, batch_idx, step_idx, file_path, rank_id, '2');
-            break;
-        default:
-            PD_THROW(
-                "function SaveWithOutputForward is not implemented for data "
-                "type");
-    }
-    return {out};
+  auto out = x.copy_to(paddle::CPUPlace(), false);
+  switch (x.type()) {
+    case paddle::DataType::FLOAT32:
+      save_with_output_kernel<float>(
+          out, batch_idx, step_idx, file_path, rank_id, '0');
+      break;
+    case paddle::DataType::INT64:
+      save_with_output_kernel<int64_t>(
+          out, batch_idx, step_idx, file_path, rank_id, '1');
+      break;
+    case paddle::DataType::INT32:
+      save_with_output_kernel<int32_t>(
+          out, batch_idx, step_idx, file_path, rank_id, '2');
+      break;
+    default:
+      PD_THROW(
+          "function SaveWithOutputForward is not implemented for data "
+          "type");
+  }
+  return {out};
 }
 
 std::vector<std::vector<int64_t>> SaveWithOutputInferShape(
     const std::vector<int64_t>& x_shape,
     const std::vector<int64_t>& batch_idx_shape,
     const std::vector<int64_t>& step_idx_shape) {
-    return {x_shape};
+  return {x_shape};
 }
 
 std::vector<paddle::DataType> SaveWithOutputInferDtype(
     const paddle::DataType& x_dtype,
     const paddle::DataType& batch_idx_dtype,
     const paddle::DataType& step_idx_dtype) {
-    return {x_dtype};
+  return {x_dtype};
 }
 
 PD_BUILD_STATIC_OP(save_with_output)
diff --git a/custom_ops/gpu_ops/save_with_output_msg.h b/custom_ops/gpu_ops/save_with_output_msg.h
index 77c8231fa9..98b6f94b7a 100644
--- a/custom_ops/gpu_ops/save_with_output_msg.h
+++ b/custom_ops/gpu_ops/save_with_output_msg.h
@@ -28,8 +28,8 @@
 #define MAX_BSZ 512
 
 struct msgdata {
-    long mtype;
-    int mtext[MAX_BSZ + 2];  // stop_flag, bsz, tokens
+  long mtype;
+  int mtext[MAX_BSZ + 2];  // stop_flag, bsz, tokens
 };
 
 void SaveOutMmsg(const paddle::Tensor& x,
diff --git a/custom_ops/gpu_ops/scaled_gemm_f8_i4_f16_gemm.cu b/custom_ops/gpu_ops/scaled_gemm_f8_i4_f16_gemm.cu
index c8d3ca8a39..fbf63cc02c 100644
--- a/custom_ops/gpu_ops/scaled_gemm_f8_i4_f16_gemm.cu
+++ b/custom_ops/gpu_ops/scaled_gemm_f8_i4_f16_gemm.cu
@@ -23,225 +23,300 @@
 using Pointer = void*;
 using ConstPointer = void const*;
 
-struct Params
-{
-    Pointer act;
-    Pointer weight;
-    Pointer scales;
-    Pointer zeros;
-    Pointer bias;
-    Pointer out;
-    float alpha;
-    int m;
-    int n;
-    int k;
-    int groupsize;
-    bool apply_alpha_in_advance;
+struct Params {
+  Pointer act;
+  Pointer weight;
+  Pointer scales;
+  Pointer zeros;
+  Pointer bias;
+  Pointer out;
+  float alpha;
+  int m;
+  int n;
+  int k;
+  int groupsize;
+  bool apply_alpha_in_advance;
 
-    Params(ConstPointer _act, ConstPointer _weight, ConstPointer _scales, ConstPointer _zeros,
-        ConstPointer _bias, Pointer _out, float _alpha, int _m, int _n, int _k, int _groupsize,
-        bool _apply_alpha_in_advance = false)
-        : act(const_cast<Pointer>(_act))
-        , weight(const_cast<Pointer>(_weight))
-        , scales(const_cast<Pointer>(_scales))
-        , zeros(const_cast<Pointer>(_zeros))
-        , bias(const_cast<Pointer>(_bias))
-        , out(_out)
-        , alpha(_alpha)
-        , m(_m)
-        , n(_n)
-        , k(_k)
-        , groupsize(_groupsize)
-        , apply_alpha_in_advance(_apply_alpha_in_advance)
-    {
-    }
+  Params(ConstPointer _act,
+         ConstPointer _weight,
+         ConstPointer _scales,
+         ConstPointer _zeros,
+         ConstPointer _bias,
+         Pointer _out,
+         float _alpha,
+         int _m,
+         int _n,
+         int _k,
+         int _groupsize,
+         bool _apply_alpha_in_advance = false)
+      : act(const_cast<Pointer>(_act)),
+        weight(const_cast<Pointer>(_weight)),
+        scales(const_cast<Pointer>(_scales)),
+        zeros(const_cast<Pointer>(_zeros)),
+        bias(const_cast<Pointer>(_bias)),
+        out(_out),
+        alpha(_alpha),
+        m(_m),
+        n(_n),
+        k(_k),
+        groupsize(_groupsize),
+        apply_alpha_in_advance(_apply_alpha_in_advance) {}
 };
 
 template <typename T, cutlass::WeightOnlyQuantOp QuantOp>
 void scaled_gemm_f8_i4_f16_launcher(Params& params,
                                     phi::Place place,
                                     cudaStream_t stream) {
-    auto runner = std::make_shared<kernels::cutlass_kernels::CutlassFpAIntBGemmRunner<__nv_fp8_e4m3, cutlass::uint4b_t, QuantOp, half, T, T>>();
-    int ws_size = runner->getWorkspaceSize(params.m, params.n, params.k);
-    phi::Allocator* allocator = paddle::GetAllocator(place);
-    auto ws = allocator->Allocate(ws_size)->ptr();
-    cutlass_extensions::CutlassGemmConfig best_config;
-    if (getenv("FLAGS_fastdeploy_op_configs")) {
-        std::string fastdeploy_op_configs = getenv("FLAGS_fastdeploy_op_configs");
-        if (fastdeploy_op_configs == "tune") {
-            auto& ConfigManager = ConfigManager::get_instance();
-            auto configs = runner->getConfigs(params.k);
-            best_config = configs[0];
-            cudaEvent_t begin, end;
-            cudaEventCreate(&begin);
-            cudaEventCreate(&end);
-            float fast_time = 1e8;
-            printf("Begin tune cutlass kernel for [%d] confis...\n", static_cast<int>(configs.size()));
-            for (int c_idx = 0; c_idx < static_cast<int>(configs.size()); ++c_idx) {
-                auto config = configs[c_idx];
-                for (int i = 0; i < 5; ++i) {
-                    runner->gemm(params.act,
-                                params.weight,
-                                params.scales,
-                                params.zeros,
-                                params.bias,
-                                params.alpha,
-                                params.out,
-                                params.m,
-                                params.n,
-                                params.k,
-                                params.groupsize,
-                                config,
-                                ws,
-                                ws_size,
-                                stream);
-                }
-                cudaEventRecord(begin, stream);
-                for (int i = 0; i < 15; ++i) {
-                    runner->gemm(params.act,
-                                params.weight,
-                                params.scales,
-                                params.zeros,
-                                params.bias,
-                                params.alpha,
-                                params.out,
-                                params.m,
-                                params.n,
-                                params.k,
-                                params.groupsize,
-                                config,
-                                ws,
-                                ws_size,
-                                stream);
-                }
-                cudaEventRecord(end, stream);
-                cudaEventSynchronize(end);
-                float time;
-                cudaEventElapsedTime(&time, begin, end);
-                // std::cout << config.toString() << "time: " << time << std::endl;
-                if (time < fast_time) {
-                    fast_time = time;
-                    best_config = config;
-                }
-            }
-            cudaEventDestroy(begin);
-            cudaEventDestroy(end);
-            ConfigManager.update("scaled_gemm_f8_i4_f16", params.m, params.n, params.k, best_config.toString());
-        } else {
-            if (!std::filesystem::exists(fastdeploy_op_configs)) {
-                PADDLE_THROW(phi::errors::Fatal("Warning: The file \"" + fastdeploy_op_configs + "\" does not exist in the specified path." ));
-            } else {
-                auto& ConfigManager = ConfigManager::get_instance(fastdeploy_op_configs);
-                auto best_config_string = ConfigManager.get_best_config("scaled_gemm_f8_i4_f16", params.m, params.n, params.k);
-                if (!best_config_string.empty()) {
-                    best_config.fromString(best_config_string);
-                }
-            }
+  auto runner = std::make_shared<
+      kernels::cutlass_kernels::CutlassFpAIntBGemmRunner<__nv_fp8_e4m3,
+                                                         cutlass::uint4b_t,
+                                                         QuantOp,
+                                                         half,
+                                                         T,
+                                                         T>>();
+  int ws_size = runner->getWorkspaceSize(params.m, params.n, params.k);
+  phi::Allocator* allocator = paddle::GetAllocator(place);
+  auto ws = allocator->Allocate(ws_size)->ptr();
+  cutlass_extensions::CutlassGemmConfig best_config;
+  if (getenv("FLAGS_fastdeploy_op_configs")) {
+    std::string fastdeploy_op_configs = getenv("FLAGS_fastdeploy_op_configs");
+    if (fastdeploy_op_configs == "tune") {
+      auto& ConfigManager = ConfigManager::get_instance();
+      auto configs = runner->getConfigs(params.k);
+      best_config = configs[0];
+      cudaEvent_t begin, end;
+      cudaEventCreate(&begin);
+      cudaEventCreate(&end);
+      float fast_time = 1e8;
+      printf("Begin tune cutlass kernel for [%d] confis...\n",
+             static_cast<int>(configs.size()));
+      for (int c_idx = 0; c_idx < static_cast<int>(configs.size()); ++c_idx) {
+        auto config = configs[c_idx];
+        for (int i = 0; i < 5; ++i) {
+          runner->gemm(params.act,
+                       params.weight,
+                       params.scales,
+                       params.zeros,
+                       params.bias,
+                       params.alpha,
+                       params.out,
+                       params.m,
+                       params.n,
+                       params.k,
+                       params.groupsize,
+                       config,
+                       ws,
+                       ws_size,
+                       stream);
         }
+        cudaEventRecord(begin, stream);
+        for (int i = 0; i < 15; ++i) {
+          runner->gemm(params.act,
+                       params.weight,
+                       params.scales,
+                       params.zeros,
+                       params.bias,
+                       params.alpha,
+                       params.out,
+                       params.m,
+                       params.n,
+                       params.k,
+                       params.groupsize,
+                       config,
+                       ws,
+                       ws_size,
+                       stream);
+        }
+        cudaEventRecord(end, stream);
+        cudaEventSynchronize(end);
+        float time;
+        cudaEventElapsedTime(&time, begin, end);
+        // std::cout << config.toString() << "time: " << time << std::endl;
+        if (time < fast_time) {
+          fast_time = time;
+          best_config = config;
+        }
+      }
+      cudaEventDestroy(begin);
+      cudaEventDestroy(end);
+      ConfigManager.update("scaled_gemm_f8_i4_f16",
+                           params.m,
+                           params.n,
+                           params.k,
+                           best_config.toString());
+    } else {
+      if (!std::filesystem::exists(fastdeploy_op_configs)) {
+        PADDLE_THROW(
+            phi::errors::Fatal("Warning: The file \"" + fastdeploy_op_configs +
+                               "\" does not exist in the specified path."));
+      } else {
+        auto& ConfigManager =
+            ConfigManager::get_instance(fastdeploy_op_configs);
+        auto best_config_string = ConfigManager.get_best_config(
+            "scaled_gemm_f8_i4_f16", params.m, params.n, params.k);
+        if (!best_config_string.empty()) {
+          best_config.fromString(best_config_string);
+        }
+      }
     }
-    runner->gemm(params.act,
-                params.weight,
-                params.scales,
-                params.zeros,
-                params.bias,
-                params.alpha,
-                params.out,
-                params.m,
-                params.n,
-                params.k,
-                params.groupsize,
-                best_config,
-                ws,
-                ws_size,
-                stream);
+  }
+  runner->gemm(params.act,
+               params.weight,
+               params.scales,
+               params.zeros,
+               params.bias,
+               params.alpha,
+               params.out,
+               params.m,
+               params.n,
+               params.k,
+               params.groupsize,
+               best_config,
+               ws,
+               ws_size,
+               stream);
 }
 
 std::vector<paddle::Tensor> scaled_gemm_f8_i4_f16(
-                                const paddle::Tensor& x,
-                                const paddle::Tensor& y,
-                                const paddle::Tensor& scale,
-                                const paddle::optional<paddle::Tensor>& zero_points,
-                                const paddle::optional<paddle::Tensor>& bias,
-                                int groupsize,
-                                float out_scale,
-                                std::string out_dtype) {
-    // if (x.dims().size() != 2 ||
-    //     y.dims().size() != 2 ||
-    //     x.dims()[x.dims().size() - 1] != y.dims()[y.dims().size() - 1] ||
-    //     x.dtype() != paddle::DataType::FLOAT8_E4M3FN ||
-    //     y.dtype() != paddle::DataType::INT8) {
-    //     PADDLE_THROW(phi::errors::Fatal(
-    //         "Only support x[M, K](fp8_e4m3fn) and y[N, K](int8) as input"));
-    // }
+    const paddle::Tensor& x,
+    const paddle::Tensor& y,
+    const paddle::Tensor& scale,
+    const paddle::optional<paddle::Tensor>& zero_points,
+    const paddle::optional<paddle::Tensor>& bias,
+    int groupsize,
+    float out_scale,
+    std::string out_dtype) {
+  // if (x.dims().size() != 2 ||
+  //     y.dims().size() != 2 ||
+  //     x.dims()[x.dims().size() - 1] != y.dims()[y.dims().size() - 1] ||
+  //     x.dtype() != paddle::DataType::FLOAT8_E4M3FN ||
+  //     y.dtype() != paddle::DataType::INT8) {
+  //     PADDLE_THROW(phi::errors::Fatal(
+  //         "Only support x[M, K](fp8_e4m3fn) and y[N, K](int8) as input"));
+  // }
 
-    int M = x.dims()[0];
-    int K = x.dims()[1];
-    int N = y.dims()[0] * 2;
+  int M = x.dims()[0];
+  int K = x.dims()[1];
+  int N = y.dims()[0] * 2;
 
-    paddle::Tensor out;
+  paddle::Tensor out;
 
-    Pointer out_ptr = nullptr;
-    ConstPointer x_ptr = nullptr;
-    ConstPointer y_ptr = nullptr;
-    ConstPointer scale_ptr = nullptr;
-    ConstPointer zero_points_ptr = nullptr;
-    ConstPointer bias_ptr = nullptr;
+  Pointer out_ptr = nullptr;
+  ConstPointer x_ptr = nullptr;
+  ConstPointer y_ptr = nullptr;
+  ConstPointer scale_ptr = nullptr;
+  ConstPointer zero_points_ptr = nullptr;
+  ConstPointer bias_ptr = nullptr;
 
-    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
-    y_ptr = reinterpret_cast<const void*>(y.data<int8_t>());
-    std::vector<int64_t> out_shape = x.shape();
-    out_shape[0] = M;
-    out_shape[1] = N;
+  x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
+  y_ptr = reinterpret_cast<const void*>(y.data<int8_t>());
+  std::vector<int64_t> out_shape = x.shape();
+  out_shape[0] = M;
+  out_shape[1] = N;
 
-    auto place = x.place();
-    cudaStream_t stream = x.stream();
+  auto place = x.place();
+  cudaStream_t stream = x.stream();
 
-    if (out_dtype == "float16") {
-        scale_ptr = reinterpret_cast<const void*>(scale.data<phi::dtype::float16>());
-        if (zero_points) {
-            zero_points_ptr =
-                reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
-                    zero_points.get().data<phi::dtype::float16>()));
-        }
-        if (bias) {
-            bias_ptr = reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
-                bias.get().data<phi::dtype::float16>()));
-        }
-        out = paddle::empty(out_shape, paddle::DataType::FLOAT16, x.place());
-        out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float16>());
-        if (groupsize > 0) {
-            Params params(x_ptr, y_ptr, scale_ptr, zero_points_ptr, bias_ptr, out_ptr, out_scale, M, N, K, groupsize);
-            scaled_gemm_f8_i4_f16_launcher<half, cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS>(params, place, stream);
-        } else {
-            Params params(x_ptr, y_ptr, scale_ptr, nullptr, bias_ptr, out_ptr, out_scale, M, N, K, K);
-            scaled_gemm_f8_i4_f16_launcher<half, cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>(params, place, stream);
-        }
-    } else if (out_dtype == "bfloat16") {
-        scale_ptr = reinterpret_cast<const void*>(scale.data<phi::dtype::bfloat16>());
-        if (zero_points) {
-            zero_points_ptr =
-                reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
-                    zero_points.get().data<phi::dtype::bfloat16>()));
-        }
-        if (bias) {
-            bias_ptr = reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
-                bias.get().data<phi::dtype::bfloat16>()));
-        }
-        out = paddle::empty(out_shape, paddle::DataType::BFLOAT16, x.place());
-        out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::bfloat16>());
-        if (groupsize > 0) {
-            Params params(x_ptr, y_ptr, scale_ptr, zero_points_ptr, bias_ptr, out_ptr, out_scale, M, N, K, groupsize);
-            scaled_gemm_f8_i4_f16_launcher<__nv_bfloat16, cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS>(params, place, stream);
-        } else {
-            Params params(x_ptr, y_ptr, scale_ptr, nullptr, bias_ptr, out_ptr, out_scale, M, N, K, K);
-            scaled_gemm_f8_i4_f16_launcher<__nv_bfloat16, cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>(params, place, stream);
-        }
-    } else {
-        PADDLE_THROW(
-            phi::errors::Fatal("fp8_int4_gemm only support bfloat16 "
-                             "and float16 output"));
+  if (out_dtype == "float16") {
+    scale_ptr =
+        reinterpret_cast<const void*>(scale.data<phi::dtype::float16>());
+    if (zero_points) {
+      zero_points_ptr =
+          reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
+              zero_points.get().data<phi::dtype::float16>()));
     }
-    return {out};
+    if (bias) {
+      bias_ptr = reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
+          bias.get().data<phi::dtype::float16>()));
+    }
+    out = paddle::empty(out_shape, paddle::DataType::FLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float16>());
+    if (groupsize > 0) {
+      Params params(x_ptr,
+                    y_ptr,
+                    scale_ptr,
+                    zero_points_ptr,
+                    bias_ptr,
+                    out_ptr,
+                    out_scale,
+                    M,
+                    N,
+                    K,
+                    groupsize);
+      scaled_gemm_f8_i4_f16_launcher<
+          half,
+          cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS>(
+          params, place, stream);
+    } else {
+      Params params(x_ptr,
+                    y_ptr,
+                    scale_ptr,
+                    nullptr,
+                    bias_ptr,
+                    out_ptr,
+                    out_scale,
+                    M,
+                    N,
+                    K,
+                    K);
+      scaled_gemm_f8_i4_f16_launcher<
+          half,
+          cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>(
+          params, place, stream);
+    }
+  } else if (out_dtype == "bfloat16") {
+    scale_ptr =
+        reinterpret_cast<const void*>(scale.data<phi::dtype::bfloat16>());
+    if (zero_points) {
+      zero_points_ptr =
+          reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
+              zero_points.get().data<phi::dtype::bfloat16>()));
+    }
+    if (bias) {
+      bias_ptr = reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
+          bias.get().data<phi::dtype::bfloat16>()));
+    }
+    out = paddle::empty(out_shape, paddle::DataType::BFLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::bfloat16>());
+    if (groupsize > 0) {
+      Params params(x_ptr,
+                    y_ptr,
+                    scale_ptr,
+                    zero_points_ptr,
+                    bias_ptr,
+                    out_ptr,
+                    out_scale,
+                    M,
+                    N,
+                    K,
+                    groupsize);
+      scaled_gemm_f8_i4_f16_launcher<
+          __nv_bfloat16,
+          cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS>(
+          params, place, stream);
+    } else {
+      Params params(x_ptr,
+                    y_ptr,
+                    scale_ptr,
+                    nullptr,
+                    bias_ptr,
+                    out_ptr,
+                    out_scale,
+                    M,
+                    N,
+                    K,
+                    K);
+      scaled_gemm_f8_i4_f16_launcher<
+          __nv_bfloat16,
+          cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>(
+          params, place, stream);
+    }
+  } else {
+    PADDLE_THROW(
+        phi::errors::Fatal("fp8_int4_gemm only support bfloat16 "
+                           "and float16 output"));
+  }
+  return {out};
 }
 
 std::vector<std::vector<int64_t>> CutlassFp8Int4GemmInferShape(
@@ -253,21 +328,21 @@ std::vector<std::vector<int64_t>> CutlassFp8Int4GemmInferShape(
     int groupsize,
     float out_scale,
     std::string out_dtype) {
-    // if (x_shape.size() != 2 ||
-    //     y_shape.size() != 2 ||
-    //     x_shape[x_shape.size()-1] != y_shape[y_shape.size()-1]) {
-    //     PADDLE_THROW(phi::errors::Fatal(
-    //         "Only support x[M, K] and y[N, K] as input"));
-    // }
+  // if (x_shape.size() != 2 ||
+  //     y_shape.size() != 2 ||
+  //     x_shape[x_shape.size()-1] != y_shape[y_shape.size()-1]) {
+  //     PADDLE_THROW(phi::errors::Fatal(
+  //         "Only support x[M, K] and y[N, K] as input"));
+  // }
 
-    int M = x_shape[0];
-    int K = x_shape[1];
-    int N = y_shape[0] * 2;
+  int M = x_shape[0];
+  int K = x_shape[1];
+  int N = y_shape[0] * 2;
 
-    std::vector<int64_t> out_shape = x_shape;
-    out_shape[0] = M;
-    out_shape[1] = N;
-    return {out_shape};
+  std::vector<int64_t> out_shape = x_shape;
+  out_shape[0] = M;
+  out_shape[1] = N;
+  return {out_shape};
 }
 
 std::vector<paddle::DataType> CutlassFp8Int4GemmInferDtype(
@@ -279,23 +354,25 @@ std::vector<paddle::DataType> CutlassFp8Int4GemmInferDtype(
     int groupsize,
     float out_scale,
     std::string out_dtype) {
-    paddle::DataType data_type;
-    if (out_dtype == "bfloat16") {
-        data_type = paddle::DataType::BFLOAT16;
-    } else if (out_dtype == "float16") {
-        data_type = paddle::DataType::FLOAT16;
-    } else {
-        PD_THROW(
-            "fp8_int4_half_gemm_fused only support bfloat16 and float16 output");
-    }
-    return {data_type};
+  paddle::DataType data_type;
+  if (out_dtype == "bfloat16") {
+    data_type = paddle::DataType::BFLOAT16;
+  } else if (out_dtype == "float16") {
+    data_type = paddle::DataType::FLOAT16;
+  } else {
+    PD_THROW(
+        "fp8_int4_half_gemm_fused only support bfloat16 and float16 output");
+  }
+  return {data_type};
 }
 
 PD_BUILD_STATIC_OP(scaled_gemm_f8_i4_f16)
-    .Inputs({"x", "y", "scale", paddle::Optional("zero_points"), paddle::Optional("bias")})
-    .Attrs({"groupsize: int",
-            "out_scale: float",
-            "out_dtype: std::string"})
+    .Inputs({"x",
+             "y",
+             "scale",
+             paddle::Optional("zero_points"),
+             paddle::Optional("bias")})
+    .Attrs({"groupsize: int", "out_scale: float", "out_dtype: std::string"})
     .Outputs({"out"})
     .SetKernelFn(PD_KERNEL(scaled_gemm_f8_i4_f16))
     .SetInferShapeFn(PD_INFER_SHAPE(CutlassFp8Int4GemmInferShape))
diff --git a/custom_ops/gpu_ops/scaled_gemm_f8_i4_f16_weight_quantize.cu b/custom_ops/gpu_ops/scaled_gemm_f8_i4_f16_weight_quantize.cu
index 978106e746..ffbe0a63fd 100644
--- a/custom_ops/gpu_ops/scaled_gemm_f8_i4_f16_weight_quantize.cu
+++ b/custom_ops/gpu_ops/scaled_gemm_f8_i4_f16_weight_quantize.cu
@@ -46,29 +46,28 @@ void per_channel_quant(int8_t* output,
     int8_t* current_quantized_weight_row = output + ii * bytes_per_out_col;
     const T* current_weight_row = input + ii * num_cols;
     for (size_t jj = 0; jj < bytes_per_out_col; ++jj) {
-        // We will pack two int4 elements per iteration of the inner loop.
-        int8_t packed_int4s = 0;
-        for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
-            const size_t input_idx = 2 * jj + packed_idx;
-            if (input_idx < num_cols) {
-                const float col_scale = static_cast<float>(scale[input_idx]);
-                const float weight_elt =
-                    static_cast<float>(current_weight_row[input_idx]);
-                const float scaled_weight = round(weight_elt / col_scale);
-                int int_weight = static_cast<int>(scaled_weight);
-                const int8_t clipped_weight = std::max(-7, std::min(7, int_weight));
-                // Kill the sign extension bits (hence 0x0F mask) then shift to
-                // upper bits if packing the second int4 and or the bits into the
-                // final result.
-                packed_int4s |= ((clipped_weight & 0x0F) << (4 * packed_idx));
-            }
+      // We will pack two int4 elements per iteration of the inner loop.
+      int8_t packed_int4s = 0;
+      for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
+        const size_t input_idx = 2 * jj + packed_idx;
+        if (input_idx < num_cols) {
+          const float col_scale = static_cast<float>(scale[input_idx]);
+          const float weight_elt =
+              static_cast<float>(current_weight_row[input_idx]);
+          const float scaled_weight = round(weight_elt / col_scale);
+          int int_weight = static_cast<int>(scaled_weight);
+          const int8_t clipped_weight = std::max(-7, std::min(7, int_weight));
+          // Kill the sign extension bits (hence 0x0F mask) then shift to
+          // upper bits if packing the second int4 and or the bits into the
+          // final result.
+          packed_int4s |= ((clipped_weight & 0x0F) << (4 * packed_idx));
         }
-        current_quantized_weight_row[jj] = packed_int4s;
+      }
+      current_quantized_weight_row[jj] = packed_int4s;
     }
   }
 }
 
-
 template <typename T, typename ScaleT>
 void group_wise_quant(int8_t* output,
                       const T* input,
@@ -81,25 +80,25 @@ void group_wise_quant(int8_t* output,
     int8_t* current_quantized_weight_row = output + ii * bytes_per_out_col;
     const T* current_weight_row = input + ii * num_cols;
     for (size_t jj = 0; jj < bytes_per_out_col; ++jj) {
-        // We will pack two int4 elements per iteration of the inner loop.
-        int8_t packed_int4s = 0;
-        for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
-          const size_t input_idx = 2 * jj + packed_idx;
-          if (input_idx < num_cols) {
-            size_t scale_cur_offset = input_idx + (ii / group_size) * num_cols;
-            const float col_scale = static_cast<float>(scale[scale_cur_offset]);
-            const float weight_elt =
-                static_cast<float>(current_weight_row[input_idx]);
-            const float scaled_weight = round(weight_elt / col_scale);
-            int int_weight = static_cast<int>(scaled_weight);
-            const int8_t clipped_weight = std::max(-7, std::min(7, int_weight));
-            // Kill the sign extension bits (hence 0x0F mask) then shift to
-            // upper bits if packing the second int4 and or the bits into the
-            // final result.
-            packed_int4s |= ((clipped_weight & 0x0F) << (4 * packed_idx));
-          }
+      // We will pack two int4 elements per iteration of the inner loop.
+      int8_t packed_int4s = 0;
+      for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
+        const size_t input_idx = 2 * jj + packed_idx;
+        if (input_idx < num_cols) {
+          size_t scale_cur_offset = input_idx + (ii / group_size) * num_cols;
+          const float col_scale = static_cast<float>(scale[scale_cur_offset]);
+          const float weight_elt =
+              static_cast<float>(current_weight_row[input_idx]);
+          const float scaled_weight = round(weight_elt / col_scale);
+          int int_weight = static_cast<int>(scaled_weight);
+          const int8_t clipped_weight = std::max(-7, std::min(7, int_weight));
+          // Kill the sign extension bits (hence 0x0F mask) then shift to
+          // upper bits if packing the second int4 and or the bits into the
+          // final result.
+          packed_int4s |= ((clipped_weight & 0x0F) << (4 * packed_idx));
         }
-        current_quantized_weight_row[jj] = packed_int4s;
+      }
+      current_quantized_weight_row[jj] = packed_int4s;
     }
   }
 }
@@ -125,101 +124,127 @@ void group_wise_scale(ScaleT* scale,
   }
 }
 
-std::vector<paddle::Tensor> Fp8Int4WeightQuantizeKernel(const paddle::Tensor &input,
-    int groupsize,
-    std::string scale_dtype) {
-    auto input_cpu = input.copy_to(paddle::CPUPlace(), false);
-    auto shape = input.shape();
-    auto k = static_cast<size_t>(shape[0]);
-    auto n = static_cast<size_t>(shape[1]);
+std::vector<paddle::Tensor> Fp8Int4WeightQuantizeKernel(
+    const paddle::Tensor& input, int groupsize, std::string scale_dtype) {
+  auto input_cpu = input.copy_to(paddle::CPUPlace(), false);
+  auto shape = input.shape();
+  auto k = static_cast<size_t>(shape[0]);
+  auto n = static_cast<size_t>(shape[1]);
 
-    paddle::Tensor scale;
-    auto packed_int4 = paddle::full({shape[0] * shape[1] / 2}, 0, paddle::DataType::INT8, paddle::CPUPlace());
-    if (scale_dtype == "bfloat16") {
-        if (groupsize > 0) {
-            scale = paddle::full({shape[0] / groupsize * shape[1]}, 1.0, paddle::DataType::BFLOAT16, paddle::CPUPlace());
-            group_wise_scale(scale.data<phi::dtype::bfloat16>(), input_cpu.data<float>(), k, n, 7.0f, groupsize);
-            group_wise_quant(packed_int4.data<int8_t>(),
-                            input_cpu.data<float>(),
-                            scale.data<phi::dtype::bfloat16>(),
-                            k,
-                            n,
-                            groupsize);
-        } else {
-            scale = paddle::full({shape[1]}, 1.0, paddle::DataType::BFLOAT16, paddle::CPUPlace());
-            per_channel_scale(scale.data<phi::dtype::bfloat16>(), input_cpu.data<float>(), k, n, 7.0f);
-            per_channel_quant(packed_int4.data<int8_t>(),
-                              input_cpu.data<float>(),
-                              scale.data<phi::dtype::bfloat16>(),
-                              k,
-                              n);
-        }
-    } else if (scale_dtype == "float16") {
-        if (groupsize > 0) {
-            scale = paddle::full({shape[0] / groupsize * shape[1]}, 1.0, paddle::DataType::FLOAT16, paddle::CPUPlace());
-            group_wise_scale(scale.data<phi::dtype::float16>(), input_cpu.data<float>(), k, n, 7.0f, groupsize);
-            group_wise_quant(packed_int4.data<int8_t>(),
-                input_cpu.data<float>(),
-                scale.data<phi::dtype::float16>(),
-                k,
-                n,
-                groupsize);
-        } else {
-            scale = paddle::full({shape[1]}, 1.0, paddle::DataType::FLOAT16, paddle::CPUPlace());
-            per_channel_scale(scale.data<phi::dtype::float16>(), input_cpu.data<float>(), k, n, 7.0f);
-            per_channel_quant(packed_int4.data<int8_t>(),
-                    input_cpu.data<float>(),
-                    scale.data<phi::dtype::float16>(),
-                    k,
-                    n);
-        }
+  paddle::Tensor scale;
+  auto packed_int4 = paddle::full(
+      {shape[0] * shape[1] / 2}, 0, paddle::DataType::INT8, paddle::CPUPlace());
+  if (scale_dtype == "bfloat16") {
+    if (groupsize > 0) {
+      scale = paddle::full({shape[0] / groupsize * shape[1]},
+                           1.0,
+                           paddle::DataType::BFLOAT16,
+                           paddle::CPUPlace());
+      group_wise_scale(scale.data<phi::dtype::bfloat16>(),
+                       input_cpu.data<float>(),
+                       k,
+                       n,
+                       7.0f,
+                       groupsize);
+      group_wise_quant(packed_int4.data<int8_t>(),
+                       input_cpu.data<float>(),
+                       scale.data<phi::dtype::bfloat16>(),
+                       k,
+                       n,
+                       groupsize);
+    } else {
+      scale = paddle::full(
+          {shape[1]}, 1.0, paddle::DataType::BFLOAT16, paddle::CPUPlace());
+      per_channel_scale(scale.data<phi::dtype::bfloat16>(),
+                        input_cpu.data<float>(),
+                        k,
+                        n,
+                        7.0f);
+      per_channel_quant(packed_int4.data<int8_t>(),
+                        input_cpu.data<float>(),
+                        scale.data<phi::dtype::bfloat16>(),
+                        k,
+                        n);
     }
+  } else if (scale_dtype == "float16") {
+    if (groupsize > 0) {
+      scale = paddle::full({shape[0] / groupsize * shape[1]},
+                           1.0,
+                           paddle::DataType::FLOAT16,
+                           paddle::CPUPlace());
+      group_wise_scale(scale.data<phi::dtype::float16>(),
+                       input_cpu.data<float>(),
+                       k,
+                       n,
+                       7.0f,
+                       groupsize);
+      group_wise_quant(packed_int4.data<int8_t>(),
+                       input_cpu.data<float>(),
+                       scale.data<phi::dtype::float16>(),
+                       k,
+                       n,
+                       groupsize);
+    } else {
+      scale = paddle::full(
+          {shape[1]}, 1.0, paddle::DataType::FLOAT16, paddle::CPUPlace());
+      per_channel_scale(scale.data<phi::dtype::float16>(),
+                        input_cpu.data<float>(),
+                        k,
+                        n,
+                        7.0f);
+      per_channel_quant(packed_int4.data<int8_t>(),
+                        input_cpu.data<float>(),
+                        scale.data<phi::dtype::float16>(),
+                        k,
+                        n);
+    }
+  }
 
-    auto out = paddle::full({shape[1] / 2, shape[0]}, 0, paddle::DataType::INT8, paddle::CPUPlace());
-    preprocess_weights_for_mixed_gemm(
-        out.data<int8_t>(),
-        packed_int4.data<int8_t>(),
-        {k, n},
-        kernels::cutlass_kernels::QuantType::W4_AFP8,
-        false);
-    return {out, scale};
+  auto out = paddle::full(
+      {shape[1] / 2, shape[0]}, 0, paddle::DataType::INT8, paddle::CPUPlace());
+  preprocess_weights_for_mixed_gemm(
+      out.data<int8_t>(),
+      packed_int4.data<int8_t>(),
+      {k, n},
+      kernels::cutlass_kernels::QuantType::W4_AFP8,
+      false);
+  return {out, scale};
 }
 
 std::vector<std::vector<int64_t>> Fp8Int4WeightQuantizeInferShape(
     const std::vector<int64_t>& input_shape,
     int groupsize,
     std::string scale_dtype) {
-    std::vector<int64_t> out_shape = {input_shape[1] / 2, input_shape[0]};
-    std::vector<int64_t> scale_shape;
-    if (groupsize > 0) {
-        scale_shape = {input_shape[0] / groupsize, input_shape[1]};
-    } else {
-        scale_shape = {input_shape[1]};
-    }
-    return {out_shape, scale_shape};
+  std::vector<int64_t> out_shape = {input_shape[1] / 2, input_shape[0]};
+  std::vector<int64_t> scale_shape;
+  if (groupsize > 0) {
+    scale_shape = {input_shape[0] / groupsize, input_shape[1]};
+  } else {
+    scale_shape = {input_shape[1]};
+  }
+  return {out_shape, scale_shape};
 }
 
 std::vector<paddle::DataType> Fp8Int4WeightQuantizeInferDtype(
     const paddle::DataType& input_type,
     int groupsize,
     std::string scale_dtype) {
-    paddle::DataType scale_data_type;
-    if (scale_dtype == "bfloat16") {
-        scale_data_type = paddle::DataType::BFLOAT16;
-    } else if (scale_dtype == "float16") {
-        scale_data_type = paddle::DataType::FLOAT16;
-    } else {
-        PD_THROW(
-            "scaled_gemm_f8_i4_f16_half_gemm_fused only support bfloat16 and float16 output");
-    }
-    return {paddle::DataType::INT8, scale_data_type};
+  paddle::DataType scale_data_type;
+  if (scale_dtype == "bfloat16") {
+    scale_data_type = paddle::DataType::BFLOAT16;
+  } else if (scale_dtype == "float16") {
+    scale_data_type = paddle::DataType::FLOAT16;
+  } else {
+    PD_THROW(
+        "scaled_gemm_f8_i4_f16_half_gemm_fused only support bfloat16 and "
+        "float16 output");
+  }
+  return {paddle::DataType::INT8, scale_data_type};
 }
 
-
 PD_BUILD_STATIC_OP(scaled_gemm_f8_i4_f16_weight_quantize)
     .Inputs({"input"})
-    .Attrs({"groupsize: int",
-            "scale_dtype: std::string"})
+    .Attrs({"groupsize: int", "scale_dtype: std::string"})
     .Outputs({"output", "scale"})
     .SetKernelFn(PD_KERNEL(Fp8Int4WeightQuantizeKernel))
     .SetInferShapeFn(PD_INFER_SHAPE(Fp8Int4WeightQuantizeInferShape))
diff --git a/custom_ops/gpu_ops/seqs2seqs.cu b/custom_ops/gpu_ops/seqs2seqs.cu
index fcd40fbb09..41e3716433 100644
--- a/custom_ops/gpu_ops/seqs2seqs.cu
+++ b/custom_ops/gpu_ops/seqs2seqs.cu
@@ -35,58 +35,56 @@ __global__ void GetSeqMapping(int* merged_seq_lens_this_time,
                               const int* cum_group_lens,
                               const int* system_lens,
                               const int max_bsz) {
-    const int bi = blockIdx.x;
-    const int ti = threadIdx.x;
-    const int group_len = group_lens[bi];
-    if (group_len <= 0) return;
-    const int start_bi = bi == 0 ? 0 : cum_group_lens[bi - 1];
-    if (ti == 0) {
-        const int* group_ids_now = group_ids + bi * max_bsz;
-        int seq_len_sum = 0;
-        int seq_len_encoder_sum = 0;
-        int dec_count = 0;
-        int system_len = 0;
-        for (int i = 0; i < group_len; i++) {
-            const int group_id = group_ids_now[i];
-            const int seq_len_this_time = seq_lens_this_time[group_id];
-            const int seq_len_encoder = seq_lens_encoder[group_id];
-            if (seq_len_encoder <= 0) {  // decoder
-                seq_len_sum += seq_len_this_time;
-                seq_len_encoder_sum += seq_len_encoder;
-                new_seq_lens_this_time[start_bi + dec_count] =
-                    seq_len_this_time;
-                new_seq_lens_encoder[start_bi + dec_count] = seq_len_encoder;
-                new_seq_lens_decoder[start_bi + dec_count] =
-                    seq_lens_decoder[group_id];
-                seq_mapping[start_bi + dec_count] = group_id;
-                system_len = system_lens[group_id];
-                new_system_lens[group_id] = system_len;
-                dec_count++;
-            } else {  // encoder
-                int encoder_bid = atomicAdd(dec_group_num, 1);
-                new_seq_lens_this_time[encoder_bid] = seq_len_this_time;
-                new_seq_lens_encoder[encoder_bid] = seq_len_encoder;
-                new_seq_lens_decoder[encoder_bid] = seq_lens_decoder[group_id];
-                seq_mapping[encoder_bid] = group_id;
-                new_system_lens[group_id] = 0;
-            }
-        }
-        if (dec_count >= 1) {
-            merged_seq_lens_this_time[bi] = seq_len_sum;
-            merged_seq_lens_encoder[bi] = seq_len_encoder_sum;
-            if (dec_count == 1) {
-                merged_seq_lens_decoder[bi] = 0;
-            } else {
-                merged_seq_lens_decoder[bi] = system_len;
-            }
-        }
-        if (dec_count <= 1) {
-            for (int i = 0; i < group_len; i++) {
-                const int group_id = group_ids_now[i];
-                new_system_lens[group_id] = 0;
-            }
-        }
+  const int bi = blockIdx.x;
+  const int ti = threadIdx.x;
+  const int group_len = group_lens[bi];
+  if (group_len <= 0) return;
+  const int start_bi = bi == 0 ? 0 : cum_group_lens[bi - 1];
+  if (ti == 0) {
+    const int* group_ids_now = group_ids + bi * max_bsz;
+    int seq_len_sum = 0;
+    int seq_len_encoder_sum = 0;
+    int dec_count = 0;
+    int system_len = 0;
+    for (int i = 0; i < group_len; i++) {
+      const int group_id = group_ids_now[i];
+      const int seq_len_this_time = seq_lens_this_time[group_id];
+      const int seq_len_encoder = seq_lens_encoder[group_id];
+      if (seq_len_encoder <= 0) {  // decoder
+        seq_len_sum += seq_len_this_time;
+        seq_len_encoder_sum += seq_len_encoder;
+        new_seq_lens_this_time[start_bi + dec_count] = seq_len_this_time;
+        new_seq_lens_encoder[start_bi + dec_count] = seq_len_encoder;
+        new_seq_lens_decoder[start_bi + dec_count] = seq_lens_decoder[group_id];
+        seq_mapping[start_bi + dec_count] = group_id;
+        system_len = system_lens[group_id];
+        new_system_lens[group_id] = system_len;
+        dec_count++;
+      } else {  // encoder
+        int encoder_bid = atomicAdd(dec_group_num, 1);
+        new_seq_lens_this_time[encoder_bid] = seq_len_this_time;
+        new_seq_lens_encoder[encoder_bid] = seq_len_encoder;
+        new_seq_lens_decoder[encoder_bid] = seq_lens_decoder[group_id];
+        seq_mapping[encoder_bid] = group_id;
+        new_system_lens[group_id] = 0;
+      }
     }
+    if (dec_count >= 1) {
+      merged_seq_lens_this_time[bi] = seq_len_sum;
+      merged_seq_lens_encoder[bi] = seq_len_encoder_sum;
+      if (dec_count == 1) {
+        merged_seq_lens_decoder[bi] = 0;
+      } else {
+        merged_seq_lens_decoder[bi] = system_len;
+      }
+    }
+    if (dec_count <= 1) {
+      for (int i = 0; i < group_len; i++) {
+        const int group_id = group_ids_now[i];
+        new_system_lens[group_id] = 0;
+      }
+    }
+  }
 }
 
 std::vector<paddle::Tensor> Seqs2Seqs(const paddle::Tensor& seq_lens_this_time,
@@ -97,55 +95,55 @@ std::vector<paddle::Tensor> Seqs2Seqs(const paddle::Tensor& seq_lens_this_time,
                                       const paddle::Tensor& dec_group_num,
                                       const paddle::Tensor& cum_group_lens,
                                       const paddle::Tensor& system_lens) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int max_bsz = seq_lens_encoder.shape()[0];
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int max_bsz = seq_lens_encoder.shape()[0];
 
-    auto merged_seq_lens_this_time = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto merged_seq_lens_encoder = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto merged_seq_lens_decoder = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto new_system_lens = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto new_seq_lens_this_time = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto new_seq_lens_encoder = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto new_seq_lens_decoder = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto seq_mapping = paddle::full(
-        {bsz, 1}, -1, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto merged_seq_lens_this_time = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto merged_seq_lens_encoder = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto merged_seq_lens_decoder = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto new_system_lens = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto new_seq_lens_this_time = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto new_seq_lens_encoder = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto new_seq_lens_decoder = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto seq_mapping = paddle::full(
+      {bsz, 1}, -1, paddle::DataType::INT32, seq_lens_this_time.place());
 
-    constexpr int BlockSize = 32;
-    // const int blockSize = (max_bsz + 32 - 1) / 32 * 32;
-    GetSeqMapping<<<bsz, BlockSize, 0, cu_stream>>>(
-        merged_seq_lens_this_time.data<int>(),
-        merged_seq_lens_encoder.data<int>(),
-        merged_seq_lens_decoder.data<int>(),
-        new_system_lens.data<int>(),
-        new_seq_lens_this_time.data<int>(),
-        new_seq_lens_encoder.data<int>(),
-        new_seq_lens_decoder.data<int>(),
-        seq_mapping.data<int>(),
-        const_cast<int*>(dec_group_num.data<int>()),
-        seq_lens_this_time.data<int>(),
-        seq_lens_encoder.data<int>(),
-        seq_lens_decoder.data<int>(),
-        group_ids.data<int>(),
-        group_lens.data<int>(),
-        cum_group_lens.data<int>(),
-        system_lens.data<int>(),
-        max_bsz);
-    return {merged_seq_lens_this_time,
-            merged_seq_lens_encoder,
-            merged_seq_lens_decoder,
-            new_system_lens,
-            new_seq_lens_this_time,
-            new_seq_lens_encoder,
-            new_seq_lens_decoder,
-            seq_mapping};  // , enc_token_num, dec_token_num};
+  constexpr int BlockSize = 32;
+  // const int blockSize = (max_bsz + 32 - 1) / 32 * 32;
+  GetSeqMapping<<<bsz, BlockSize, 0, cu_stream>>>(
+      merged_seq_lens_this_time.data<int>(),
+      merged_seq_lens_encoder.data<int>(),
+      merged_seq_lens_decoder.data<int>(),
+      new_system_lens.data<int>(),
+      new_seq_lens_this_time.data<int>(),
+      new_seq_lens_encoder.data<int>(),
+      new_seq_lens_decoder.data<int>(),
+      seq_mapping.data<int>(),
+      const_cast<int*>(dec_group_num.data<int>()),
+      seq_lens_this_time.data<int>(),
+      seq_lens_encoder.data<int>(),
+      seq_lens_decoder.data<int>(),
+      group_ids.data<int>(),
+      group_lens.data<int>(),
+      cum_group_lens.data<int>(),
+      system_lens.data<int>(),
+      max_bsz);
+  return {merged_seq_lens_this_time,
+          merged_seq_lens_encoder,
+          merged_seq_lens_decoder,
+          new_system_lens,
+          new_seq_lens_this_time,
+          new_seq_lens_encoder,
+          new_seq_lens_decoder,
+          seq_mapping};  // , enc_token_num, dec_token_num};
 }
 
 std::vector<std::vector<int64_t>> Seqs2SeqsInferShape(
@@ -157,15 +155,15 @@ std::vector<std::vector<int64_t>> Seqs2SeqsInferShape(
     const std::vector<int64_t>& dec_group_num_shape,
     const std::vector<int64_t>& cum_group_lens_shape,
     const std::vector<int64_t>& system_lens_shape) {
-    int64_t bsz = seq_lens_this_time_shape[0];
-    return {{bsz, 1},
-            {bsz, 1},
-            {bsz, 1},
-            {bsz, 1},
-            {bsz, 1},
-            {bsz, 1},
-            {bsz, 1},
-            {bsz, 1}};
+  int64_t bsz = seq_lens_this_time_shape[0];
+  return {{bsz, 1},
+          {bsz, 1},
+          {bsz, 1},
+          {bsz, 1},
+          {bsz, 1},
+          {bsz, 1},
+          {bsz, 1},
+          {bsz, 1}};
 }
 
 std::vector<paddle::DataType> Seqs2SeqsInferDtype(
@@ -177,14 +175,14 @@ std::vector<paddle::DataType> Seqs2SeqsInferDtype(
     const paddle::DataType& dec_group_num_dtype,
     const paddle::DataType& cum_group_lens_dtype,
     const paddle::DataType& system_lens_dtype) {
-    return {seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype};
+  return {seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype};
 }
 
 PD_BUILD_STATIC_OP(seqs2seqs)
diff --git a/custom_ops/gpu_ops/set_flags.cu b/custom_ops/gpu_ops/set_flags.cu
index 7fcbc3925d..b120e1ab8b 100644
--- a/custom_ops/gpu_ops/set_flags.cu
+++ b/custom_ops/gpu_ops/set_flags.cu
@@ -19,13 +19,13 @@
 #endif
 
 __device__ bool is_in_list(const int64_t id, const int64_t *ids, int bs_id) {
-    bool is_in_list = false;
-    for (int i = 0; i < bs_id; i++) {
-        if (ids[i] == id) {
-            return true;
-        }
+  bool is_in_list = false;
+  for (int i = 0; i < bs_id; i++) {
+    if (ids[i] == id) {
+      return true;
     }
-    return is_in_list;
+  }
+  return is_in_list;
 }
 
 __global__ void set_value_by_id(const bool *stop_flags,
@@ -33,46 +33,45 @@ __global__ void set_value_by_id(const bool *stop_flags,
                                 bool *stop_flags_out,
                                 int bs,
                                 int bs_id) {
-    int tid = threadIdx.x;
-    if (tid < bs && !is_in_list(tid, ids, bs_id)) {
-        stop_flags_out[tid] = true;
-    }
+  int tid = threadIdx.x;
+  if (tid < bs && !is_in_list(tid, ids, bs_id)) {
+    stop_flags_out[tid] = true;
+  }
 }
 
 std::vector<paddle::Tensor> SetFlags(const paddle::Tensor &stop_flags,
                                      const paddle::Tensor &gather_id) {
-    PD_CHECK(gather_id.dtype() == paddle::DataType::INT64);
-    PD_CHECK(stop_flags.dtype() == paddle::DataType::BOOL);
-    auto cu_stream = stop_flags.stream();
-    std::vector<int64_t> flag_shape = stop_flags.shape();
-    std::vector<int64_t> id_shape = gather_id.shape();
-    auto stop_flags_out =
-        stop_flags.copy_to(stop_flags.place(), false);  // gpu -> gpu
-    if (flag_shape[0] == id_shape[0]) {
-        return {stop_flags_out};
-    }
-    int flag_bs = flag_shape[0];
-    int id_bs = id_shape[0];
-    int block_size = (flag_bs + 32 - 1) / 32 * 32;
-    set_value_by_id<<<1, block_size, 0, cu_stream>>>(
-        stop_flags.data<bool>(),
-        gather_id.data<int64_t>(),
-        stop_flags_out.data<bool>(),
-        flag_bs,
-        id_bs);
+  PD_CHECK(gather_id.dtype() == paddle::DataType::INT64);
+  PD_CHECK(stop_flags.dtype() == paddle::DataType::BOOL);
+  auto cu_stream = stop_flags.stream();
+  std::vector<int64_t> flag_shape = stop_flags.shape();
+  std::vector<int64_t> id_shape = gather_id.shape();
+  auto stop_flags_out =
+      stop_flags.copy_to(stop_flags.place(), false);  // gpu -> gpu
+  if (flag_shape[0] == id_shape[0]) {
     return {stop_flags_out};
+  }
+  int flag_bs = flag_shape[0];
+  int id_bs = id_shape[0];
+  int block_size = (flag_bs + 32 - 1) / 32 * 32;
+  set_value_by_id<<<1, block_size, 0, cu_stream>>>(stop_flags.data<bool>(),
+                                                   gather_id.data<int64_t>(),
+                                                   stop_flags_out.data<bool>(),
+                                                   flag_bs,
+                                                   id_bs);
+  return {stop_flags_out};
 }
 
 std::vector<std::vector<int64_t>> SetFlagsInferShape(
     const std::vector<int64_t> &stop_flags_shape,
     const std::vector<int64_t> &gather_id_shape) {
-    return {stop_flags_shape};
+  return {stop_flags_shape};
 }
 
 std::vector<paddle::DataType> SetFlagsInferDtype(
     const paddle::DataType &stop_flags_dtype,
     const paddle::DataType &gather_id_dtype) {
-    return {stop_flags_dtype};
+  return {stop_flags_dtype};
 }
 
 PD_BUILD_STATIC_OP(set_flags)
diff --git a/custom_ops/gpu_ops/set_mask_value.cu b/custom_ops/gpu_ops/set_mask_value.cu
index 3d72ae1b5c..336facf271 100644
--- a/custom_ops/gpu_ops/set_mask_value.cu
+++ b/custom_ops/gpu_ops/set_mask_value.cu
@@ -23,100 +23,110 @@ class PDTraits;
 
 template <>
 class PDTraits<paddle::DataType::FLOAT32> {
-public:
+ public:
   typedef float DataType;
   typedef float data_t;
 };
 
 template <>
 class PDTraits<paddle::DataType::FLOAT16> {
-public:
+ public:
   typedef half DataType;
   typedef paddle::float16 data_t;
 };
 
 template <>
 class PDTraits<paddle::DataType::BFLOAT16> {
-public:
+ public:
   typedef __nv_bfloat16 DataType;
   typedef paddle::bfloat16 data_t;
 };
 
 template <typename T>
-__global__ void set_value_by_id(const int *seq_lens, const bool *stop_flags, T *output_data, int *sequence_lengths, int bs, int length) {
-    int tid = threadIdx.x;
-    if (tid < bs) {
-        T *output_data_now = output_data + tid * length;
-        output_data_now[seq_lens[tid]] = 1.0;
-        if (stop_flags[tid]) {
-            sequence_lengths[tid] = 0;
-        }
+__global__ void set_value_by_id(const int* seq_lens,
+                                const bool* stop_flags,
+                                T* output_data,
+                                int* sequence_lengths,
+                                int bs,
+                                int length) {
+  int tid = threadIdx.x;
+  if (tid < bs) {
+    T* output_data_now = output_data + tid * length;
+    output_data_now[seq_lens[tid]] = 1.0;
+    if (stop_flags[tid]) {
+      sequence_lengths[tid] = 0;
     }
+  }
 }
 
 template <paddle::DataType D>
-std::vector<paddle::Tensor> set_mask_value(const paddle::Tensor& input_data, const paddle::Tensor& stop_flags, const paddle::Tensor& seq_lens) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+std::vector<paddle::Tensor> set_mask_value(const paddle::Tensor& input_data,
+                                           const paddle::Tensor& stop_flags,
+                                           const paddle::Tensor& seq_lens) {
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    PD_CHECK(seq_lens.dtype() == paddle::DataType::INT32);
-    PD_CHECK(stop_flags.dtype() == paddle::DataType::BOOL);
-    auto cu_stream = input_data.stream();
-    std::vector<int64_t> input_data_shape = input_data.shape();
-    std::vector<int64_t> seq_lens_shape = seq_lens.shape();
-    auto sequence_lengths = seq_lens.copy_to(seq_lens.place(), false);
+  PD_CHECK(seq_lens.dtype() == paddle::DataType::INT32);
+  PD_CHECK(stop_flags.dtype() == paddle::DataType::BOOL);
+  auto cu_stream = input_data.stream();
+  std::vector<int64_t> input_data_shape = input_data.shape();
+  std::vector<int64_t> seq_lens_shape = seq_lens.shape();
+  auto sequence_lengths = seq_lens.copy_to(seq_lens.place(), false);
 
-    int input_bs = input_data_shape[0];
-    int length = input_data_shape[3];
-    int seq_bs = seq_lens_shape[0];
+  int input_bs = input_data_shape[0];
+  int length = input_data_shape[3];
+  int seq_bs = seq_lens_shape[0];
 
-    int block_size = (input_bs + 32 - 1) / 32 * 32;
-    set_value_by_id<<<1, block_size, 0, cu_stream>>>(seq_lens.data<int>(),
-                                                     stop_flags.data<bool>(),
-                                                     reinterpret_cast<DataType_*>(const_cast<data_t*>(input_data.data<data_t>())),
-                                                     sequence_lengths.data<int>(), seq_bs, length);
-    return {sequence_lengths};
+  int block_size = (input_bs + 32 - 1) / 32 * 32;
+  set_value_by_id<<<1, block_size, 0, cu_stream>>>(
+      seq_lens.data<int>(),
+      stop_flags.data<bool>(),
+      reinterpret_cast<DataType_*>(
+          const_cast<data_t*>(input_data.data<data_t>())),
+      sequence_lengths.data<int>(),
+      seq_bs,
+      length);
+  return {sequence_lengths};
 }
 
-std::vector<paddle::Tensor> SetMaskValue(const paddle::Tensor& input_data, const paddle::Tensor& stop_flags, const paddle::Tensor& seq_lens) {
-    switch (input_data.type()) {
-        case paddle::DataType::BFLOAT16: {
-            return set_mask_value<paddle::DataType::BFLOAT16>(
-                input_data,
-                stop_flags,
-                seq_lens
-            );
-        }
-        case paddle::DataType::FLOAT16: {
-            return set_mask_value<paddle::DataType::FLOAT16>(
-                input_data,
-                stop_flags,
-                seq_lens
-            );
-        }
-        case paddle::DataType::FLOAT32: {
-            return set_mask_value<paddle::DataType::FLOAT32>(
-                input_data,
-                stop_flags,
-                seq_lens
-            );
-        }
-        default: {
-            PD_THROW(
-                "NOT supported data type. "
-                "Only float16, bfloat16 and float32 are supported. ");
-            break;
-        }
+std::vector<paddle::Tensor> SetMaskValue(const paddle::Tensor& input_data,
+                                         const paddle::Tensor& stop_flags,
+                                         const paddle::Tensor& seq_lens) {
+  switch (input_data.type()) {
+    case paddle::DataType::BFLOAT16: {
+      return set_mask_value<paddle::DataType::BFLOAT16>(
+          input_data, stop_flags, seq_lens);
     }
+    case paddle::DataType::FLOAT16: {
+      return set_mask_value<paddle::DataType::FLOAT16>(
+          input_data, stop_flags, seq_lens);
+    }
+    case paddle::DataType::FLOAT32: {
+      return set_mask_value<paddle::DataType::FLOAT32>(
+          input_data, stop_flags, seq_lens);
+    }
+    default: {
+      PD_THROW(
+          "NOT supported data type. "
+          "Only float16, bfloat16 and float32 are supported. ");
+      break;
+    }
+  }
 }
 
-std::vector<std::vector<int64_t>> SetMaskValueInferShape(const std::vector<int64_t>& input_data_shape, const std::vector<int64_t>& stop_flags_shape, const std::vector<int64_t>& seq_lens_shape) {
-    return {seq_lens_shape};
+std::vector<std::vector<int64_t>> SetMaskValueInferShape(
+    const std::vector<int64_t>& input_data_shape,
+    const std::vector<int64_t>& stop_flags_shape,
+    const std::vector<int64_t>& seq_lens_shape) {
+  return {seq_lens_shape};
 }
 
-std::vector<paddle::DataType> SetMaskValueInferDtype(const paddle::DataType& input_data_dtype, const paddle::DataType& stop_flags_dtype, const paddle::DataType& seq_lens_dtype) {
-    return {seq_lens_dtype};
+std::vector<paddle::DataType> SetMaskValueInferDtype(
+    const paddle::DataType& input_data_dtype,
+    const paddle::DataType& stop_flags_dtype,
+    const paddle::DataType& seq_lens_dtype) {
+  return {seq_lens_dtype};
 }
 
 PD_BUILD_STATIC_OP(set_mask_value)
diff --git a/custom_ops/gpu_ops/share_external_data.cu b/custom_ops/gpu_ops/share_external_data.cu
index 8b204ccc3f..8fd462582e 100644
--- a/custom_ops/gpu_ops/share_external_data.cu
+++ b/custom_ops/gpu_ops/share_external_data.cu
@@ -13,21 +13,20 @@
 // limitations under the License.
 
 #include "helper.h"
-#include<stdlib.h>
-#include<string.h>
-#include<sys/types.h>
-#include<sys/stat.h>
-#include<unistd.h>
-#include<fcntl.h>
-#include<sys/mman.h>
-#include<stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <unistd.h>
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <stdio.h>
 #include "cuda_multiprocess.h"
 #include "paddle/phi/core/tensor_meta.h"
 
-
-std::vector<paddle::Tensor> ShareExternalData(paddle::Tensor& input,
+std::vector<paddle::Tensor> ShareExternalData(paddle::Tensor &input,
                                               const std::string shm_name,
-                                              const std::vector<int>& shape) {
+                                              const std::vector<int> &shape) {
   volatile shmStruct *shm = NULL;
   sharedMemoryInfo info;
   if (sharedMemoryOpen(shm_name.c_str(), sizeof(shmStruct), &info) != 0) {
@@ -40,8 +39,8 @@ std::vector<paddle::Tensor> ShareExternalData(paddle::Tensor& input,
 #ifdef PADDLE_WITH_HIP
   checkCudaErrors(
       hipIpcOpenMemHandle(&ptr,
-                           *(hipIpcMemHandle_t *)&shm->memHandle,  // NOLINT
-                           hipIpcMemLazyEnablePeerAccess));
+                          *(hipIpcMemHandle_t *)&shm->memHandle,  // NOLINT
+                          hipIpcMemLazyEnablePeerAccess));
 #else
   checkCudaErrors(
       cudaIpcOpenMemHandle(&ptr,
@@ -49,11 +48,7 @@ std::vector<paddle::Tensor> ShareExternalData(paddle::Tensor& input,
                            cudaIpcMemLazyEnablePeerAccess));
 #endif
 
-  paddle::Tensor tmp_tensor = paddle::from_blob(
-    ptr,
-    shape,
-    input.type()
-  );
+  paddle::Tensor tmp_tensor = paddle::from_blob(ptr, shape, input.type());
   sharedMemoryClose(&info);
   return {tmp_tensor};
 }
diff --git a/custom_ops/gpu_ops/speculate_decoding/draft_model/draft_model_postprocess.cu b/custom_ops/gpu_ops/speculate_decoding/draft_model/draft_model_postprocess.cu
index 068d42b6c3..fa486e93da 100644
--- a/custom_ops/gpu_ops/speculate_decoding/draft_model/draft_model_postprocess.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/draft_model/draft_model_postprocess.cu
@@ -25,44 +25,40 @@ __global__ void draft_model_update_seq_lens_this_time_kernel(
     const bool* base_model_stop_flags,
     int bsz,
     int base_model_draft_token_len) {
-    int tid = threadIdx.x;
-    if (tid < bsz) {
-        if (!base_model_stop_flags[tid] &&
-            base_model_seq_lens_encoder[tid] == 0) {
-            const int64_t* base_model_draft_tokens_now =
-                base_model_draft_tokens + tid * base_model_draft_token_len;
-            int token_num = 0;
+  int tid = threadIdx.x;
+  if (tid < bsz) {
+    if (!base_model_stop_flags[tid] && base_model_seq_lens_encoder[tid] == 0) {
+      const int64_t* base_model_draft_tokens_now =
+          base_model_draft_tokens + tid * base_model_draft_token_len;
+      int token_num = 0;
 
-            for (int i = 0; i < base_model_draft_token_len; ++i) {
-                if (base_model_draft_tokens_now[i] != -1) {
-                    token_num++;
-                }
-            }
-            base_model_seq_lens_this_time[tid] = token_num;
-        } else if (base_model_stop_flags[tid]) {
-            base_model_seq_lens_this_time[tid] = 0;
+      for (int i = 0; i < base_model_draft_token_len; ++i) {
+        if (base_model_draft_tokens_now[i] != -1) {
+          token_num++;
         }
+      }
+      base_model_seq_lens_this_time[tid] = token_num;
+    } else if (base_model_stop_flags[tid]) {
+      base_model_seq_lens_this_time[tid] = 0;
     }
+  }
 }
 
 void DraftModelPostprocess(const paddle::Tensor& base_model_draft_tokens,
                            const paddle::Tensor& base_model_seq_lens_this_time,
                            const paddle::Tensor& base_model_seq_lens_encoder,
                            const paddle::Tensor& base_model_stop_flags) {
-    int real_bsz = base_model_seq_lens_this_time.shape()[0];
-    auto cu_stream = base_model_seq_lens_this_time.stream();
-    int block_size = (real_bsz + 32 - 1) / 32 * 32;
-    int base_model_draft_token_len = base_model_draft_tokens.shape()[1];
-    draft_model_update_seq_lens_this_time_kernel<<<1,
-                                                   block_size,
-                                                   0,
-                                                   cu_stream>>>(
-        base_model_draft_tokens.data<int64_t>(),
-        const_cast<int*>(base_model_seq_lens_this_time.data<int>()),
-        base_model_seq_lens_encoder.data<int>(),
-        base_model_stop_flags.data<bool>(),
-        real_bsz,
-        base_model_draft_token_len);
+  int real_bsz = base_model_seq_lens_this_time.shape()[0];
+  auto cu_stream = base_model_seq_lens_this_time.stream();
+  int block_size = (real_bsz + 32 - 1) / 32 * 32;
+  int base_model_draft_token_len = base_model_draft_tokens.shape()[1];
+  draft_model_update_seq_lens_this_time_kernel<<<1, block_size, 0, cu_stream>>>(
+      base_model_draft_tokens.data<int64_t>(),
+      const_cast<int*>(base_model_seq_lens_this_time.data<int>()),
+      base_model_seq_lens_encoder.data<int>(),
+      base_model_stop_flags.data<bool>(),
+      real_bsz,
+      base_model_draft_token_len);
 }
 
 PD_BUILD_STATIC_OP(draft_model_postprocess)
diff --git a/custom_ops/gpu_ops/speculate_decoding/draft_model/draft_model_set_value_by_flags.cu b/custom_ops/gpu_ops/speculate_decoding/draft_model/draft_model_set_value_by_flags.cu
index 37e98f050f..0e51778696 100644
--- a/custom_ops/gpu_ops/speculate_decoding/draft_model/draft_model_set_value_by_flags.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/draft_model/draft_model_set_value_by_flags.cu
@@ -22,21 +22,21 @@ __global__ void update_pre_ids_kernel(const int64_t* draft_tokens,
                                       int bs,
                                       int pre_id_length,
                                       int max_draft_token) {
-    int tid = threadIdx.x;
-    if (tid < bs && seq_lens_this_time[tid] != 0 && !stop_flags[tid]) {
-        int64_t* pre_ids_all_now = pre_ids_all + tid * pre_id_length;
-        const int64_t* draft_token_now = draft_tokens + tid * max_draft_token;
-        const int seq_len_this_time = seq_lens_this_time[tid];
-        if (step_idx[tid] - 1 > 0 /*Decoder Step*/) {
-            for (int i = 0; i < seq_len_this_time; ++i) {
-                pre_ids_all_now[step_idx[tid] - i] =
-                    draft_token_now[seq_len_this_time - 1 - i];
-            }
-        } else if (step_idx[tid] == 1 /*Encoder Step*/) {
-            pre_ids_all_now[1] = draft_token_now[0];
-        }
-        seq_lens_this_time[tid] = 1;
+  int tid = threadIdx.x;
+  if (tid < bs && seq_lens_this_time[tid] != 0 && !stop_flags[tid]) {
+    int64_t* pre_ids_all_now = pre_ids_all + tid * pre_id_length;
+    const int64_t* draft_token_now = draft_tokens + tid * max_draft_token;
+    const int seq_len_this_time = seq_lens_this_time[tid];
+    if (step_idx[tid] - 1 > 0 /*Decoder Step*/) {
+      for (int i = 0; i < seq_len_this_time; ++i) {
+        pre_ids_all_now[step_idx[tid] - i] =
+            draft_token_now[seq_len_this_time - 1 - i];
+      }
+    } else if (step_idx[tid] == 1 /*Encoder Step*/) {
+      pre_ids_all_now[1] = draft_token_now[0];
     }
+    seq_lens_this_time[tid] = 1;
+  }
 }
 
 void SpeculateDraftModelUpdate(const paddle::Tensor& draft_tokens,
@@ -46,21 +46,21 @@ void SpeculateDraftModelUpdate(const paddle::Tensor& draft_tokens,
                                const paddle::Tensor& seq_lens_encoder,
                                const paddle::Tensor& seq_lens_decoder,
                                const paddle::Tensor& step_idx) {
-    int64_t real_bs = seq_lens_this_time.shape()[0];
-    int64_t pre_id_length = pre_ids_all.shape()[1];
-    auto cu_stream = seq_lens_this_time.stream();
-    int64_t max_draft_token = draft_tokens.shape()[1];
+  int64_t real_bs = seq_lens_this_time.shape()[0];
+  int64_t pre_id_length = pre_ids_all.shape()[1];
+  auto cu_stream = seq_lens_this_time.stream();
+  int64_t max_draft_token = draft_tokens.shape()[1];
 
-    int block_size = (real_bs + 32 - 1) / 32 * 32;
-    update_pre_ids_kernel<<<1, block_size, 0, cu_stream>>>(
-        draft_tokens.data<int64_t>(),
-        const_cast<int64_t*>(pre_ids_all.data<int64_t>()),
-        stop_flags.data<bool>(),
-        const_cast<int*>(seq_lens_this_time.data<int>()),
-        step_idx.data<int64_t>(),
-        real_bs,
-        pre_id_length,
-        max_draft_token);
+  int block_size = (real_bs + 32 - 1) / 32 * 32;
+  update_pre_ids_kernel<<<1, block_size, 0, cu_stream>>>(
+      draft_tokens.data<int64_t>(),
+      const_cast<int64_t*>(pre_ids_all.data<int64_t>()),
+      stop_flags.data<bool>(),
+      const_cast<int*>(seq_lens_this_time.data<int>()),
+      step_idx.data<int64_t>(),
+      real_bs,
+      pre_id_length,
+      max_draft_token);
 }
 
 PD_BUILD_STATIC_OP(draft_model_set_value_by_flags)
diff --git a/custom_ops/gpu_ops/speculate_decoding/draft_model/eagle_get_hidden_states.cu b/custom_ops/gpu_ops/speculate_decoding/draft_model/eagle_get_hidden_states.cu
index 43501cdbd4..a325894517 100644
--- a/custom_ops/gpu_ops/speculate_decoding/draft_model/eagle_get_hidden_states.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/draft_model/eagle_get_hidden_states.cu
@@ -18,19 +18,18 @@
 
 // #define DEBUG_EAGLE_KERNEL
 
-__global__ void ComputeOrderKernel(
-                const int* seq_lens_this_time,
-                const int* seq_lens_encoder,
-                const int* base_model_seq_lens_this_time,
-                const int* base_model_seq_lens_encoder,
-                const int* accept_nums,
-                int* position_map,
-                int* output_token_num,
-                const int bsz,
-                const int actual_draft_token_num,
-                const int input_token_num) {
-  int in_offset = 0;     // input_offset(long)
-  int out_offset = 0;    // output_offset(short)
+__global__ void ComputeOrderKernel(const int* seq_lens_this_time,
+                                   const int* seq_lens_encoder,
+                                   const int* base_model_seq_lens_this_time,
+                                   const int* base_model_seq_lens_encoder,
+                                   const int* accept_nums,
+                                   int* position_map,
+                                   int* output_token_num,
+                                   const int bsz,
+                                   const int actual_draft_token_num,
+                                   const int input_token_num) {
+  int in_offset = 0;   // input_offset(long)
+  int out_offset = 0;  // output_offset(short)
   if (threadIdx.x == 0) {
     for (int i = 0; i < bsz; ++i) {
       int cur_base_model_seq_lens_this_time = base_model_seq_lens_this_time[i];
@@ -39,7 +38,13 @@ __global__ void ComputeOrderKernel(
       int accept_num = accept_nums[i];
       int cur_seq_lens_encoder = seq_lens_encoder[i];
 #ifdef DEBUG_EAGLE_KERNEL
-        printf("batch %d: cur_base_model_seq_lens_this_time%d. cur_seq_lens_this_time%d, accept_num %d\n", i, cur_base_model_seq_lens_this_time, cur_seq_lens_this_time, accept_num);
+      printf(
+          "batch %d: cur_base_model_seq_lens_this_time%d. "
+          "cur_seq_lens_this_time%d, accept_num %d\n",
+          i,
+          cur_base_model_seq_lens_this_time,
+          cur_seq_lens_this_time,
+          accept_num);
 #endif
       // 1. eagle encoder. Base step=1
       if (cur_seq_lens_encoder > 0) {
@@ -49,35 +54,40 @@ __global__ void ComputeOrderKernel(
         for (int j = 0; j < cur_seq_lens_encoder; j++) {
           position_map[in_offset++] = out_offset++;
         }
-      // 2. Base model stop at last verify-step.
-      } else if (cur_base_model_seq_lens_this_time != 0 && cur_seq_lens_this_time == 0) {
+        // 2. Base model stop at last verify-step.
+      } else if (cur_base_model_seq_lens_this_time != 0 &&
+                 cur_seq_lens_this_time == 0) {
 #ifdef DEBUG_EAGLE_KERNEL
         printf("batch %d: base=0. draft !=0 \n", i);
 #endif
 
         in_offset += cur_base_model_seq_lens_this_time;
-      // 4. stopped
-      } else if (cur_base_model_seq_lens_this_time == 0 && cur_seq_lens_this_time == 0) /* end */ {
+        // 4. stopped
+      } else if (cur_base_model_seq_lens_this_time == 0 &&
+                 cur_seq_lens_this_time == 0) /* end */ {
       } else {
         for (int i = 0; i < accept_num; i++) {
           position_map[in_offset++] = out_offset++;
         }
         in_offset += cur_base_model_seq_lens_this_time - accept_num;
-// (liuzichang): Temporary Reserved for debug
-//         if (accept_num <= actual_draft_token_num) /*Accept partial draft tokens*/ {
-// #ifdef DEBUG_EAGLE_KERNEL
-//         printf("batch %d: accept_num <= actual_draft_token_num \n", i);
-// #endif
-//           position_map[in_offset + accept_num - 1] = out_offset++;
-//           in_offset += cur_base_model_seq_lens_this_time;
-//         } else /*Accept all draft tokens*/ {
-// #ifdef DEBUG_EAGLE_KERNEL
-//         printf("batch %d: accept_num > actual_draft_token_num \n", i);
-// #endif
-//           position_map[in_offset + accept_num - 2] = out_offset++;
-//           position_map[in_offset + accept_num - 1] = out_offset++;
-//           in_offset += cur_base_model_seq_lens_this_time;
-//         }
+        // (liuzichang): Temporary Reserved for debug
+        //         if (accept_num <= actual_draft_token_num) /*Accept partial
+        //         draft tokens*/ {
+        // #ifdef DEBUG_EAGLE_KERNEL
+        //         printf("batch %d: accept_num <= actual_draft_token_num \n",
+        //         i);
+        // #endif
+        //           position_map[in_offset + accept_num - 1] = out_offset++;
+        //           in_offset += cur_base_model_seq_lens_this_time;
+        //         } else /*Accept all draft tokens*/ {
+        // #ifdef DEBUG_EAGLE_KERNEL
+        //         printf("batch %d: accept_num > actual_draft_token_num \n",
+        //         i);
+        // #endif
+        //           position_map[in_offset + accept_num - 2] = out_offset++;
+        //           position_map[in_offset + accept_num - 1] = out_offset++;
+        //           in_offset += cur_base_model_seq_lens_this_time;
+        //         }
       }
     }
     output_token_num[0] = out_offset;
@@ -89,26 +99,23 @@ __global__ void ComputeOrderKernel(
     printf("\n");
 #endif
   }
-
 }
 
 template <typename T, int VecSize>
-__global__ void rebuildHiddenStatesKernel(
-                  const T* input,
-                  const int* position_map,
-                  T* out,
-                  const int dim_embed,
-                  const int elem_cnt) {
+__global__ void rebuildHiddenStatesKernel(const T* input,
+                                          const int* position_map,
+                                          T* out,
+                                          const int dim_embed,
+                                          const int elem_cnt) {
   using LoadT = AlignedVector<T, VecSize>;
   LoadT src_vec;
 
   int global_thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
-  for (int elem_idx = global_thread_idx * VecSize; elem_idx < elem_cnt; elem_idx += blockDim.x * gridDim.x * VecSize) {
-
+  for (int elem_idx = global_thread_idx * VecSize; elem_idx < elem_cnt;
+       elem_idx += blockDim.x * gridDim.x * VecSize) {
     int ori_token_idx = elem_idx / dim_embed;
     int token_idx = position_map[ori_token_idx];
     if (token_idx >= 0) {
-
       int offset = elem_idx % dim_embed;
       if (token_idx == 0) {
       }
@@ -118,20 +125,17 @@ __global__ void rebuildHiddenStatesKernel(
   }
 }
 
-
-
 template <paddle::DataType D>
 std::vector<paddle::Tensor> DispatchDtype(
-                                const paddle::Tensor& input,
-                                const paddle::Tensor& seq_lens_this_time,
-                                const paddle::Tensor& seq_lens_encoder,
-                                const paddle::Tensor& seq_lens_decoder,
-                                const paddle::Tensor& stop_flags,
-                                const paddle::Tensor& accept_nums,
-                                const paddle::Tensor& base_model_seq_lens_this_time,
-                                const paddle::Tensor& base_model_seq_lens_encoder,
-                                const int actual_draft_token_num) {
-
+    const paddle::Tensor& input,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& stop_flags,
+    const paddle::Tensor& accept_nums,
+    const paddle::Tensor& base_model_seq_lens_this_time,
+    const paddle::Tensor& base_model_seq_lens_encoder,
+    const int actual_draft_token_num) {
   typedef PDTraits<D> traits_;
   typedef typename traits_::DataType DataType_;
   typedef typename traits_::data_t data_t;
@@ -143,24 +147,31 @@ std::vector<paddle::Tensor> DispatchDtype(
 
   int bsz = seq_lens_this_time.shape()[0];
 
-  auto position_map = paddle::empty({input_token_num}, seq_lens_this_time.dtype(), input.place());
-  cudaMemsetAsync(position_map.data<int>(), 0xFF, input_token_num * sizeof(seq_lens_this_time.dtype()), seq_lens_this_time.stream());
+  auto position_map = paddle::empty(
+      {input_token_num}, seq_lens_this_time.dtype(), input.place());
+  cudaMemsetAsync(position_map.data<int>(),
+                  0xFF,
+                  input_token_num * sizeof(seq_lens_this_time.dtype()),
+                  seq_lens_this_time.stream());
 
-  auto output_token_num  = paddle::empty({1}, seq_lens_this_time.dtype(), seq_lens_this_time.place());
-  ComputeOrderKernel<<<1,1>>>(seq_lens_this_time.data<int>(),
-                              seq_lens_encoder.data<int>(),
-                              base_model_seq_lens_this_time.data<int>(),
-                              base_model_seq_lens_encoder.data<int>(),
-                              accept_nums.data<int>(),
-                              position_map.data<int>(),
-                              output_token_num.data<int>(),
-                              bsz,
-                              actual_draft_token_num,
-                              input_token_num);
+  auto output_token_num = paddle::empty(
+      {1}, seq_lens_this_time.dtype(), seq_lens_this_time.place());
+  ComputeOrderKernel<<<1, 1>>>(seq_lens_this_time.data<int>(),
+                               seq_lens_encoder.data<int>(),
+                               base_model_seq_lens_this_time.data<int>(),
+                               base_model_seq_lens_encoder.data<int>(),
+                               accept_nums.data<int>(),
+                               position_map.data<int>(),
+                               output_token_num.data<int>(),
+                               bsz,
+                               actual_draft_token_num,
+                               input_token_num);
 
-  int output_token_num_cpu = output_token_num.copy_to(paddle::CPUPlace(), false).data<int>()[0];
+  int output_token_num_cpu =
+      output_token_num.copy_to(paddle::CPUPlace(), false).data<int>()[0];
 
-  auto out = paddle::empty({output_token_num_cpu, dim_embed}, input.dtype(), input.place());
+  auto out = paddle::empty(
+      {output_token_num_cpu, dim_embed}, input.dtype(), input.place());
 
   constexpr int packSize = VEC_16B / (sizeof(DataType_));
   int elem_cnt = input_token_num * dim_embed;
@@ -175,69 +186,67 @@ std::vector<paddle::Tensor> DispatchDtype(
 
   constexpr int thread_per_block = 128;
 
-  rebuildHiddenStatesKernel<DataType_, packSize><<<grid_size, thread_per_block>>>(
-    reinterpret_cast<const DataType_*>(input.data<data_t>()),
-    position_map.data<int>(),
-    reinterpret_cast<DataType_*>(out.data<data_t>()),
-    dim_embed,
-    elem_cnt);
+  rebuildHiddenStatesKernel<DataType_, packSize>
+      <<<grid_size, thread_per_block>>>(
+          reinterpret_cast<const DataType_*>(input.data<data_t>()),
+          position_map.data<int>(),
+          reinterpret_cast<DataType_*>(out.data<data_t>()),
+          dim_embed,
+          elem_cnt);
 
   return {out};
 }
 
-
 std::vector<paddle::Tensor> EagleGetHiddenStates(
-                                const paddle::Tensor& input,
-                                const paddle::Tensor& seq_lens_this_time,
-                                const paddle::Tensor& seq_lens_encoder,
-                                const paddle::Tensor& seq_lens_decoder,
-                                const paddle::Tensor& stop_flags,
-                                const paddle::Tensor& accept_nums,
-                                const paddle::Tensor& base_model_seq_lens_this_time,
-                                const paddle::Tensor& base_model_seq_lens_encoder,
-                                const int actual_draft_token_num) {
+    const paddle::Tensor& input,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& stop_flags,
+    const paddle::Tensor& accept_nums,
+    const paddle::Tensor& base_model_seq_lens_this_time,
+    const paddle::Tensor& base_model_seq_lens_encoder,
+    const int actual_draft_token_num) {
   switch (input.dtype()) {
     case paddle::DataType::FLOAT16: {
       return DispatchDtype<paddle::DataType::FLOAT16>(
-              input,
-              seq_lens_this_time,
-              seq_lens_encoder,
-              seq_lens_decoder,
-              stop_flags,
-              accept_nums,
-              base_model_seq_lens_this_time,
-              base_model_seq_lens_encoder,
-              actual_draft_token_num);
+          input,
+          seq_lens_this_time,
+          seq_lens_encoder,
+          seq_lens_decoder,
+          stop_flags,
+          accept_nums,
+          base_model_seq_lens_this_time,
+          base_model_seq_lens_encoder,
+          actual_draft_token_num);
     }
     case paddle::DataType::BFLOAT16: {
       return DispatchDtype<paddle::DataType::BFLOAT16>(
-              input,
-              seq_lens_this_time,
-              seq_lens_encoder,
-              seq_lens_decoder,
-              stop_flags,
-              accept_nums,
-              base_model_seq_lens_this_time,
-              base_model_seq_lens_encoder,
-              actual_draft_token_num);
+          input,
+          seq_lens_this_time,
+          seq_lens_encoder,
+          seq_lens_decoder,
+          stop_flags,
+          accept_nums,
+          base_model_seq_lens_this_time,
+          base_model_seq_lens_encoder,
+          actual_draft_token_num);
     }
     default: {
       PD_THROW("Not support this data type");
     }
   }
-
 }
 
-
 PD_BUILD_STATIC_OP(eagle_get_hidden_states)
-  .Inputs({"input",
-            "seq_lens_this_time",
-            "seq_lens_encoder",
-            "seq_lens_decoder",
-            "stop_flags",
-            "accept_nums",
-            "base_model_seq_lens_this_time",
-            "base_model_seq_lens_encoder"})
-  .Attrs({"actual_draft_token_num: int"})
-  .Outputs({"out"})
-  .SetKernelFn(PD_KERNEL(EagleGetHiddenStates));
+    .Inputs({"input",
+             "seq_lens_this_time",
+             "seq_lens_encoder",
+             "seq_lens_decoder",
+             "stop_flags",
+             "accept_nums",
+             "base_model_seq_lens_this_time",
+             "base_model_seq_lens_encoder"})
+    .Attrs({"actual_draft_token_num: int"})
+    .Outputs({"out"})
+    .SetKernelFn(PD_KERNEL(EagleGetHiddenStates));
diff --git a/custom_ops/gpu_ops/speculate_decoding/draft_model/eagle_get_self_hidden_states.cu b/custom_ops/gpu_ops/speculate_decoding/draft_model/eagle_get_self_hidden_states.cu
index 878926f3ba..0bb70f6088 100644
--- a/custom_ops/gpu_ops/speculate_decoding/draft_model/eagle_get_self_hidden_states.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/draft_model/eagle_get_self_hidden_states.cu
@@ -16,58 +16,64 @@
 #include "paddle/extension.h"
 #include "helper.h"
 
-
 // #define DEBUG_EAGLE_KERNEL
 
-__global__ void computeOrderKernel(
-                const int* last_seq_lens_this_time,
-                const int* seq_lens_this_time,
-                const int64_t* step_idx,
-                int* src_map,
-                int* output_token_num,
-                int bsz) {
+__global__ void computeOrderKernel(const int* last_seq_lens_this_time,
+                                   const int* seq_lens_this_time,
+                                   const int64_t* step_idx,
+                                   int* src_map,
+                                   int* output_token_num,
+                                   int bsz) {
   int in_offset = 0;
   int out_offset = 0;
   if (threadIdx.x == 0) {
     for (int i = 0; i < bsz; ++i) {
-        int cur_seq_lens_this_time = seq_lens_this_time[i];
-        int cur_last_seq_lens_this_time = last_seq_lens_this_time[i];
+      int cur_seq_lens_this_time = seq_lens_this_time[i];
+      int cur_last_seq_lens_this_time = last_seq_lens_this_time[i];
 #ifdef DEBUG_EAGLE_KERNEL
-        printf("batch %d: cur_seq_lens_this_time:%d. cur_last_seq_lens_this_time:%d\n", i, cur_seq_lens_this_time, cur_last_seq_lens_this_time);
+      printf(
+          "batch %d: cur_seq_lens_this_time:%d. "
+          "cur_last_seq_lens_this_time:%d\n",
+          i,
+          cur_seq_lens_this_time,
+          cur_last_seq_lens_this_time);
 #endif
-        // 1. encoder
-        if (step_idx[i] == 1 && cur_seq_lens_this_time > 0) {
+      // 1. encoder
+      if (step_idx[i] == 1 && cur_seq_lens_this_time > 0) {
 #ifdef DEBUG_EAGLE_KERNEL
         printf("batch %d last_step is encoder \n", i);
 #endif
-            in_offset += 1;
-            src_map[out_offset++] = in_offset - 1;
+        in_offset += 1;
+        src_map[out_offset++] = in_offset - 1;
 #ifdef DEBUG_EAGLE_KERNEL
-        printf("batch %d finish. src_map[%d]=%d \n", i, out_offset - 1, in_offset - 1);
+        printf("batch %d finish. src_map[%d]=%d \n",
+               i,
+               out_offset - 1,
+               in_offset - 1);
 #endif
         // 2. decoder
-        } else if (cur_seq_lens_this_time > 0) /* =1 */ {
+      } else if (cur_seq_lens_this_time > 0) /* =1 */ {
 #ifdef DEBUG_EAGLE_KERNEL
         printf("batch %d is decoder\n", i);
 #endif
-            in_offset += cur_last_seq_lens_this_time;
-            src_map[out_offset++] = in_offset - 1;
+        in_offset += cur_last_seq_lens_this_time;
+        src_map[out_offset++] = in_offset - 1;
         // 3. stop
-        } else {
-          // first token end
-          if (step_idx[i] == 1) {
+      } else {
+        // first token end
+        if (step_idx[i] == 1) {
 #ifdef DEBUG_EAGLE_KERNEL
-        printf("batch %d finished in first token \n", i);
+          printf("batch %d finished in first token \n", i);
 #endif
-            in_offset += cur_last_seq_lens_this_time > 0 ? 1 : 0;
+          in_offset += cur_last_seq_lens_this_time > 0 ? 1 : 0;
           // normal end
-          } else {
+        } else {
 #ifdef DEBUG_EAGLE_KERNEL
-        printf("batch %d finished in non-first token \n", i);
+          printf("batch %d finished in non-first token \n", i);
 #endif
-            in_offset += cur_last_seq_lens_this_time;
-          }
+          in_offset += cur_last_seq_lens_this_time;
         }
+      }
     }
     output_token_num[0] = out_offset;
 #ifdef DEBUG_EAGLE_KERNEL
@@ -77,22 +83,18 @@ __global__ void computeOrderKernel(
     }
     printf("\n");
 #endif
-
   }
 }
 
-template<typename T, int PackSize>
+template <typename T, int PackSize>
 __global__ void rebuildSelfHiddenStatesKernel(
-                    const T* input,
-                    int* src_map,
-                    T* output,
-                    int dim_embed,
-                    int elem_cnt) {
+    const T* input, int* src_map, T* output, int dim_embed, int elem_cnt) {
   using LoadT = AlignedVector<T, PackSize>;
   LoadT src_vec;
 
   int global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
-  for (int elem_id = global_thread_idx * PackSize; elem_id < elem_cnt; elem_id += blockDim.x * gridDim.x * PackSize) {
+  for (int elem_id = global_thread_idx * PackSize; elem_id < elem_cnt;
+       elem_id += blockDim.x * gridDim.x * PackSize) {
     int output_token_idx = elem_id / dim_embed;
     int input_token_idx = src_map[output_token_idx];
     int offset = elem_id % dim_embed;
@@ -101,13 +103,12 @@ __global__ void rebuildSelfHiddenStatesKernel(
   }
 }
 
-
-template<paddle::DataType D>
+template <paddle::DataType D>
 std::vector<paddle::Tensor> DispatchDtype(
-                    const paddle::Tensor input,
-                    const paddle::Tensor last_seq_lens_this_time,
-                    const paddle::Tensor seq_lens_this_time,
-                    const paddle::Tensor step_idx) {
+    const paddle::Tensor input,
+    const paddle::Tensor last_seq_lens_this_time,
+    const paddle::Tensor seq_lens_this_time,
+    const paddle::Tensor step_idx) {
   typedef PDTraits<D> traits_;
   typedef typename traits_::DataType DataType_;
   typedef typename traits_::data_t data_t;
@@ -115,24 +116,31 @@ std::vector<paddle::Tensor> DispatchDtype(
   int input_token_num = input.shape()[0];
   int dim_embed = input.shape()[1];
   int bsz = seq_lens_this_time.shape()[0];
-  auto src_map = paddle::full({input_token_num}, -1, seq_lens_this_time.dtype(), seq_lens_this_time.place());
-  auto output_token_num  = paddle::full({1}, 0, seq_lens_this_time.dtype(), seq_lens_this_time.place());
+  auto src_map = paddle::full({input_token_num},
+                              -1,
+                              seq_lens_this_time.dtype(),
+                              seq_lens_this_time.place());
+  auto output_token_num = paddle::full(
+      {1}, 0, seq_lens_this_time.dtype(), seq_lens_this_time.place());
 
   computeOrderKernel<<<1, 1, 0, seq_lens_this_time.stream()>>>(
-              last_seq_lens_this_time.data<int>(),
-              seq_lens_this_time.data<int>(),
-              step_idx.data<int64_t>(),
-              src_map.data<int>(),
-              output_token_num.data<int>(),
-              bsz);
+      last_seq_lens_this_time.data<int>(),
+      seq_lens_this_time.data<int>(),
+      step_idx.data<int64_t>(),
+      src_map.data<int>(),
+      output_token_num.data<int>(),
+      bsz);
 
-  int output_token_num_cpu = output_token_num.copy_to(paddle::CPUPlace(), false).data<int>()[0];
+  int output_token_num_cpu =
+      output_token_num.copy_to(paddle::CPUPlace(), false).data<int>()[0];
 
-  auto out = paddle::full({output_token_num_cpu, dim_embed}, -1, input.type(), input.place());
+  auto out = paddle::full(
+      {output_token_num_cpu, dim_embed}, -1, input.type(), input.place());
 
   constexpr int packSize = VEC_16B / (sizeof(DataType_));
   int elem_cnt = output_token_num_cpu * dim_embed;
-  // printf("output_token_num: %d, dim_embed: %d, cnt: %d. packSize: %d\n", output_token_num_cpu, dim_embed,elem_cnt, packSize);
+  // printf("output_token_num: %d, dim_embed: %d, cnt: %d. packSize: %d\n",
+  // output_token_num_cpu, dim_embed,elem_cnt, packSize);
   assert(elem_cnt % packSize == 0);
 
   int pack_num = elem_cnt / packSize;
@@ -143,48 +151,36 @@ std::vector<paddle::Tensor> DispatchDtype(
 
   constexpr int threadPerBlock = 128;
 
-  rebuildSelfHiddenStatesKernel<DataType_, packSize><<<grid_size, threadPerBlock, 0, input.stream()>>>(
-    reinterpret_cast<const DataType_*>(input.data<data_t>()),
-    src_map.data<int>(),
-    reinterpret_cast<DataType_*>(out.data<data_t>()),
-    dim_embed,
-    elem_cnt);
-
+  rebuildSelfHiddenStatesKernel<DataType_, packSize>
+      <<<grid_size, threadPerBlock, 0, input.stream()>>>(
+          reinterpret_cast<const DataType_*>(input.data<data_t>()),
+          src_map.data<int>(),
+          reinterpret_cast<DataType_*>(out.data<data_t>()),
+          dim_embed,
+          elem_cnt);
 
   return {out};
-
 }
 
-
-
 std::vector<paddle::Tensor> EagleGetSelfHiddenStates(
-                    const paddle::Tensor& input,
-                    const paddle::Tensor& last_seq_lens_this_time,
-                    const paddle::Tensor& seq_lens_this_time,
-                    const paddle::Tensor& step_idx) {
+    const paddle::Tensor& input,
+    const paddle::Tensor& last_seq_lens_this_time,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& step_idx) {
   switch (input.dtype()) {
     case paddle::DataType::BFLOAT16:
       return DispatchDtype<paddle::DataType::BFLOAT16>(
-        input,
-        last_seq_lens_this_time,
-        seq_lens_this_time,
-        step_idx);
+          input, last_seq_lens_this_time, seq_lens_this_time, step_idx);
     case paddle::DataType::FLOAT16:
       return DispatchDtype<paddle::DataType::FLOAT16>(
-        input,
-        last_seq_lens_this_time,
-        seq_lens_this_time,
-        step_idx);
+          input, last_seq_lens_this_time, seq_lens_this_time, step_idx);
     default:
-        PD_THROW("Not support this data type");
+      PD_THROW("Not support this data type");
   }
 }
 
-
 PD_BUILD_STATIC_OP(eagle_get_self_hidden_states)
-  .Inputs({"input",
-           "last_seq_lens_this_time",
-           "seq_lens_this_time",
-           "step_idx"})
-  .Outputs({"out"})
-  .SetKernelFn(PD_KERNEL(EagleGetSelfHiddenStates));
+    .Inputs(
+        {"input", "last_seq_lens_this_time", "seq_lens_this_time", "step_idx"})
+    .Outputs({"out"})
+    .SetKernelFn(PD_KERNEL(EagleGetSelfHiddenStates));
diff --git a/custom_ops/gpu_ops/speculate_decoding/draft_model/hydra_fetch_hidden_states.cu b/custom_ops/gpu_ops/speculate_decoding/draft_model/hydra_fetch_hidden_states.cu
index db121890df..acb3c34eac 100644
--- a/custom_ops/gpu_ops/speculate_decoding/draft_model/hydra_fetch_hidden_states.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/draft_model/hydra_fetch_hidden_states.cu
@@ -23,23 +23,23 @@ __global__ void HydraFetchHiddenStatesKernel(const T* output_hidden_states,
                                              const int bsz,
                                              const int max_seq_len,
                                              const int hidden_size) {
-    const int token_id = blockIdx.x;
-    const int ori_token_id = token_id + output_padding_offset[token_id];
-    const int bid = ori_token_id / max_seq_len;
-    const int start_ori_token_id = bid * max_seq_len;
-    const int local_token_id = ori_token_id - start_ori_token_id;
+  const int token_id = blockIdx.x;
+  const int ori_token_id = token_id + output_padding_offset[token_id];
+  const int bid = ori_token_id / max_seq_len;
+  const int start_ori_token_id = bid * max_seq_len;
+  const int local_token_id = ori_token_id - start_ori_token_id;
 
-    if (local_token_id != accept_token_num[bid] - 1) return;
+  if (local_token_id != accept_token_num[bid] - 1) return;
 
-    using LoadT = AlignedVector<T, VecSize>;
+  using LoadT = AlignedVector<T, VecSize>;
 
-    LoadT vec;
+  LoadT vec;
 
-    for (int idx = threadIdx.x * VecSize; idx < hidden_size;
-         idx += blockDim.x * VecSize) {
-        Load(&output_hidden_states[token_id * hidden_size + idx], &vec);
-        Store(vec, &hidden_states[bid * hidden_size + idx]);
-    }
+  for (int idx = threadIdx.x * VecSize; idx < hidden_size;
+       idx += blockDim.x * VecSize) {
+    Load(&output_hidden_states[token_id * hidden_size + idx], &vec);
+    Store(vec, &hidden_states[bid * hidden_size + idx]);
+  }
 }
 
 template <paddle::DataType D>
@@ -49,32 +49,32 @@ std::vector<paddle::Tensor> HydraFetchHiddenStatesImpl(
     const paddle::Tensor& accept_token_num,
     const int bsz,
     const int max_seq_length) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
 
-    auto cu_stream = output_hidden_states.stream();
+  auto cu_stream = output_hidden_states.stream();
 
-    auto output_token_num = output_hidden_states.shape()[0];
-    auto hidden_size = output_hidden_states.shape()[1];
+  auto output_token_num = output_hidden_states.shape()[0];
+  auto hidden_size = output_hidden_states.shape()[1];
 
-    auto hidden_states = paddle::full({bsz, hidden_size},
-                                      0,
-                                      output_hidden_states.dtype(),
-                                      output_hidden_states.place());
+  auto hidden_states = paddle::full({bsz, hidden_size},
+                                    0,
+                                    output_hidden_states.dtype(),
+                                    output_hidden_states.place());
 
-    constexpr int VecSize = 16 / sizeof(data_t);
+  constexpr int VecSize = 16 / sizeof(data_t);
 
-    HydraFetchHiddenStatesKernel<data_t, VecSize>
-        <<<output_token_num, 256, 0, cu_stream>>>(
-            output_hidden_states.data<data_t>(),
-            output_padding_offset.data<int>(),
-            accept_token_num.data<int>(),
-            hidden_states.data<data_t>(),
-            bsz,
-            max_seq_length,
-            hidden_size);
-    return {hidden_states};  // , enc_token_num, dec_token_num};
+  HydraFetchHiddenStatesKernel<data_t, VecSize>
+      <<<output_token_num, 256, 0, cu_stream>>>(
+          output_hidden_states.data<data_t>(),
+          output_padding_offset.data<int>(),
+          accept_token_num.data<int>(),
+          hidden_states.data<data_t>(),
+          bsz,
+          max_seq_length,
+          hidden_size);
+  return {hidden_states};  // , enc_token_num, dec_token_num};
 }
 
 std::vector<paddle::Tensor> HydraFetchHiddenStates(
@@ -83,30 +83,30 @@ std::vector<paddle::Tensor> HydraFetchHiddenStates(
     const paddle::Tensor& accept_token_num,
     const int bsz,
     const int max_seq_length) {
-    switch (output_hidden_states.dtype()) {
-        case paddle::DataType::BFLOAT16: {
-            return HydraFetchHiddenStatesImpl<paddle::DataType::BFLOAT16>(
-                output_hidden_states,
-                output_padding_offset,
-                accept_token_num,
-                bsz,
-                max_seq_length);
-        }
-        case paddle::DataType::FLOAT16: {
-            return HydraFetchHiddenStatesImpl<paddle::DataType::FLOAT16>(
-                output_hidden_states,
-                output_padding_offset,
-                accept_token_num,
-                bsz,
-                max_seq_length);
-        }
-        default: {
-            PD_THROW(
-                "NOT supported data type. "
-                "Only float16, bfloat16 are supported. ");
-            break;
-        }
+  switch (output_hidden_states.dtype()) {
+    case paddle::DataType::BFLOAT16: {
+      return HydraFetchHiddenStatesImpl<paddle::DataType::BFLOAT16>(
+          output_hidden_states,
+          output_padding_offset,
+          accept_token_num,
+          bsz,
+          max_seq_length);
     }
+    case paddle::DataType::FLOAT16: {
+      return HydraFetchHiddenStatesImpl<paddle::DataType::FLOAT16>(
+          output_hidden_states,
+          output_padding_offset,
+          accept_token_num,
+          bsz,
+          max_seq_length);
+    }
+    default: {
+      PD_THROW(
+          "NOT supported data type. "
+          "Only float16, bfloat16 are supported. ");
+      break;
+    }
+  }
 }
 
 std::vector<std::vector<int64_t>> HydraFetchHiddenStatesInferShape(
@@ -115,14 +115,14 @@ std::vector<std::vector<int64_t>> HydraFetchHiddenStatesInferShape(
     const std::vector<int64_t>& accept_token_num_shape,
     const int bsz,
     const int max_seq_length) {
-    return {{bsz, output_hidden_states_shape[1]}};
+  return {{bsz, output_hidden_states_shape[1]}};
 }
 
 std::vector<paddle::DataType> HydraFetchHiddenStatesInferDtype(
     const paddle::DataType& output_hidden_states_dtype,
     const paddle::DataType& output_padding_offset_dtype,
     const paddle::DataType& accept_token_num_dtype) {
-    return {output_hidden_states_dtype};
+  return {output_hidden_states_dtype};
 }
 
 PD_BUILD_STATIC_OP(hydra_fetch_hidden_states)
diff --git a/custom_ops/gpu_ops/speculate_decoding/draft_model/mtp_step_paddle.cu b/custom_ops/gpu_ops/speculate_decoding/draft_model/mtp_step_paddle.cu
index 9aa1a56090..c6a616e6f4 100644
--- a/custom_ops/gpu_ops/speculate_decoding/draft_model/mtp_step_paddle.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/draft_model/mtp_step_paddle.cu
@@ -13,146 +13,152 @@
 // limitations under the License.
 
 #include "helper.h"  // NOLINT
-template<int NUM_THREADS, int MAX_BATCH_SIZE=256>
-__global__ void mtp_free_and_dispatch_block(
-    bool *base_model_stop_flags,
-    bool *stop_flags,
-    bool *batch_drop,
-    int *seq_lens_this_time,
-    int *seq_lens_decoder,
-    int *block_tables,
-    int *encoder_block_lens,
-    int *used_list_len,
-    int *free_list,
-    int *free_list_len,
-    const int bsz,
-    const int block_size,
-    const int block_num_per_seq,
-    const int max_draft_tokens) {
+template <int NUM_THREADS, int MAX_BATCH_SIZE = 256>
+__global__ void mtp_free_and_dispatch_block(bool *base_model_stop_flags,
+                                            bool *stop_flags,
+                                            bool *batch_drop,
+                                            int *seq_lens_this_time,
+                                            int *seq_lens_decoder,
+                                            int *block_tables,
+                                            int *encoder_block_lens,
+                                            int *used_list_len,
+                                            int *free_list,
+                                            int *free_list_len,
+                                            const int bsz,
+                                            const int block_size,
+                                            const int block_num_per_seq,
+                                            const int max_draft_tokens) {
+  typedef cub::BlockReduce<cub::KeyValuePair<int, int>, NUM_THREADS>
+      BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
 
-    typedef cub::BlockReduce<cub::KeyValuePair<int, int>, NUM_THREADS> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ int need_block_len;
+  __shared__ int need_block_list[MAX_BATCH_SIZE];
+  const int tid = threadIdx.x;
 
-    __shared__ int need_block_len;
-    __shared__ int need_block_list[MAX_BATCH_SIZE];
-    const int tid = threadIdx.x;
-
-    if (tid < bsz) {
-        if (tid == 0) {
-            need_block_len = 0;
-        }
-        need_block_list[tid] = 0;
-        int *block_table_now = block_tables + tid * block_num_per_seq;
-        if (base_model_stop_flags[tid] || batch_drop[tid]) {
-            // 回收block块
-            const int encoder_block_len = encoder_block_lens[tid];
-            const int decoder_used_len = used_list_len[tid];
-            if (decoder_used_len > 0) {
-                const int ori_free_list_len =
-                    atomicAdd(free_list_len, decoder_used_len);
-#ifdef DEBUG_STEP
-                printf(
-                    "free block seq_id: %d, free block num: %d, "
-                    "encoder_block_len: %d, ori_free_list_len: %d\n",
-                    tid,
-                    decoder_used_len,
-                    encoder_block_len,
-                    ori_free_list_len);
-#endif
-                for (int i = 0; i < decoder_used_len; i++) {
-                    free_list[ori_free_list_len + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                encoder_block_lens[tid] = 0;
-                used_list_len[tid] = 0;
-            }
-        }
+  if (tid < bsz) {
+    if (tid == 0) {
+      need_block_len = 0;
     }
-    __syncthreads();
-    if (tid < bsz) {
-        int *block_table_now = block_tables + tid * block_num_per_seq;
-        int max_possible_block_idx = (seq_lens_decoder[tid] + max_draft_tokens + 1) / block_size;
-        if (!base_model_stop_flags[tid] && !batch_drop[tid] && max_possible_block_idx < block_num_per_seq &&
-                   block_table_now[max_possible_block_idx] == -1) {
-            int ori_need_block_len = atomicAdd(&need_block_len, 1);
-            need_block_list[ori_need_block_len] = tid;
-            // 统计需要分配block的位置和总数
-            // const int ori_free_list_len = atomicSub(free_list_len, 1);
-            // block_table_now[(seq_lens_decoder[tid] + max_draft_tokens + 1) / block_size] =
-            //                 free_list[ori_free_list_len - 1];
-            // used_list_len[tid] += 1;
+    need_block_list[tid] = 0;
+    int *block_table_now = block_tables + tid * block_num_per_seq;
+    if (base_model_stop_flags[tid] || batch_drop[tid]) {
+      // 回收block块
+      const int encoder_block_len = encoder_block_lens[tid];
+      const int decoder_used_len = used_list_len[tid];
+      if (decoder_used_len > 0) {
+        const int ori_free_list_len =
+            atomicAdd(free_list_len, decoder_used_len);
+#ifdef DEBUG_STEP
+        printf(
+            "free block seq_id: %d, free block num: %d, "
+            "encoder_block_len: %d, ori_free_list_len: %d\n",
+            tid,
+            decoder_used_len,
+            encoder_block_len,
+            ori_free_list_len);
+#endif
+        for (int i = 0; i < decoder_used_len; i++) {
+          free_list[ori_free_list_len + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        encoder_block_lens[tid] = 0;
+        used_list_len[tid] = 0;
+      }
+    }
+  }
+  __syncthreads();
+  if (tid < bsz) {
+    int *block_table_now = block_tables + tid * block_num_per_seq;
+    int max_possible_block_idx =
+        (seq_lens_decoder[tid] + max_draft_tokens + 1) / block_size;
+    if (!base_model_stop_flags[tid] && !batch_drop[tid] &&
+        max_possible_block_idx < block_num_per_seq &&
+        block_table_now[max_possible_block_idx] == -1) {
+      int ori_need_block_len = atomicAdd(&need_block_len, 1);
+      need_block_list[ori_need_block_len] = tid;
+      // 统计需要分配block的位置和总数
+      // const int ori_free_list_len = atomicSub(free_list_len, 1);
+      // block_table_now[(seq_lens_decoder[tid] + max_draft_tokens + 1) /
+      // block_size] =
+      //                 free_list[ori_free_list_len - 1];
+      // used_list_len[tid] += 1;
 
 #ifdef DEBUG_STEP
-            printf("seq_id: %d need block\n", tid);
+      printf("seq_id: %d need block\n", tid);
 #endif
-        }
     }
-    __syncthreads();
+  }
+  __syncthreads();
+#ifdef DEBUG_STEP
+  if (tid == 0) {
+    printf("need_block_len:%d, free_list_len: %d\n",
+           need_block_len,
+           free_list_len[0]);
+  }
+#endif
+  // 这里直接从 bid 0 开始遍历
+  while (need_block_len > free_list_len[0]) {
 #ifdef DEBUG_STEP
     if (tid == 0) {
-        printf("need_block_len:%d, free_list_len: %d\n", need_block_len, free_list_len[0]);
+      printf("in while need_block_len:%d, free_list_len: %d\n",
+             need_block_len,
+             free_list_len[0]);
     }
 #endif
-    // 这里直接从 bid 0 开始遍历
-    while (need_block_len > free_list_len[0]) {
+    const int used_block_num =
+        tid < bsz && !base_model_stop_flags[tid] ? used_list_len[tid] : 0;
+    cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
+    kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
+
+    if (tid == 0) {
+      const int encoder_block_len = encoder_block_lens[kv_pair.key];
+      int *block_table_now = block_tables + kv_pair.key * block_num_per_seq;
+      for (int i = 0; i < kv_pair.value; i++) {
+        free_list[free_list_len[0] + i] =
+            block_table_now[encoder_block_len + i];
+        block_table_now[encoder_block_len + i] = -1;
+      }
+      const int ori_free_list_len = atomicAdd(free_list_len, kv_pair.value);
+
+      printf(
+          "MTP STEP need_block_len: %d. free_list_len: %d."
+          "Drop bid: %d, free block num: %d, "
+          "encoder_block_len: %d,"
+          "After drop free_list_len %d \n",
+          need_block_len,
+          ori_free_list_len,
+          kv_pair.key,
+          kv_pair.value,
+          encoder_block_len,
+          free_list_len[0]);
+      stop_flags[kv_pair.key] = true;
+      batch_drop[kv_pair.key] = true;
+      seq_lens_this_time[kv_pair.key] = 0;
+      seq_lens_decoder[kv_pair.key] = 0;
+      used_list_len[kv_pair.key] = 0;
+    }
+    __syncthreads();
+  }
+
+  if (tid < need_block_len) {
+    const int need_block_id = need_block_list[tid];
+    // 这里必须用 batch_drop, 不能用 stop_flags
+    if (!batch_drop[need_block_id]) {
+      used_list_len[need_block_id] += 1;
+      const int ori_free_list_len = atomicSub(free_list_len, 1);
+      int *block_table_now = block_tables + need_block_id * block_num_per_seq;
 #ifdef DEBUG_STEP
-        if (tid == 0) {
-            printf("in while need_block_len:%d, free_list_len: %d\n", need_block_len, free_list_len[0]);
-        }
+      printf("bid: %d allocate block_id %d. seq_lens_decoder:%d \n",
+             need_block_id,
+             free_list[ori_free_list_len - 1],
+             seq_lens_decoder[need_block_id]);
 #endif
-        const int used_block_num =
-            tid < bsz && !base_model_stop_flags[tid] ? used_list_len[tid] : 0;
-        cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
-        kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
-
-        if (tid == 0) {
-            const int encoder_block_len = encoder_block_lens[kv_pair.key];
-            int *block_table_now =
-                block_tables + kv_pair.key * block_num_per_seq;
-            for (int i = 0; i < kv_pair.value; i++) {
-                free_list[free_list_len[0] + i] =
-                    block_table_now[encoder_block_len + i];
-                block_table_now[encoder_block_len + i] = -1;
-            }
-            const int ori_free_list_len = atomicAdd(free_list_len, kv_pair.value);
-
-            printf(
-                "MTP STEP need_block_len: %d. free_list_len: %d."
-                "Drop bid: %d, free block num: %d, "
-                "encoder_block_len: %d,"
-                "After drop free_list_len %d \n",
-                need_block_len,
-                ori_free_list_len,
-                kv_pair.key,
-                kv_pair.value,
-                encoder_block_len,
-                free_list_len[0]);
-            stop_flags[kv_pair.key] = true;
-            batch_drop[kv_pair.key] = true;
-            seq_lens_this_time[kv_pair.key] = 0;
-            seq_lens_decoder[kv_pair.key] = 0;
-            used_list_len[kv_pair.key] = 0;
-        }
-        __syncthreads();
-    }
-
-    if (tid < need_block_len) {
-        const int need_block_id = need_block_list[tid];
-        // 这里必须用 batch_drop, 不能用 stop_flags
-        if (!batch_drop[need_block_id]) {
-            used_list_len[need_block_id] += 1;
-            const int ori_free_list_len = atomicSub(free_list_len, 1);
-            int *block_table_now =
-                    block_tables + need_block_id * block_num_per_seq;
-#ifdef DEBUG_STEP
-            printf("bid: %d allocate block_id %d. seq_lens_decoder:%d \n", need_block_id, free_list[ori_free_list_len - 1], seq_lens_decoder[need_block_id]);
-#endif
-            block_table_now[(seq_lens_decoder[need_block_id] +
-                                 max_draft_tokens + 1) /
-                                block_size] = free_list[ori_free_list_len - 1];
-        }
+      block_table_now[(seq_lens_decoder[need_block_id] + max_draft_tokens + 1) /
+                      block_size] = free_list[ori_free_list_len - 1];
     }
+  }
 }
 
 void MTPStepPaddle(
@@ -169,34 +175,35 @@ void MTPStepPaddle(
     const paddle::Tensor &free_list_len,
     const int block_size,
     const int max_draft_tokens) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    constexpr int BlockSize = 512;  // bsz <= 256
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  constexpr int BlockSize = 512;  // bsz <= 256
 #ifdef DEBUG_STEP
-    printf(
-        "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
-        "%d\n",
-        bsz,
-        block_num_per_seq,
-        length,
-        max_decoder_block_num);
+  printf(
+      "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
+      "%d\n",
+      bsz,
+      block_num_per_seq,
+      length,
+      max_decoder_block_num);
 #endif
-    mtp_free_and_dispatch_block<BlockSize, BlockSize><<<1, BlockSize, 0, cu_stream>>>(
-        const_cast<bool *>(base_model_stop_flags.data<bool>()),
-        const_cast<bool *>(stop_flags.data<bool>()),
-        const_cast<bool *>(batch_drop.data<bool>()),
-        const_cast<int *>(seq_lens_this_time.data<int>()),
-        const_cast<int *>(seq_lens_decoder.data<int>()),
-        const_cast<int *>(block_tables.data<int>()),
-        const_cast<int *>(encoder_block_lens.data<int>()),
-        const_cast<int *>(used_list_len.data<int>()),
-        const_cast<int *>(free_list.data<int>()),
-        const_cast<int *>(free_list_len.data<int>()),
-        bsz,
-        block_size,
-        block_num_per_seq,
-        max_draft_tokens);
+  mtp_free_and_dispatch_block<BlockSize, BlockSize>
+      <<<1, BlockSize, 0, cu_stream>>>(
+          const_cast<bool *>(base_model_stop_flags.data<bool>()),
+          const_cast<bool *>(stop_flags.data<bool>()),
+          const_cast<bool *>(batch_drop.data<bool>()),
+          const_cast<int *>(seq_lens_this_time.data<int>()),
+          const_cast<int *>(seq_lens_decoder.data<int>()),
+          const_cast<int *>(block_tables.data<int>()),
+          const_cast<int *>(encoder_block_lens.data<int>()),
+          const_cast<int *>(used_list_len.data<int>()),
+          const_cast<int *>(free_list.data<int>()),
+          const_cast<int *>(free_list_len.data<int>()),
+          bsz,
+          block_size,
+          block_num_per_seq,
+          max_draft_tokens);
 }
 
 PD_BUILD_STATIC_OP(mtp_step_paddle)
@@ -211,8 +218,7 @@ PD_BUILD_STATIC_OP(mtp_step_paddle)
              "used_list_len",
              "free_list",
              "free_list_len"})
-    .Attrs({"block_size: int",
-            "max_draft_tokens: int"})
+    .Attrs({"block_size: int", "max_draft_tokens: int"})
     .Outputs({"block_tables_out",
               "stop_flags_out",
               "used_list_len_out",
diff --git a/custom_ops/gpu_ops/speculate_decoding/draft_model/ngram_match_mixed.cu b/custom_ops/gpu_ops/speculate_decoding/draft_model/ngram_match_mixed.cu
index 9b5f4b5b7c..529bfd9ab0 100644
--- a/custom_ops/gpu_ops/speculate_decoding/draft_model/ngram_match_mixed.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/draft_model/ngram_match_mixed.cu
@@ -24,191 +24,208 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-
 int sum_mixed(const int *value, int num) {
-    int sum_value = 0;
-    for (int i = 0; i <= num; i++) {
-        sum_value += value[i];
-    }
-    return sum_value;
+  int sum_value = 0;
+  for (int i = 0; i <= num; i++) {
+    sum_value += value[i];
+  }
+  return sum_value;
 }
 
 void find_candidate_pred_tokens_mixed(const int64_t *input_ids,
-        const int64_t *input_ids_len,
-        const int64_t *pre_ids,
-        const int64_t *step_idx,
-        const int *draft_token_num,
-        int64_t *draft_tokens,
-        int32_t *seq_lens_this_time,
-        int32_t *seq_lens_decoder,
-        int64_t *max_dec_len,
-        int64_t input_ids_stride,
-        int64_t pre_ids_stride,
-        int64_t draft_tokens_stride,
-        int64_t max_batch_size,
-        int max_ngram_size = 3,
-        int min_ngram_size = 1,
-        const int max_draft_tokens = 10) {
-    int threshold = 1024;
-    // dynamic in future
-    char *env_var = getenv("SPEC_TOKENUM_THRESHOLD");
-    if (env_var) {
-        threshold = std::stoi(env_var);
+                                      const int64_t *input_ids_len,
+                                      const int64_t *pre_ids,
+                                      const int64_t *step_idx,
+                                      const int *draft_token_num,
+                                      int64_t *draft_tokens,
+                                      int32_t *seq_lens_this_time,
+                                      int32_t *seq_lens_decoder,
+                                      int64_t *max_dec_len,
+                                      int64_t input_ids_stride,
+                                      int64_t pre_ids_stride,
+                                      int64_t draft_tokens_stride,
+                                      int64_t max_batch_size,
+                                      int max_ngram_size = 3,
+                                      int min_ngram_size = 1,
+                                      const int max_draft_tokens = 10) {
+  int threshold = 1024;
+  // dynamic in future
+  char *env_var = getenv("SPEC_TOKENUM_THRESHOLD");
+  if (env_var) {
+    threshold = std::stoi(env_var);
+  }
+  int unprocessed_batch_size = 0;
+  for (int batch_idx = 0; batch_idx < max_batch_size; batch_idx++) {
+    if (seq_lens_decoder[batch_idx] > 0) {
+      unprocessed_batch_size++;
     }
-    int unprocessed_batch_size = 0;
-    for (int batch_idx = 0; batch_idx < max_batch_size; batch_idx++) {
-        if (seq_lens_decoder[batch_idx] > 0) {
-            unprocessed_batch_size++;
-        }
+  }
+  for (int batch_idx = 0; batch_idx < max_batch_size; batch_idx++) {
+    const int ori_seq_len_this_time = seq_lens_this_time[batch_idx];
+    int max_draft_tokens_query = std::min(
+        static_cast<int64_t>(max_draft_tokens - ori_seq_len_this_time + 1),
+        max_dec_len[batch_idx] - step_idx[batch_idx] - 1);
+
+    if (ori_seq_len_this_time == 0 || max_draft_tokens_query <= 0) {
+      continue;
     }
-    for (int batch_idx = 0; batch_idx < max_batch_size; batch_idx++) {
-        const int ori_seq_len_this_time = seq_lens_this_time[batch_idx];
-        int max_draft_tokens_query = std::min(static_cast<int64_t>(
-            max_draft_tokens - ori_seq_len_this_time + 1), max_dec_len[batch_idx] - step_idx[batch_idx] - 1);
 
-        if (ori_seq_len_this_time == 0 || max_draft_tokens_query <= 0) {
-            continue;
-        }
+    const int64_t *cur_input_ids = input_ids + batch_idx * input_ids_stride;
+    int64_t *cur_draft_tokens = draft_tokens + batch_idx * draft_tokens_stride;
+    const int64_t *cur_pre_ids = pre_ids + batch_idx * pre_ids_stride;
+    const int64_t cur_step_idx = step_idx[batch_idx];
+    const int64_t cur_input_ids_len = input_ids_len[batch_idx];
+    unprocessed_batch_size--;
 
-        const int64_t *cur_input_ids = input_ids + batch_idx * input_ids_stride;
-        int64_t *cur_draft_tokens = draft_tokens + batch_idx * draft_tokens_stride;
-        const int64_t *cur_pre_ids = pre_ids + batch_idx * pre_ids_stride;
-        const int64_t cur_step_idx = step_idx[batch_idx];
-        const int64_t cur_input_ids_len = input_ids_len[batch_idx];
-        unprocessed_batch_size--;
+    auto sum_token_num = sum_mixed(seq_lens_this_time, batch_idx);
+    int left_min_token_num = unprocessed_batch_size;
 
-        auto sum_token_num = sum_mixed(seq_lens_this_time, batch_idx);
-        int left_min_token_num = unprocessed_batch_size;
-
-        if (sum_token_num + max_draft_tokens_query + left_min_token_num > threshold) {
-            int tmp_max_draft_tokens = threshold - sum_token_num - left_min_token_num;
-            max_draft_tokens_query = std::min(max_draft_tokens_query, tmp_max_draft_tokens);
-        }
-
-        if (sum_token_num + left_min_token_num >= threshold - 1) {
-            continue;
-        }
-        bool match_global = false;
-        // apply ngram_match in input_ids
-        for (int ngram_size = max_ngram_size; ngram_size >= min_ngram_size && !match_global; --ngram_size) {
-            // Extract the last n tokens as our search ngram
-            if (cur_step_idx < ngram_size) {
-                continue;
-            }
-            const int64_t *ngram = cur_pre_ids + (cur_step_idx + 1 - ngram_size);
-
-            // Iterate through sliding windows of size ngram_size
-            // bool match_input = false;
-            for (int64_t i = 0; i <= cur_input_ids_len - ngram_size && !match_global; ++i) {
-                // Check if the current window matches the ngram
-                bool match_local = true;
-                for (int j = 0; j < ngram_size; j++) {
-                    if (ngram[j] != cur_input_ids[i + j]) {
-                        match_local = false;
-                        break;
-                    }
-                }
-                if (match_local) {
-                    int64_t start_idx = i + ngram_size;
-                    int64_t end_idx = std::min(start_idx + max_draft_tokens_query, cur_input_ids_len);
-                    if (start_idx >= end_idx)
-                        continue;
-
-                    int64_t cur_draft_token_num = end_idx - start_idx;
-
-                    seq_lens_this_time[batch_idx] = ori_seq_len_this_time + cur_draft_token_num;
-                    memcpy(cur_draft_tokens + ori_seq_len_this_time, cur_input_ids + start_idx, sizeof(int64_t) * cur_draft_token_num);
-                    // To break the current batch_idx for-loop
-                    match_global = true;
-                    break;
-                }
-            }
-            // apply ngram_match in generated tokens
-            if (!match_global) {
-                for (int64_t i = 0; i <= cur_step_idx - ngram_size && !match_global; ++i) {
-                    // Check if the current window matches the ngram
-                    bool match_local = true;
-
-                    for (int j = 0; j < ngram_size; j++) {
-                        if (ngram[j] != cur_pre_ids[i + j]) {
-                            match_local = false;
-                            break;
-                        }
-                    }
-                    if (match_local) {
-                        int64_t start_idx = i + ngram_size;
-                        int64_t end_idx = std::min(start_idx + max_draft_tokens_query, cur_step_idx);
-
-                        int64_t cur_draft_token_num = end_idx - start_idx;
-
-                        if (start_idx >= end_idx)
-                            continue;
-                        // printf("match in Output with Ngram_size %d. %lld:[%lld,%lld]\n",ngram_size, cur_draft_token_num, start_idx, end_idx);
-
-                        seq_lens_this_time[batch_idx] = ori_seq_len_this_time + cur_draft_token_num;
-                        memcpy(cur_draft_tokens + ori_seq_len_this_time, cur_pre_ids + start_idx, sizeof(int64_t) * cur_draft_token_num);
-                        match_global = true;
-                        break;
-                    }
-                }
-            }
-        }
+    if (sum_token_num + max_draft_tokens_query + left_min_token_num >
+        threshold) {
+      int tmp_max_draft_tokens = threshold - sum_token_num - left_min_token_num;
+      max_draft_tokens_query =
+          std::min(max_draft_tokens_query, tmp_max_draft_tokens);
     }
+
+    if (sum_token_num + left_min_token_num >= threshold - 1) {
+      continue;
+    }
+    bool match_global = false;
+    // apply ngram_match in input_ids
+    for (int ngram_size = max_ngram_size;
+         ngram_size >= min_ngram_size && !match_global;
+         --ngram_size) {
+      // Extract the last n tokens as our search ngram
+      if (cur_step_idx < ngram_size) {
+        continue;
+      }
+      const int64_t *ngram = cur_pre_ids + (cur_step_idx + 1 - ngram_size);
+
+      // Iterate through sliding windows of size ngram_size
+      // bool match_input = false;
+      for (int64_t i = 0; i <= cur_input_ids_len - ngram_size && !match_global;
+           ++i) {
+        // Check if the current window matches the ngram
+        bool match_local = true;
+        for (int j = 0; j < ngram_size; j++) {
+          if (ngram[j] != cur_input_ids[i + j]) {
+            match_local = false;
+            break;
+          }
+        }
+        if (match_local) {
+          int64_t start_idx = i + ngram_size;
+          int64_t end_idx =
+              std::min(start_idx + max_draft_tokens_query, cur_input_ids_len);
+          if (start_idx >= end_idx) continue;
+
+          int64_t cur_draft_token_num = end_idx - start_idx;
+
+          seq_lens_this_time[batch_idx] =
+              ori_seq_len_this_time + cur_draft_token_num;
+          memcpy(cur_draft_tokens + ori_seq_len_this_time,
+                 cur_input_ids + start_idx,
+                 sizeof(int64_t) * cur_draft_token_num);
+          // To break the current batch_idx for-loop
+          match_global = true;
+          break;
+        }
+      }
+      // apply ngram_match in generated tokens
+      if (!match_global) {
+        for (int64_t i = 0; i <= cur_step_idx - ngram_size && !match_global;
+             ++i) {
+          // Check if the current window matches the ngram
+          bool match_local = true;
+
+          for (int j = 0; j < ngram_size; j++) {
+            if (ngram[j] != cur_pre_ids[i + j]) {
+              match_local = false;
+              break;
+            }
+          }
+          if (match_local) {
+            int64_t start_idx = i + ngram_size;
+            int64_t end_idx =
+                std::min(start_idx + max_draft_tokens_query, cur_step_idx);
+
+            int64_t cur_draft_token_num = end_idx - start_idx;
+
+            if (start_idx >= end_idx) continue;
+            // printf("match in Output with Ngram_size %d.
+            // %lld:[%lld,%lld]\n",ngram_size, cur_draft_token_num, start_idx,
+            // end_idx);
+
+            seq_lens_this_time[batch_idx] =
+                ori_seq_len_this_time + cur_draft_token_num;
+            memcpy(cur_draft_tokens + ori_seq_len_this_time,
+                   cur_pre_ids + start_idx,
+                   sizeof(int64_t) * cur_draft_token_num);
+            match_global = true;
+            break;
+          }
+        }
+      }
+    }
+  }
 }
 
 void HybridMtpNgram(const paddle::Tensor &input_ids,
-        const paddle::Tensor &input_ids_len,
-        const paddle::Tensor &pre_ids,
-        const paddle::Tensor &step_idx,
-        const paddle::Tensor &draft_token_num,
-        const paddle::Tensor &draft_tokens,
-        const paddle::Tensor &seq_lens_this_time,
-        const paddle::Tensor &seq_lens_decoder,
-        const paddle::Tensor &max_dec_len,
-        const int max_ngram_size,
-        const int min_ngram_size,
-        const int max_draft_tokens) {
+                    const paddle::Tensor &input_ids_len,
+                    const paddle::Tensor &pre_ids,
+                    const paddle::Tensor &step_idx,
+                    const paddle::Tensor &draft_token_num,
+                    const paddle::Tensor &draft_tokens,
+                    const paddle::Tensor &seq_lens_this_time,
+                    const paddle::Tensor &seq_lens_decoder,
+                    const paddle::Tensor &max_dec_len,
+                    const int max_ngram_size,
+                    const int min_ngram_size,
+                    const int max_draft_tokens) {
+  auto input_ids_shape = input_ids.shape();
+  const int64_t input_ids_stride = input_ids_shape[1];
 
-    auto input_ids_shape = input_ids.shape();
-    const int64_t input_ids_stride = input_ids_shape[1];
+  auto pre_ids_shape = pre_ids.shape();
+  const int64_t pre_ids_stride = pre_ids_shape[1];
 
-    auto pre_ids_shape = pre_ids.shape();
-    const int64_t pre_ids_stride = pre_ids_shape[1];
+  auto draft_tokens_shape = draft_tokens.shape();
+  const int64_t draft_tokens_stride = draft_tokens_shape[1];
 
-    auto draft_tokens_shape = draft_tokens.shape();
-    const int64_t draft_tokens_stride = draft_tokens_shape[1];
+  const int64_t max_batch_size = seq_lens_this_time.shape()[0];
 
-    const int64_t max_batch_size = seq_lens_this_time.shape()[0];
-
-    find_candidate_pred_tokens_mixed(input_ids.data<int64_t>(),
-            input_ids_len.data<int64_t>(),
-            pre_ids.data<int64_t>(),
-            step_idx.data<int64_t>(),
-            draft_token_num.data<int>(),
-            const_cast<int64_t *>(draft_tokens.data<int64_t>()),
-            const_cast<int32_t *>(seq_lens_this_time.data<int32_t>()),
-            const_cast<int32_t *>(seq_lens_decoder.data<int32_t>()),
-            const_cast<int64_t *>(max_dec_len.data<int64_t>()),
-            input_ids_stride,
-            pre_ids_stride,
-            draft_tokens_stride,
-            max_batch_size,
-            max_ngram_size,
-            min_ngram_size,
-            max_draft_tokens);
+  find_candidate_pred_tokens_mixed(
+      input_ids.data<int64_t>(),
+      input_ids_len.data<int64_t>(),
+      pre_ids.data<int64_t>(),
+      step_idx.data<int64_t>(),
+      draft_token_num.data<int>(),
+      const_cast<int64_t *>(draft_tokens.data<int64_t>()),
+      const_cast<int32_t *>(seq_lens_this_time.data<int32_t>()),
+      const_cast<int32_t *>(seq_lens_decoder.data<int32_t>()),
+      const_cast<int64_t *>(max_dec_len.data<int64_t>()),
+      input_ids_stride,
+      pre_ids_stride,
+      draft_tokens_stride,
+      max_batch_size,
+      max_ngram_size,
+      min_ngram_size,
+      max_draft_tokens);
 }
 
 PD_BUILD_STATIC_OP(hybrid_mtp_ngram)
-        .Inputs({"input_ids",
-                "input_ids_len",
-                "pre_ids",
-                "step_idx",
-                "draft_token_num",
-                "draft_tokens",
-                "seq_lens_this_time",
-                "seq_lens_decoder",
-                "max_dec_len"})
-        .Attrs({"max_ngram_size: int", "min_ngram_size: int", "max_draft_tokens: int"})
-        .Outputs({"draft_tokens_out", "seq_lens_this_time_out"})
-        .SetKernelFn(PD_KERNEL(HybridMtpNgram))
-        .SetInplaceMap({{"draft_tokens", "draft_tokens_out"}, {"seq_lens_this_time", "seq_lens_this_time_out"}});
+    .Inputs({"input_ids",
+             "input_ids_len",
+             "pre_ids",
+             "step_idx",
+             "draft_token_num",
+             "draft_tokens",
+             "seq_lens_this_time",
+             "seq_lens_decoder",
+             "max_dec_len"})
+    .Attrs({"max_ngram_size: int",
+            "min_ngram_size: int",
+            "max_draft_tokens: int"})
+    .Outputs({"draft_tokens_out", "seq_lens_this_time_out"})
+    .SetKernelFn(PD_KERNEL(HybridMtpNgram))
+    .SetInplaceMap({{"draft_tokens", "draft_tokens_out"},
+                    {"seq_lens_this_time", "seq_lens_this_time_out"}});
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_calcu_accept_ratio.cu b/custom_ops/gpu_ops/speculate_decoding/speculate_calcu_accept_ratio.cu
index 6d5736e05e..98b8f35a2a 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_calcu_accept_ratio.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_calcu_accept_ratio.cu
@@ -22,30 +22,29 @@ __global__ void CalculateKernel(int32_t* sum_draft_num,
                                 const int32_t* seq_lens_decoder,
                                 const bool* stop_flags,
                                 int real_bsz) {
-    int tid = threadIdx.x;
-    int draft_num = 0, accept_num = 0;
-    if (tid < real_bsz) {
-        if (seq_lens_decoder[tid] > 0 &&
-            seq_lens_this_time[tid] != seq_lens_decoder[tid]) {
-            draft_num = seq_lens_this_time[tid] - 1;
-            accept_num = accept_nums[tid] - 1;
-        } else if (seq_lens_this_time[tid] > 0 &&
-                   stop_flags[tid]) {  // last step
-            draft_num = seq_lens_this_time[tid] - 1;
-            accept_num = accept_nums[tid];
-        }
+  int tid = threadIdx.x;
+  int draft_num = 0, accept_num = 0;
+  if (tid < real_bsz) {
+    if (seq_lens_decoder[tid] > 0 &&
+        seq_lens_this_time[tid] != seq_lens_decoder[tid]) {
+      draft_num = seq_lens_this_time[tid] - 1;
+      accept_num = accept_nums[tid] - 1;
+    } else if (seq_lens_this_time[tid] > 0 && stop_flags[tid]) {  // last step
+      draft_num = seq_lens_this_time[tid] - 1;
+      accept_num = accept_nums[tid];
     }
-    __syncthreads();
-    typedef cub::BlockReduce<int64_t, THREADBLOCK_SIZE> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
+  }
+  __syncthreads();
+  typedef cub::BlockReduce<int64_t, THREADBLOCK_SIZE> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
 
-    int draft_nums_sum = BlockReduce(temp_storage).Sum(draft_num);
-    int accept_nums_sum = BlockReduce(temp_storage).Sum(accept_num);
+  int draft_nums_sum = BlockReduce(temp_storage).Sum(draft_num);
+  int accept_nums_sum = BlockReduce(temp_storage).Sum(accept_num);
 
-    if (tid == 0 && draft_nums_sum != 0) {
-        sum_draft_num[0] += draft_nums_sum;
-        sum_accept_num[0] += accept_nums_sum;
-    }
+  if (tid == 0 && draft_nums_sum != 0) {
+    sum_draft_num[0] += draft_nums_sum;
+    sum_accept_num[0] += accept_nums_sum;
+  }
 }
 
 void Calculate(const paddle::Tensor& sum_draft_num,
@@ -54,17 +53,17 @@ void Calculate(const paddle::Tensor& sum_draft_num,
                const paddle::Tensor& seq_lens_this_time,
                const paddle::Tensor& seq_lens_decoder,
                const paddle::Tensor& stop_flags) {
-    int real_bsz = seq_lens_this_time.shape()[0];
-    constexpr int BLOCK_SIZE = 512;
+  int real_bsz = seq_lens_this_time.shape()[0];
+  constexpr int BLOCK_SIZE = 512;
 
-    CalculateKernel<BLOCK_SIZE><<<1, BLOCK_SIZE, 0, accept_nums.stream()>>>(
-        const_cast<int*>(sum_draft_num.data<int32_t>()),
-        const_cast<int*>(sum_accept_num.data<int32_t>()),
-        accept_nums.data<int32_t>(),
-        seq_lens_this_time.data<int32_t>(),
-        seq_lens_decoder.data<int32_t>(),
-        stop_flags.data<bool>(),
-        real_bsz);
+  CalculateKernel<BLOCK_SIZE><<<1, BLOCK_SIZE, 0, accept_nums.stream()>>>(
+      const_cast<int*>(sum_draft_num.data<int32_t>()),
+      const_cast<int*>(sum_accept_num.data<int32_t>()),
+      accept_nums.data<int32_t>(),
+      seq_lens_this_time.data<int32_t>(),
+      seq_lens_decoder.data<int32_t>(),
+      stop_flags.data<bool>(),
+      real_bsz);
 }
 
 PD_BUILD_STATIC_OP(speculate_calcu_accept_ratio)
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_clear_accept_nums.cu b/custom_ops/gpu_ops/speculate_decoding/speculate_clear_accept_nums.cu
index b4c9b176da..ca8ae85380 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_clear_accept_nums.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_clear_accept_nums.cu
@@ -17,19 +17,19 @@
 __global__ void speculate_clear_accept_nums_kernel(int* accept_num,
                                                    const int* seq_lens_decoder,
                                                    const int max_bsz) {
-    const int bid = threadIdx.x;
-    if (bid >= max_bsz) return;
-    accept_num[bid] = seq_lens_decoder[bid] == 0 ? 0 : accept_num[bid];
+  const int bid = threadIdx.x;
+  if (bid >= max_bsz) return;
+  accept_num[bid] = seq_lens_decoder[bid] == 0 ? 0 : accept_num[bid];
 }
 
 void SpeculateClearAcceptNums(const paddle::Tensor& accept_num,
                               const paddle::Tensor& seq_lens_decoder) {
-    // printf("enter clear \n");
-    const int max_bsz = seq_lens_decoder.shape()[0];
-    speculate_clear_accept_nums_kernel<<<1, 1024, 0, accept_num.stream()>>>(
-        const_cast<int*>(accept_num.data<int>()),
-        seq_lens_decoder.data<int>(),
-        max_bsz);
+  // printf("enter clear \n");
+  const int max_bsz = seq_lens_decoder.shape()[0];
+  speculate_clear_accept_nums_kernel<<<1, 1024, 0, accept_num.stream()>>>(
+      const_cast<int*>(accept_num.data<int>()),
+      seq_lens_decoder.data<int>(),
+      max_bsz);
 }
 
 PD_BUILD_STATIC_OP(speculate_clear_accept_nums)
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_get_output.cc b/custom_ops/gpu_ops/speculate_decoding/speculate_get_output.cc
index 77dd9d607e..c2cb95e219 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_get_output.cc
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_get_output.cc
@@ -30,62 +30,58 @@ void SpeculateGetOutput(const paddle::Tensor& x,
                         bool wait_flag,
                         int msg_queue_id,
                         bool get_each_rank) {
-    if (!get_each_rank && rank_id > 0) {
-        return;
-    }
-
-    if (const char* inference_msg_queue_id_env_p =
-            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
-        std::string inference_msg_queue_id_env_str(
-            inference_msg_queue_id_env_p);
-        int inference_msg_queue_id_from_env =
-            std::stoi(inference_msg_queue_id_env_str);
-#ifdef GET_OUTPUT_DEBUG
-        std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
-                  << inference_msg_queue_id_from_env << std::endl;
-#endif
-        msg_queue_id = inference_msg_queue_id_from_env;
-    }
-
-    static struct speculate_msgdata msg_rcv;
-
-    static key_t key = ftok("./", msg_queue_id);
-
-    static int msgid = msgget(key, IPC_CREAT | 0666);
-
-    int64_t* out_data = const_cast<int64_t*>(x.data<int64_t>());
-    int ret = -1;
-    if (!wait_flag) {
-        ret = msgrcv(msgid,
-                     &msg_rcv,
-                     (MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2) * 4,
-                     0,
-                     IPC_NOWAIT);
-    } else {
-        ret = msgrcv(msgid,
-                     &msg_rcv,
-                     (MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2) * 4,
-                     0,
-                     0);
-    }
-    if (ret == -1) {
-        out_data[0] = -2;
-        out_data[1] = 0;
-        return;
-    }
-    int bsz = msg_rcv.mtext[1];
-
-    for (int64_t i = 0; i < MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2; i++) {
-        out_data[i] = (int64_t)msg_rcv.mtext[i];
-    }
+  if (!get_each_rank && rank_id > 0) {
     return;
+  }
+
+  if (const char* inference_msg_queue_id_env_p =
+          std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+    std::string inference_msg_queue_id_env_str(inference_msg_queue_id_env_p);
+    int inference_msg_queue_id_from_env =
+        std::stoi(inference_msg_queue_id_env_str);
+#ifdef GET_OUTPUT_DEBUG
+    std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+              << inference_msg_queue_id_from_env << std::endl;
+#endif
+    msg_queue_id = inference_msg_queue_id_from_env;
+  }
+
+  static struct speculate_msgdata msg_rcv;
+
+  static key_t key = ftok("./", msg_queue_id);
+
+  static int msgid = msgget(key, IPC_CREAT | 0666);
+
+  int64_t* out_data = const_cast<int64_t*>(x.data<int64_t>());
+  int ret = -1;
+  if (!wait_flag) {
+    ret = msgrcv(msgid,
+                 &msg_rcv,
+                 (MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2) * 4,
+                 0,
+                 IPC_NOWAIT);
+  } else {
+    ret = msgrcv(
+        msgid, &msg_rcv, (MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2) * 4, 0, 0);
+  }
+  if (ret == -1) {
+    out_data[0] = -2;
+    out_data[1] = 0;
+    return;
+  }
+  int bsz = msg_rcv.mtext[1];
+
+  for (int64_t i = 0; i < MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2; i++) {
+    out_data[i] = (int64_t)msg_rcv.mtext[i];
+  }
+  return;
 }
 
 void SpeculateGetOutputStatic(const paddle::Tensor& x,
                               int64_t rank_id,
                               bool wait_flag,
                               bool get_each_rank) {
-    SpeculateGetOutput(x, rank_id, wait_flag, 1, get_each_rank);
+  SpeculateGetOutput(x, rank_id, wait_flag, 1, get_each_rank);
 }
 
 void SpeculateGetOutputDynamic(const paddle::Tensor& x,
@@ -93,7 +89,7 @@ void SpeculateGetOutputDynamic(const paddle::Tensor& x,
                                bool wait_flag,
                                int msg_queue_id,
                                bool get_each_rank) {
-    SpeculateGetOutput(x, rank_id, wait_flag, msg_queue_id, get_each_rank);
+  SpeculateGetOutput(x, rank_id, wait_flag, msg_queue_id, get_each_rank);
 }
 
 PD_BUILD_STATIC_OP(speculate_get_output)
@@ -105,7 +101,10 @@ PD_BUILD_STATIC_OP(speculate_get_output)
 
 PD_BUILD_STATIC_OP(speculate_get_output_dynamic)
     .Inputs({"x"})
-    .Attrs({"rank_id: int64_t", "wait_flag: bool", "msg_queue_id: int", "get_each_rank: bool"})
+    .Attrs({"rank_id: int64_t",
+            "wait_flag: bool",
+            "msg_queue_id: int",
+            "get_each_rank: bool"})
     .Outputs({"x_out"})
     .SetInplaceMap({{"x", "x_out"}})
     .SetKernelFn(PD_KERNEL(SpeculateGetOutputDynamic));
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_get_output_with_topk.cc b/custom_ops/gpu_ops/speculate_decoding/speculate_get_output_with_topk.cc
index c87b27d1e2..d64d53f846 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_get_output_with_topk.cc
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_get_output_with_topk.cc
@@ -28,15 +28,15 @@
 #define MAX_DRAFT_TOKEN_NUM 6
 
 struct batch_msgdata {
-    int tokens[MAX_DRAFT_TOKEN_NUM * (K + 1)];
-    float scores[MAX_DRAFT_TOKEN_NUM * (K + 1)];
-    int ranks[MAX_DRAFT_TOKEN_NUM];
+  int tokens[MAX_DRAFT_TOKEN_NUM * (K + 1)];
+  float scores[MAX_DRAFT_TOKEN_NUM * (K + 1)];
+  int ranks[MAX_DRAFT_TOKEN_NUM];
 };
 
 struct msgdata {
-    long mtype;
-    int meta[3 + MAX_BSZ];  // stop_flag, message_flag, bsz, batch_token_nums
-    batch_msgdata mtext[MAX_BSZ];
+  long mtype;
+  int meta[3 + MAX_BSZ];  // stop_flag, message_flag, bsz, batch_token_nums
+  batch_msgdata mtext[MAX_BSZ];
 };
 
 void SpeculateGetOutMmsgTopK(const paddle::Tensor& output_tokens,
@@ -45,110 +45,105 @@ void SpeculateGetOutMmsgTopK(const paddle::Tensor& output_tokens,
                              int real_k,
                              int64_t rank_id,
                              bool wait_flag) {
-    struct msgdata msg_rcv;
-    int msg_queue_id = 1;
+  struct msgdata msg_rcv;
+  int msg_queue_id = 1;
 
-    if (const char* inference_msg_queue_id_env_p =
-            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
-        std::string inference_msg_queue_id_env_str(
-            inference_msg_queue_id_env_p);
-        int inference_msg_queue_id_from_env =
-            std::stoi(inference_msg_queue_id_env_str);
+  if (const char* inference_msg_queue_id_env_p =
+          std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+    std::string inference_msg_queue_id_env_str(inference_msg_queue_id_env_p);
+    int inference_msg_queue_id_from_env =
+        std::stoi(inference_msg_queue_id_env_str);
 #ifdef SPECULATE_GET_WITH_OUTPUT_DEBUG
-        std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
-                  << inference_msg_queue_id_from_env << std::endl;
+    std::cout << "Your INFERENCE_MSG_QUEUE_ID is: "
+              << inference_msg_queue_id_from_env << std::endl;
 #endif
-        msg_queue_id = inference_msg_queue_id_from_env;
-    }
-    static key_t key = ftok("/dev/shm", msg_queue_id);
+    msg_queue_id = inference_msg_queue_id_from_env;
+  }
+  static key_t key = ftok("/dev/shm", msg_queue_id);
 
-    static int msgid = msgget(key, IPC_CREAT | 0666);
+  static int msgid = msgget(key, IPC_CREAT | 0666);
 #ifdef SPECULATE_GET_WITH_OUTPUT_DEBUG
-    std::cout << "get_output_key: " << key << std::endl;
-    std::cout << "get_output msgid: " << msgid << std::endl;
+  std::cout << "get_output_key: " << key << std::endl;
+  std::cout << "get_output msgid: " << msgid << std::endl;
 #endif
 
-    int64_t* output_tokens_data =
-        const_cast<int64_t*>(output_tokens.data<int64_t>());
-    float* output_scores_data = const_cast<float*>(output_scores.data<float>());
-    int64_t* output_ranks_data =
-        const_cast<int64_t*>(output_ranks.data<int64_t>());
-    int ret = -1;
-    if (!wait_flag) {
-        ret = msgrcv(
-            msgid, &msg_rcv, sizeof(msg_rcv) - sizeof(long), 0, IPC_NOWAIT);
-    } else {
-        ret = msgrcv(msgid, &msg_rcv, sizeof(msg_rcv) - sizeof(long), 0, 0);
-    }
-    if (ret == -1) {
-        // read none
-        output_tokens_data[0] = -2;  // stop_flag
-        output_tokens_data[1] = 0;   // message_flag, Target: 3, Draft: 4
-        output_tokens_data[2] = 0;   // bsz
-        return;
-    }
-
-    int bsz = msg_rcv.meta[2];
-    output_tokens_data[0] = (int64_t)msg_rcv.meta[0];
-    output_tokens_data[1] = (int64_t)msg_rcv.meta[1];
-    output_tokens_data[2] = (int64_t)msg_rcv.meta[2];
-
-    int output_tokens_offset = 3 + MAX_BSZ;
-    for (int i = 0; i < bsz; i++) {
-        int cur_token_num = msg_rcv.meta[3 + i];
-        output_tokens_data[3 + i] = (int64_t)cur_token_num;  // batch_token_nums
-
-        auto* cur_output_token = output_tokens_data + output_tokens_offset +
-                                 i * (MAX_DRAFT_TOKEN_NUM * (K + 1));
-        auto* cur_output_score =
-            output_scores_data + i * (MAX_DRAFT_TOKEN_NUM * (K + 1));
-        auto* cur_batch_msg_rcv = &msg_rcv.mtext[i];
-        for (int j = 0; j < cur_token_num; j++) {
-            for (int k = 0; k < real_k + 1; k++) {
-                cur_output_token[j * (K + 1) + k] =
-                    (int64_t)cur_batch_msg_rcv->tokens[j * (K + 1) + k];
-                cur_output_score[j * (K + 1) + k] =
-                    cur_batch_msg_rcv->scores[j * (K + 1) + k];
-            }
-            output_ranks_data[i * MAX_DRAFT_TOKEN_NUM + j] =
-                (int64_t)cur_batch_msg_rcv->ranks[j];
-        }
-    }
-#ifdef SPECULATE_GET_WITH_OUTPUT_DEBUG
-    std::cout << "msg data: " << std::endl;
-    std::cout << "stop_flag: " << output_tokens_data[0]
-              << ", message_flag: " << output_tokens_data[1]
-              << ", bsz: " << output_tokens_data[2] << std::endl;
-    for (int i = 0; i < output_tokens_data[2]; i++) {
-        int cur_token_num = output_tokens_data[3 + i];
-        std::cout << "batch " << i << " token_num: " << cur_token_num
-                  << std::endl;
-        for (int j = 0; j < cur_token_num; j++) {
-            std::cout << "tokens: ";
-            for (int k = 0; k < K + 1; k++) {
-                std::cout
-                    << output_tokens_data[output_tokens_offset +
-                                          i * MAX_DRAFT_TOKEN_NUM * (K + 1) +
-                                          j * (K + 1) + k]
-                    << " ";
-            }
-            std::cout << std::endl;
-            std::cout << "scores: ";
-            for (int k = 0; k < K + 1; k++) {
-                std::cout
-                    << output_scores_data[i * MAX_DRAFT_TOKEN_NUM * (K + 1) +
-                                          j * (K + 1) + k]
-                    << " ";
-            }
-            std::cout << std::endl;
-            std::cout << "ranks: "
-                      << output_ranks_data[i * MAX_DRAFT_TOKEN_NUM + j]
-                      << std::endl;
-        }
-    }
-    std::cout << std::endl;
-#endif
+  int64_t* output_tokens_data =
+      const_cast<int64_t*>(output_tokens.data<int64_t>());
+  float* output_scores_data = const_cast<float*>(output_scores.data<float>());
+  int64_t* output_ranks_data =
+      const_cast<int64_t*>(output_ranks.data<int64_t>());
+  int ret = -1;
+  if (!wait_flag) {
+    ret =
+        msgrcv(msgid, &msg_rcv, sizeof(msg_rcv) - sizeof(long), 0, IPC_NOWAIT);
+  } else {
+    ret = msgrcv(msgid, &msg_rcv, sizeof(msg_rcv) - sizeof(long), 0, 0);
+  }
+  if (ret == -1) {
+    // read none
+    output_tokens_data[0] = -2;  // stop_flag
+    output_tokens_data[1] = 0;   // message_flag, Target: 3, Draft: 4
+    output_tokens_data[2] = 0;   // bsz
     return;
+  }
+
+  int bsz = msg_rcv.meta[2];
+  output_tokens_data[0] = (int64_t)msg_rcv.meta[0];
+  output_tokens_data[1] = (int64_t)msg_rcv.meta[1];
+  output_tokens_data[2] = (int64_t)msg_rcv.meta[2];
+
+  int output_tokens_offset = 3 + MAX_BSZ;
+  for (int i = 0; i < bsz; i++) {
+    int cur_token_num = msg_rcv.meta[3 + i];
+    output_tokens_data[3 + i] = (int64_t)cur_token_num;  // batch_token_nums
+
+    auto* cur_output_token = output_tokens_data + output_tokens_offset +
+                             i * (MAX_DRAFT_TOKEN_NUM * (K + 1));
+    auto* cur_output_score =
+        output_scores_data + i * (MAX_DRAFT_TOKEN_NUM * (K + 1));
+    auto* cur_batch_msg_rcv = &msg_rcv.mtext[i];
+    for (int j = 0; j < cur_token_num; j++) {
+      for (int k = 0; k < real_k + 1; k++) {
+        cur_output_token[j * (K + 1) + k] =
+            (int64_t)cur_batch_msg_rcv->tokens[j * (K + 1) + k];
+        cur_output_score[j * (K + 1) + k] =
+            cur_batch_msg_rcv->scores[j * (K + 1) + k];
+      }
+      output_ranks_data[i * MAX_DRAFT_TOKEN_NUM + j] =
+          (int64_t)cur_batch_msg_rcv->ranks[j];
+    }
+  }
+#ifdef SPECULATE_GET_WITH_OUTPUT_DEBUG
+  std::cout << "msg data: " << std::endl;
+  std::cout << "stop_flag: " << output_tokens_data[0]
+            << ", message_flag: " << output_tokens_data[1]
+            << ", bsz: " << output_tokens_data[2] << std::endl;
+  for (int i = 0; i < output_tokens_data[2]; i++) {
+    int cur_token_num = output_tokens_data[3 + i];
+    std::cout << "batch " << i << " token_num: " << cur_token_num << std::endl;
+    for (int j = 0; j < cur_token_num; j++) {
+      std::cout << "tokens: ";
+      for (int k = 0; k < K + 1; k++) {
+        std::cout << output_tokens_data[output_tokens_offset +
+                                        i * MAX_DRAFT_TOKEN_NUM * (K + 1) +
+                                        j * (K + 1) + k]
+                  << " ";
+      }
+      std::cout << std::endl;
+      std::cout << "scores: ";
+      for (int k = 0; k < K + 1; k++) {
+        std::cout << output_scores_data[i * MAX_DRAFT_TOKEN_NUM * (K + 1) +
+                                        j * (K + 1) + k]
+                  << " ";
+      }
+      std::cout << std::endl;
+      std::cout << "ranks: " << output_ranks_data[i * MAX_DRAFT_TOKEN_NUM + j]
+                << std::endl;
+    }
+  }
+  std::cout << std::endl;
+#endif
+  return;
 }
 
 PD_BUILD_STATIC_OP(speculate_get_output_topk)
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_get_seq_lens_output.cu b/custom_ops/gpu_ops/speculate_decoding/speculate_get_seq_lens_output.cu
index 9da1793192..672fdab793 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_get_seq_lens_output.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_get_seq_lens_output.cu
@@ -23,53 +23,53 @@ __global__ void SpeculateGetSeqLensOutputKernel(int* seq_lens_output,
                                                 const int* seq_lens_encoder,
                                                 const int* seq_lens_decoder,
                                                 const int real_bsz) {
-    for (int bid = threadIdx.x; bid < real_bsz; bid += blockDim.x) {
-        if (seq_lens_this_time[bid] == 0) {
-            continue;
-        } else if (seq_lens_this_time[bid] == 1) {
-            seq_lens_output[bid] = 1;
-        } else if (seq_lens_encoder[bid] != 0) {
-            seq_lens_output[bid] = 1;
-        } else {
-            seq_lens_output[bid] = seq_lens_this_time[bid];
-        }
+  for (int bid = threadIdx.x; bid < real_bsz; bid += blockDim.x) {
+    if (seq_lens_this_time[bid] == 0) {
+      continue;
+    } else if (seq_lens_this_time[bid] == 1) {
+      seq_lens_output[bid] = 1;
+    } else if (seq_lens_encoder[bid] != 0) {
+      seq_lens_output[bid] = 1;
+    } else {
+      seq_lens_output[bid] = seq_lens_this_time[bid];
     }
+  }
 }
 
 std::vector<paddle::Tensor> SpeculateGetSeqLensOutput(
     const paddle::Tensor& seq_lens_this_time,
     const paddle::Tensor& seq_lens_encoder,
     const paddle::Tensor& seq_lens_decoder) {
-    auto cu_stream = seq_lens_this_time.stream();
-    std::vector<int64_t> seq_lens_this_time_shape = seq_lens_this_time.shape();
-    const int bsz = seq_lens_this_time_shape[0];
+  auto cu_stream = seq_lens_this_time.stream();
+  std::vector<int64_t> seq_lens_this_time_shape = seq_lens_this_time.shape();
+  const int bsz = seq_lens_this_time_shape[0];
 
-    auto seq_lens_output = paddle::full(
-        {bsz}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto seq_lens_output = paddle::full(
+      {bsz}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
 
-    SpeculateGetSeqLensOutputKernel<<<1, 256, 0, cu_stream>>>(
-        seq_lens_output.data<int>(),
-        seq_lens_this_time.data<int>(),
-        seq_lens_encoder.data<int>(),
-        seq_lens_decoder.data<int>(),
-        bsz);
+  SpeculateGetSeqLensOutputKernel<<<1, 256, 0, cu_stream>>>(
+      seq_lens_output.data<int>(),
+      seq_lens_this_time.data<int>(),
+      seq_lens_encoder.data<int>(),
+      seq_lens_decoder.data<int>(),
+      bsz);
 
-    return {seq_lens_output};
+  return {seq_lens_output};
 }
 
 std::vector<std::vector<int64_t>> SpeculateGetSeqLensOutputInferShape(
     const std::vector<int64_t>& seq_lens_this_time_shape,
     const std::vector<int64_t>& seq_lens_encoder_shape,
     const std::vector<int64_t>& seq_lens_decoder_shape) {
-    int64_t bsz = seq_lens_this_time_shape[0];
-    return {{bsz}};
+  int64_t bsz = seq_lens_this_time_shape[0];
+  return {{bsz}};
 }
 
 std::vector<paddle::DataType> SpeculateGetSeqLensOutputInferDtype(
     const paddle::DataType& seq_lens_this_time_dtype,
     const paddle::DataType& seq_lens_encoder_dtype,
     const paddle::DataType& seq_lens_decoder_dtype) {
-    return {seq_lens_this_time_dtype};
+  return {seq_lens_this_time_dtype};
 }
 
 PD_BUILD_STATIC_OP(speculate_get_seq_lens_output)
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_logprob_utils.cu b/custom_ops/gpu_ops/speculate_decoding/speculate_logprob_utils.cu
index a20773bf57..37dabeae07 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_logprob_utils.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_logprob_utils.cu
@@ -19,13 +19,13 @@ __global__ void get_token_num_per_batch_kernel(int* next_token_num,
                                                const int* seq_lens_this_time,
                                                const int* seq_lens_encoder,
                                                const int real_bsz) {
-    int bid = threadIdx.x;
-    if (bid < real_bsz) {
-        next_token_num[bid] =
-            seq_lens_encoder[bid] > 0 ? 1 : seq_lens_this_time[bid];
-        batch_token_num[bid] =
-            seq_lens_encoder[bid] > 0 ? 2 : seq_lens_this_time[bid];
-    }
+  int bid = threadIdx.x;
+  if (bid < real_bsz) {
+    next_token_num[bid] =
+        seq_lens_encoder[bid] > 0 ? 1 : seq_lens_this_time[bid];
+    batch_token_num[bid] =
+        seq_lens_encoder[bid] > 0 ? 2 : seq_lens_this_time[bid];
+  }
 }
 
 template <int VecSize>
@@ -38,32 +38,29 @@ __global__ void speculate_get_logits_kernel(float* draft_logits,
                                             const int* seq_lens_encoder,
                                             const int vocab_size,
                                             const int real_bsz) {
-    AlignedVector<float, VecSize> src_vec;
-    const int bid = blockIdx.x;
-    const int tid = threadIdx.x;
-    if (bid < real_bsz) {
-        auto* draft_logits_now =
-            draft_logits + cu_batch_token_offset[bid] * vocab_size;
-        auto* logits_now = logits + cu_next_token_offset[bid] * vocab_size;
-        for (int i = tid * VecSize; i < vocab_size; i += blockDim.x * VecSize) {
-            if (seq_lens_encoder[bid] > 0) {
-                Load<float, VecSize>(&first_token_logits[bid * vocab_size + i],
-                                     &src_vec);
-                Store<float, VecSize>(src_vec, &draft_logits_now[i]);
+  AlignedVector<float, VecSize> src_vec;
+  const int bid = blockIdx.x;
+  const int tid = threadIdx.x;
+  if (bid < real_bsz) {
+    auto* draft_logits_now =
+        draft_logits + cu_batch_token_offset[bid] * vocab_size;
+    auto* logits_now = logits + cu_next_token_offset[bid] * vocab_size;
+    for (int i = tid * VecSize; i < vocab_size; i += blockDim.x * VecSize) {
+      if (seq_lens_encoder[bid] > 0) {
+        Load<float, VecSize>(&first_token_logits[bid * vocab_size + i],
+                             &src_vec);
+        Store<float, VecSize>(src_vec, &draft_logits_now[i]);
 
-                Load<float, VecSize>(&logits_now[i], &src_vec);
-                Store<float, VecSize>(src_vec,
-                                      &draft_logits_now[vocab_size + i]);
-            } else {
-                for (int j = 0; j < seq_lens_this_time[bid]; j++) {
-                    Load<float, VecSize>(&logits_now[j * vocab_size + i],
-                                         &src_vec);
-                    Store<float, VecSize>(
-                        src_vec, &draft_logits_now[j * vocab_size + i]);
-                }
-            }
+        Load<float, VecSize>(&logits_now[i], &src_vec);
+        Store<float, VecSize>(src_vec, &draft_logits_now[vocab_size + i]);
+      } else {
+        for (int j = 0; j < seq_lens_this_time[bid]; j++) {
+          Load<float, VecSize>(&logits_now[j * vocab_size + i], &src_vec);
+          Store<float, VecSize>(src_vec, &draft_logits_now[j * vocab_size + i]);
         }
+      }
     }
+  }
 }
 
 void SpeculateGetLogits(const paddle::Tensor& draft_logits,
@@ -75,67 +72,66 @@ void SpeculateGetLogits(const paddle::Tensor& draft_logits,
                         const paddle::Tensor& first_token_logits,
                         const paddle::Tensor& seq_lens_this_time,
                         const paddle::Tensor& seq_lens_encoder) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int vocab_size = logits.shape()[1];
-    const int real_bsz = seq_lens_this_time.shape()[0];
+  auto cu_stream = seq_lens_this_time.stream();
+  const int vocab_size = logits.shape()[1];
+  const int real_bsz = seq_lens_this_time.shape()[0];
 
-    get_token_num_per_batch_kernel<<<1, 512, 0, cu_stream>>>(
-        const_cast<int*>(next_token_num.data<int>()),
-        const_cast<int*>(batch_token_num.data<int>()),
-        seq_lens_this_time.data<int>(),
-        seq_lens_encoder.data<int>(),
-        real_bsz);
+  get_token_num_per_batch_kernel<<<1, 512, 0, cu_stream>>>(
+      const_cast<int*>(next_token_num.data<int>()),
+      const_cast<int*>(batch_token_num.data<int>()),
+      seq_lens_this_time.data<int>(),
+      seq_lens_encoder.data<int>(),
+      real_bsz);
 
-    void* temp_storage1 = nullptr;
-    size_t temp_storage_bytes1 = 0;
-    cub::DeviceScan::InclusiveSum(
-        temp_storage1,
-        temp_storage_bytes1,
-        batch_token_num.data<int>(),
-        const_cast<int*>(&cu_batch_token_offset.data<int>()[1]),
-        real_bsz,
-        cu_stream);
-    cudaMalloc(&temp_storage1, temp_storage_bytes1);
-    cub::DeviceScan::InclusiveSum(
-        temp_storage1,
-        temp_storage_bytes1,
-        batch_token_num.data<int>(),
-        const_cast<int*>(&cu_batch_token_offset.data<int>()[1]),
-        real_bsz,
-        cu_stream);
+  void* temp_storage1 = nullptr;
+  size_t temp_storage_bytes1 = 0;
+  cub::DeviceScan::InclusiveSum(
+      temp_storage1,
+      temp_storage_bytes1,
+      batch_token_num.data<int>(),
+      const_cast<int*>(&cu_batch_token_offset.data<int>()[1]),
+      real_bsz,
+      cu_stream);
+  cudaMalloc(&temp_storage1, temp_storage_bytes1);
+  cub::DeviceScan::InclusiveSum(
+      temp_storage1,
+      temp_storage_bytes1,
+      batch_token_num.data<int>(),
+      const_cast<int*>(&cu_batch_token_offset.data<int>()[1]),
+      real_bsz,
+      cu_stream);
 
-    void* temp_storage2 = nullptr;
-    size_t temp_storage_bytes2 = 0;
-    cub::DeviceScan::InclusiveSum(
-        temp_storage2,
-        temp_storage_bytes2,
-        next_token_num.data<int>(),
-        const_cast<int*>(&cu_next_token_offset.data<int>()[1]),
-        real_bsz,
-        cu_stream);
-    cudaMalloc(&temp_storage2, temp_storage_bytes2);
-    cub::DeviceScan::InclusiveSum(
-        temp_storage2,
-        temp_storage_bytes2,
-        next_token_num.data<int>(),
-        const_cast<int*>(&cu_next_token_offset.data<int>()[1]),
-        real_bsz,
-        cu_stream);
+  void* temp_storage2 = nullptr;
+  size_t temp_storage_bytes2 = 0;
+  cub::DeviceScan::InclusiveSum(
+      temp_storage2,
+      temp_storage_bytes2,
+      next_token_num.data<int>(),
+      const_cast<int*>(&cu_next_token_offset.data<int>()[1]),
+      real_bsz,
+      cu_stream);
+  cudaMalloc(&temp_storage2, temp_storage_bytes2);
+  cub::DeviceScan::InclusiveSum(
+      temp_storage2,
+      temp_storage_bytes2,
+      next_token_num.data<int>(),
+      const_cast<int*>(&cu_next_token_offset.data<int>()[1]),
+      real_bsz,
+      cu_stream);
 
-    constexpr int PackSize = VEC_16B / sizeof(float);
-    dim3 grid_dim(real_bsz);
-    dim3 block_dim(128);
-    speculate_get_logits_kernel<PackSize>
-        <<<grid_dim, block_dim, 0, cu_stream>>>(
-            const_cast<float*>(draft_logits.data<float>()),
-            logits.data<float>(),
-            first_token_logits.data<float>(),
-            cu_next_token_offset.data<int>(),
-            cu_batch_token_offset.data<int>(),
-            seq_lens_this_time.data<int>(),
-            seq_lens_encoder.data<int>(),
-            vocab_size,
-            real_bsz);
+  constexpr int PackSize = VEC_16B / sizeof(float);
+  dim3 grid_dim(real_bsz);
+  dim3 block_dim(128);
+  speculate_get_logits_kernel<PackSize><<<grid_dim, block_dim, 0, cu_stream>>>(
+      const_cast<float*>(draft_logits.data<float>()),
+      logits.data<float>(),
+      first_token_logits.data<float>(),
+      cu_next_token_offset.data<int>(),
+      cu_batch_token_offset.data<int>(),
+      seq_lens_this_time.data<int>(),
+      seq_lens_encoder.data<int>(),
+      vocab_size,
+      real_bsz);
 }
 
 __global__ void speculate_insert_first_token_kernel(
@@ -148,19 +144,19 @@ __global__ void speculate_insert_first_token_kernel(
     const int* seq_lens_encoder,
     const int max_draft_tokens,
     const int real_bsz) {
-    const int bid = threadIdx.x;
+  const int bid = threadIdx.x;
 
-    auto* token_ids_now = token_ids + cu_batch_token_offset[bid];
-    auto* accept_tokens_now = accept_tokens + bid * max_draft_tokens;
-    auto* next_tokens_now = next_tokens + cu_next_token_offset[bid];
-    if (seq_lens_encoder[bid] != 0) {
-        token_ids_now[0] = accept_tokens_now[0];
-        token_ids_now[1] = next_tokens_now[0];
-    } else {
-        for (int i = 0; i < seq_lens_this_time[bid]; i++) {
-            token_ids_now[i] = next_tokens_now[i];
-        }
+  auto* token_ids_now = token_ids + cu_batch_token_offset[bid];
+  auto* accept_tokens_now = accept_tokens + bid * max_draft_tokens;
+  auto* next_tokens_now = next_tokens + cu_next_token_offset[bid];
+  if (seq_lens_encoder[bid] != 0) {
+    token_ids_now[0] = accept_tokens_now[0];
+    token_ids_now[1] = next_tokens_now[0];
+  } else {
+    for (int i = 0; i < seq_lens_this_time[bid]; i++) {
+      token_ids_now[i] = next_tokens_now[i];
     }
+  }
 }
 
 void SpeculateInsertFirstToken(const paddle::Tensor& token_ids,
@@ -170,20 +166,20 @@ void SpeculateInsertFirstToken(const paddle::Tensor& token_ids,
                                const paddle::Tensor& cu_batch_token_offset,
                                const paddle::Tensor& seq_lens_this_time,
                                const paddle::Tensor& seq_lens_encoder) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int max_draft_tokens = accept_tokens.shape()[1];
-    const int real_bsz = seq_lens_this_time.shape()[0];
+  auto cu_stream = seq_lens_this_time.stream();
+  const int max_draft_tokens = accept_tokens.shape()[1];
+  const int real_bsz = seq_lens_this_time.shape()[0];
 
-    speculate_insert_first_token_kernel<<<1, real_bsz, 0, cu_stream>>>(
-        const_cast<int64_t*>(token_ids.data<int64_t>()),
-        accept_tokens.data<int64_t>(),
-        next_tokens.data<int64_t>(),
-        cu_next_token_offset.data<int>(),
-        cu_batch_token_offset.data<int>(),
-        seq_lens_this_time.data<int>(),
-        seq_lens_encoder.data<int>(),
-        max_draft_tokens,
-        real_bsz);
+  speculate_insert_first_token_kernel<<<1, real_bsz, 0, cu_stream>>>(
+      const_cast<int64_t*>(token_ids.data<int64_t>()),
+      accept_tokens.data<int64_t>(),
+      next_tokens.data<int64_t>(),
+      cu_next_token_offset.data<int>(),
+      cu_batch_token_offset.data<int>(),
+      seq_lens_this_time.data<int>(),
+      seq_lens_encoder.data<int>(),
+      max_draft_tokens,
+      real_bsz);
 }
 
 template <int VecSize>
@@ -197,27 +193,26 @@ __global__ void speculate_get_target_logits_kernel(
     const int* accept_num,
     const int vocab_size,
     const int real_bsz) {
-    AlignedVector<float, VecSize> src_vec;
-    const int bid = blockIdx.x;
-    const int tid = threadIdx.x;
-    if (bid < real_bsz) {
-        auto* target_logtis_now =
-            target_logtis + cu_batch_token_offset[bid] * vocab_size;
-        auto* logits_now = logits + ori_cu_batch_token_offset[bid] * vocab_size;
-        for (int i = tid * VecSize; i < vocab_size; i += blockDim.x * VecSize) {
-            if (seq_lens_encoder[bid] > 0) {
-                Load<float, VecSize>(&logits_now[i], &src_vec);
-                Store<float, VecSize>(src_vec, &target_logtis_now[i]);
-            } else {
-                for (int j = 0; j < accept_num[bid]; j++) {
-                    Load<float, VecSize>(&logits_now[j * vocab_size + i],
-                                         &src_vec);
-                    Store<float, VecSize>(
-                        src_vec, &target_logtis_now[j * vocab_size + i]);
-                }
-            }
+  AlignedVector<float, VecSize> src_vec;
+  const int bid = blockIdx.x;
+  const int tid = threadIdx.x;
+  if (bid < real_bsz) {
+    auto* target_logtis_now =
+        target_logtis + cu_batch_token_offset[bid] * vocab_size;
+    auto* logits_now = logits + ori_cu_batch_token_offset[bid] * vocab_size;
+    for (int i = tid * VecSize; i < vocab_size; i += blockDim.x * VecSize) {
+      if (seq_lens_encoder[bid] > 0) {
+        Load<float, VecSize>(&logits_now[i], &src_vec);
+        Store<float, VecSize>(src_vec, &target_logtis_now[i]);
+      } else {
+        for (int j = 0; j < accept_num[bid]; j++) {
+          Load<float, VecSize>(&logits_now[j * vocab_size + i], &src_vec);
+          Store<float, VecSize>(src_vec,
+                                &target_logtis_now[j * vocab_size + i]);
         }
+      }
     }
+  }
 }
 
 void SpeculateGetTargetLogits(const paddle::Tensor& target_logits,
@@ -227,24 +222,24 @@ void SpeculateGetTargetLogits(const paddle::Tensor& target_logits,
                               const paddle::Tensor& seq_lens_this_time,
                               const paddle::Tensor& seq_lens_encoder,
                               const paddle::Tensor& accept_num) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int vocab_size = logits.shape()[1];
-    const int real_bsz = seq_lens_this_time.shape()[0];
+  auto cu_stream = seq_lens_this_time.stream();
+  const int vocab_size = logits.shape()[1];
+  const int real_bsz = seq_lens_this_time.shape()[0];
 
-    constexpr int PackSize = VEC_16B / sizeof(float);
-    dim3 grid_dim(real_bsz);
-    dim3 block_dim(128);
-    speculate_get_target_logits_kernel<PackSize>
-        <<<grid_dim, block_dim, 0, cu_stream>>>(
-            const_cast<float*>(target_logits.data<float>()),
-            logits.data<float>(),
-            cu_batch_token_offset.data<int>(),
-            ori_cu_batch_token_offset.data<int>(),
-            seq_lens_this_time.data<int>(),
-            seq_lens_encoder.data<int>(),
-            accept_num.data<int>(),
-            vocab_size,
-            real_bsz);
+  constexpr int PackSize = VEC_16B / sizeof(float);
+  dim3 grid_dim(real_bsz);
+  dim3 block_dim(128);
+  speculate_get_target_logits_kernel<PackSize>
+      <<<grid_dim, block_dim, 0, cu_stream>>>(
+          const_cast<float*>(target_logits.data<float>()),
+          logits.data<float>(),
+          cu_batch_token_offset.data<int>(),
+          ori_cu_batch_token_offset.data<int>(),
+          seq_lens_this_time.data<int>(),
+          seq_lens_encoder.data<int>(),
+          accept_num.data<int>(),
+          vocab_size,
+          real_bsz);
 }
 
 PD_BUILD_STATIC_OP(speculate_get_logits)
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_msg.h b/custom_ops/gpu_ops/speculate_decoding/speculate_msg.h
index dbc177b4cb..a212bb9c1b 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_msg.h
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_msg.h
@@ -26,7 +26,7 @@
 
 // TODO: replace all msgdata in speculate-decoding
 struct speculate_msgdata {
-    long mtype;
-    int mtext[MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ +
-              2];  // stop_flag, bsz, tokens
+  long mtype;
+  int mtext[MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ +
+            2];  // stop_flag, bsz, tokens
 };
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_step.cu b/custom_ops/gpu_ops/speculate_decoding/speculate_step.cu
index d8d4dc3f55..d797d196ba 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_step.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_step.cu
@@ -19,15 +19,15 @@
 __device__ bool speculate_free_and_dispatch_block(const int &qid,
                                                   int *need_block_list,
                                                   const int &need_block_len) {
-    bool res = false;
-    for (int i = 0; i < need_block_len; i++) {
-        if (qid == need_block_list[i]) {
-            res = true;
-            need_block_list[i] = -1;
-            break;
-        }
+  bool res = false;
+  for (int i = 0; i < need_block_len; i++) {
+    if (qid == need_block_list[i]) {
+      res = true;
+      need_block_list[i] = -1;
+      break;
     }
-    return res;
+  }
+  return res;
 }
 
 __global__ void speculate_free_and_dispatch_block(
@@ -53,185 +53,184 @@ __global__ void speculate_free_and_dispatch_block(
     const int block_num_per_seq,
     const int max_decoder_block_num,
     const int max_draft_tokens) {
-    typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
-    __shared__ bool step_max_block_flag;
-    __shared__ int in_need_block_list_len;
-    const int tid = threadIdx.x;
-    if (tid < bsz) {
-        if (tid == 0) {
-            step_max_block_flag = false;
-            in_need_block_list_len = 0;
-        }
-        int *block_table_now = block_tables + tid * block_num_per_seq;
-        int max_possible_block_idx = (seq_lens_decoder[tid] + max_draft_tokens + 1 ) / block_size;
-        if (stop_flags[tid] && !is_block_step[tid]) {
-            // 回收block块
-            first_token_ids[tid] = -1;
-            const int encoder_block_len = encoder_block_lens[tid];
-            const int decoder_used_len = used_list_len[tid];
-            if (decoder_used_len > 0) {
-                const int ori_free_list_len =
-                    atomicAdd(free_list_len, decoder_used_len);
-#ifdef DEBUG_STEP
-                printf(
-                    "free block seq_id: %d, free block num: %d, "
-                    "encoder_block_len: %d, ori_free_list_len: %d\n",
-                    tid,
-                    decoder_used_len,
-                    encoder_block_len,
-                    ori_free_list_len);
-#endif
-                for (int i = 0; i < decoder_used_len; i++) {
-                    free_list[ori_free_list_len + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                encoder_block_lens[tid] = 0;
-                used_list_len[tid] = 0;
-            }
-        } else if (seq_lens_this_time[tid] != 0 && max_possible_block_idx < block_num_per_seq &&
-                   block_table_now[(seq_lens_decoder[tid] + max_draft_tokens +
-                                    1) /
-                                   block_size] == -1) {
-            // 统计需要分配block的位置和总数
-            const int ori_need_block_len = atomicAdd(need_block_len, 1);
-            need_block_list[ori_need_block_len] = tid;
-#ifdef DEBUG_STEP
-            printf("seq_id: %d need block\n", tid);
-#endif
-        }
-    }
-    __syncthreads();
-
-    while (need_block_len[0] > free_list_len[0]) {
-#ifdef DEBUG_STEP
-        if (tid == 0) {
-            printf("need_block_len: %d, free_list_len: %d\n",
-                   need_block_len[0],
-                   free_list_len[0]);
-        }
-#endif
-        // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
-        const int used_block_num =
-            tid < bsz && !is_block_step[tid] &&
-                    (step_max_block_flag ||
-                     used_list_len[tid] != max_decoder_block_num)
-                ? used_list_len[tid]
-                : 0;
-        cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
-        kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
-
-        if (tid == 0) {
-            if (kv_pair.value == 0) {
-                step_max_block_flag = true;
-            } else {
-                const int encoder_block_len = encoder_block_lens[kv_pair.key];
-                // #ifdef DEBUG_STEP
-                printf("max_id: %d, max_num: %d, encoder_block_len: %d\n",
-                       kv_pair.key,
-                       kv_pair.value,
-                       encoder_block_len);
-                // #endif
-                int *block_table_now =
-                    block_tables + kv_pair.key * block_num_per_seq;
-                for (int i = 0; i < kv_pair.value; i++) {
-                    free_list[free_list_len[0] + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                step_block_list[step_len[0]] = kv_pair.key;
-                if (speculate_free_and_dispatch_block(
-                        kv_pair.key,
-                        need_block_list,
-                        need_block_len[0] + in_need_block_list_len)) {
-                    need_block_len[0] -= 1;
-                    in_need_block_list_len += 1;
-                }
-                step_len[0] += 1;
-                free_list_len[0] += kv_pair.value;
-                stop_flags[kv_pair.key] = true;
-                is_block_step[kv_pair.key] = true;
-                seq_lens_this_time[kv_pair.key] = 0;
-                seq_lens_decoder[kv_pair.key] = 0;
-                // Note(@wufeisheng): when step, accept num will not be 0 so
-                // that next step even if this batch member is stepped, save
-                // output still stream output, so accept num should be set to 0
-                accept_num[kv_pair.key] = 0;
-            }
-        }
-        __syncthreads();
-    }
-
-    // 为需要block的位置分配block，每个位置分配一个block
-    if (tid < need_block_len[0] + in_need_block_list_len) {
-        const int need_block_id = need_block_list[tid];
-        if (need_block_id != -1) {
-            if (!stop_flags[need_block_id]) {
-                // 如果需要的位置正好是上一步中被释放的位置，不做处理
-                used_list_len[need_block_id] += 1;
-                const int ori_free_list_len = atomicSub(free_list_len, 1);
-                int *block_table_now =
-                    block_tables + need_block_id * block_num_per_seq;
-#ifdef DEBUG_STEP
-                printf("need_block_id %d\n", need_block_id);
-                printf("ori_free_list_len %d\n", ori_free_list_len);
-                printf("max_draft_tokens %d\n", max_draft_tokens);
-                printf("seq_lens_decoder[need_block_id] %d\n",
-                       seq_lens_decoder[need_block_id]);
-                printf("free_list[ori_free_list_len - 1] %d\n",
-                       free_list[ori_free_list_len - 1]);
-#endif
-                block_table_now[(seq_lens_decoder[need_block_id] +
-                                 max_draft_tokens + 1) /
-                                block_size] = free_list[ori_free_list_len - 1];
-            }
-            need_block_list[tid] = -1;
-        }
-    }
-    __syncthreads();
-
-    // 计算可以复原的query id
+  typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ bool step_max_block_flag;
+  __shared__ int in_need_block_list_len;
+  const int tid = threadIdx.x;
+  if (tid < bsz) {
     if (tid == 0) {
-        int ori_free_list_len = free_list_len[0];
-        int ori_step_len = step_len[0];
-        if (ori_step_len > 0) {
-            int ori_step_block_id = step_block_list[ori_step_len - 1];
-            int tmp_used_len = used_list_len[ori_step_block_id];
-            // 比之前调度时多分配一个block，防止马上恢复刚调度的query(比如回收的seq_id在need_block_list中）
-            const int max_decoder_block_num_this_seq =
-                max_decoder_block_num - encoder_block_lens[ori_step_block_id];
-            int used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
-                               ? tmp_used_len + 1
-                               : max_decoder_block_num_this_seq;
-            if (ori_step_len > 0 && ori_free_list_len >= used_len) {
-                // #ifdef DEBUG_STEP
-                printf(
-                    "recover seq_id: %d, free_list_len: %d, used_list_len: "
-                    "%d\n",
-                    ori_step_block_id,
-                    ori_free_list_len,
-                    used_len);
-                // #endif
-                recover_block_list[recover_len[0]] = ori_step_block_id;
-                is_block_step[ori_step_block_id] = false;
-                used_list_len[ori_step_block_id] = used_len;
-                ori_free_list_len -= used_len;
-                step_block_list[ori_step_len - 1] = -1;
-                step_len[0] -= 1;
-                recover_len[0] += 1;
-                ori_step_len = step_len[0];
-                if (ori_step_len > 0) {
-                    ori_step_block_id = step_block_list[ori_step_len - 1];
-                    tmp_used_len = used_list_len[ori_step_block_id];
-                    used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
-                                   ? tmp_used_len + 1
-                                   : max_decoder_block_num_this_seq;
-                }
-            }
-        }
-        need_block_len[0] = 0;
+      step_max_block_flag = false;
+      in_need_block_list_len = 0;
     }
+    int *block_table_now = block_tables + tid * block_num_per_seq;
+    int max_possible_block_idx =
+        (seq_lens_decoder[tid] + max_draft_tokens + 1) / block_size;
+    if (stop_flags[tid] && !is_block_step[tid]) {
+      // 回收block块
+      first_token_ids[tid] = -1;
+      const int encoder_block_len = encoder_block_lens[tid];
+      const int decoder_used_len = used_list_len[tid];
+      if (decoder_used_len > 0) {
+        const int ori_free_list_len =
+            atomicAdd(free_list_len, decoder_used_len);
+#ifdef DEBUG_STEP
+        printf(
+            "free block seq_id: %d, free block num: %d, "
+            "encoder_block_len: %d, ori_free_list_len: %d\n",
+            tid,
+            decoder_used_len,
+            encoder_block_len,
+            ori_free_list_len);
+#endif
+        for (int i = 0; i < decoder_used_len; i++) {
+          free_list[ori_free_list_len + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        encoder_block_lens[tid] = 0;
+        used_list_len[tid] = 0;
+      }
+    } else if (seq_lens_this_time[tid] != 0 &&
+               max_possible_block_idx < block_num_per_seq &&
+               block_table_now[(seq_lens_decoder[tid] + max_draft_tokens + 1) /
+                               block_size] == -1) {
+      // 统计需要分配block的位置和总数
+      const int ori_need_block_len = atomicAdd(need_block_len, 1);
+      need_block_list[ori_need_block_len] = tid;
+#ifdef DEBUG_STEP
+      printf("seq_id: %d need block\n", tid);
+#endif
+    }
+  }
+  __syncthreads();
+
+  while (need_block_len[0] > free_list_len[0]) {
+#ifdef DEBUG_STEP
+    if (tid == 0) {
+      printf("need_block_len: %d, free_list_len: %d\n",
+             need_block_len[0],
+             free_list_len[0]);
+    }
+#endif
+    // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
+    const int used_block_num =
+        tid < bsz && !is_block_step[tid] &&
+                (step_max_block_flag ||
+                 used_list_len[tid] != max_decoder_block_num)
+            ? used_list_len[tid]
+            : 0;
+    cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
+    kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
+
+    if (tid == 0) {
+      if (kv_pair.value == 0) {
+        step_max_block_flag = true;
+      } else {
+        const int encoder_block_len = encoder_block_lens[kv_pair.key];
+        // #ifdef DEBUG_STEP
+        printf("max_id: %d, max_num: %d, encoder_block_len: %d\n",
+               kv_pair.key,
+               kv_pair.value,
+               encoder_block_len);
+        // #endif
+        int *block_table_now = block_tables + kv_pair.key * block_num_per_seq;
+        for (int i = 0; i < kv_pair.value; i++) {
+          free_list[free_list_len[0] + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        step_block_list[step_len[0]] = kv_pair.key;
+        if (speculate_free_and_dispatch_block(
+                kv_pair.key,
+                need_block_list,
+                need_block_len[0] + in_need_block_list_len)) {
+          need_block_len[0] -= 1;
+          in_need_block_list_len += 1;
+        }
+        step_len[0] += 1;
+        free_list_len[0] += kv_pair.value;
+        stop_flags[kv_pair.key] = true;
+        is_block_step[kv_pair.key] = true;
+        seq_lens_this_time[kv_pair.key] = 0;
+        seq_lens_decoder[kv_pair.key] = 0;
+        // Note(@wufeisheng): when step, accept num will not be 0 so
+        // that next step even if this batch member is stepped, save
+        // output still stream output, so accept num should be set to 0
+        accept_num[kv_pair.key] = 0;
+      }
+    }
+    __syncthreads();
+  }
+
+  // 为需要block的位置分配block，每个位置分配一个block
+  if (tid < need_block_len[0] + in_need_block_list_len) {
+    const int need_block_id = need_block_list[tid];
+    if (need_block_id != -1) {
+      if (!stop_flags[need_block_id]) {
+        // 如果需要的位置正好是上一步中被释放的位置，不做处理
+        used_list_len[need_block_id] += 1;
+        const int ori_free_list_len = atomicSub(free_list_len, 1);
+        int *block_table_now = block_tables + need_block_id * block_num_per_seq;
+#ifdef DEBUG_STEP
+        printf("need_block_id %d\n", need_block_id);
+        printf("ori_free_list_len %d\n", ori_free_list_len);
+        printf("max_draft_tokens %d\n", max_draft_tokens);
+        printf("seq_lens_decoder[need_block_id] %d\n",
+               seq_lens_decoder[need_block_id]);
+        printf("free_list[ori_free_list_len - 1] %d\n",
+               free_list[ori_free_list_len - 1]);
+#endif
+        block_table_now[(seq_lens_decoder[need_block_id] + max_draft_tokens +
+                         1) /
+                        block_size] = free_list[ori_free_list_len - 1];
+      }
+      need_block_list[tid] = -1;
+    }
+  }
+  __syncthreads();
+
+  // 计算可以复原的query id
+  if (tid == 0) {
+    int ori_free_list_len = free_list_len[0];
+    int ori_step_len = step_len[0];
+    if (ori_step_len > 0) {
+      int ori_step_block_id = step_block_list[ori_step_len - 1];
+      int tmp_used_len = used_list_len[ori_step_block_id];
+      // 比之前调度时多分配一个block，防止马上恢复刚调度的query(比如回收的seq_id在need_block_list中）
+      const int max_decoder_block_num_this_seq =
+          max_decoder_block_num - encoder_block_lens[ori_step_block_id];
+      int used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
+                         ? tmp_used_len + 1
+                         : max_decoder_block_num_this_seq;
+      if (ori_step_len > 0 && ori_free_list_len >= used_len) {
+        // #ifdef DEBUG_STEP
+        printf(
+            "recover seq_id: %d, free_list_len: %d, used_list_len: "
+            "%d\n",
+            ori_step_block_id,
+            ori_free_list_len,
+            used_len);
+        // #endif
+        recover_block_list[recover_len[0]] = ori_step_block_id;
+        is_block_step[ori_step_block_id] = false;
+        used_list_len[ori_step_block_id] = used_len;
+        ori_free_list_len -= used_len;
+        step_block_list[ori_step_len - 1] = -1;
+        step_len[0] -= 1;
+        recover_len[0] += 1;
+        ori_step_len = step_len[0];
+        if (ori_step_len > 0) {
+          ori_step_block_id = step_block_list[ori_step_len - 1];
+          tmp_used_len = used_list_len[ori_step_block_id];
+          used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
+                         ? tmp_used_len + 1
+                         : max_decoder_block_num_this_seq;
+        }
+      }
+    }
+    need_block_len[0] = 0;
+  }
 }
 
 // 根据上一步计算出的可以复原的query_id进行状态恢复
@@ -256,58 +255,57 @@ __global__ void speculate_recover_block(int *recover_block_list,  // [bsz]
                                         const int block_num_per_seq,
                                         const int length,
                                         const int pre_id_length) {
-    const int bid = blockIdx.x;
-    const int tid = threadIdx.x;
-    __shared__ int ori_free_list_len;
-    if (bid < recover_len[0]) {
-        const int recover_id = recover_block_list[bid];
-        const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
-        const int step_idx_now = step_idx[recover_id];
-        const int seq_len = ori_seq_len_encoder + step_idx_now;
-        const int encoder_block_len = encoder_block_lens[recover_id];
-        const int decoder_used_len = used_list_len[recover_id];
-        int *block_table_now = block_tables + recover_id * block_num_per_seq;
-        int64_t *input_ids_now = input_ids + recover_id * length;
-        int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
-        if (tid == 0) {
-            seq_lens_this_time[recover_id] = seq_len;
-            seq_lens_encoder[recover_id] = seq_len;
-            stop_flags[recover_id] = false;
-            // input_ids_now[ori_seq_len_encoder + step_idx_now - 1] =
-            //     next_tokens[recover_id];  // next tokens
-            input_ids_now[0] =
-                first_token_ids[recover_id];  // set first prompt token
-            const int ori_free_list_len_tid0 =
-                atomicSub(free_list_len, decoder_used_len);
-            ori_free_list_len = ori_free_list_len_tid0;
+  const int bid = blockIdx.x;
+  const int tid = threadIdx.x;
+  __shared__ int ori_free_list_len;
+  if (bid < recover_len[0]) {
+    const int recover_id = recover_block_list[bid];
+    const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
+    const int step_idx_now = step_idx[recover_id];
+    const int seq_len = ori_seq_len_encoder + step_idx_now;
+    const int encoder_block_len = encoder_block_lens[recover_id];
+    const int decoder_used_len = used_list_len[recover_id];
+    int *block_table_now = block_tables + recover_id * block_num_per_seq;
+    int64_t *input_ids_now = input_ids + recover_id * length;
+    int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
+    if (tid == 0) {
+      seq_lens_this_time[recover_id] = seq_len;
+      seq_lens_encoder[recover_id] = seq_len;
+      stop_flags[recover_id] = false;
+      // input_ids_now[ori_seq_len_encoder + step_idx_now - 1] =
+      //     next_tokens[recover_id];  // next tokens
+      input_ids_now[0] = first_token_ids[recover_id];  // set first prompt token
+      const int ori_free_list_len_tid0 =
+          atomicSub(free_list_len, decoder_used_len);
+      ori_free_list_len = ori_free_list_len_tid0;
 #ifdef DEBUG_STEP
-            printf(
-                "seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, "
-                "seq_len: %d, ori_free_list_len_tid0: %d, "
-                "ori_free_list_len: %d\n",
-                recover_id,
-                ori_seq_len_encoder,
-                step_idx_now,
-                seq_len,
-                ori_free_list_len_tid0,
-                ori_free_list_len);
+      printf(
+          "seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, "
+          "seq_len: %d, ori_free_list_len_tid0: %d, "
+          "ori_free_list_len: %d\n",
+          recover_id,
+          ori_seq_len_encoder,
+          step_idx_now,
+          seq_len,
+          ori_free_list_len_tid0,
+          ori_free_list_len);
 #endif
-        }
-        __syncthreads();
-        // 恢复block table
-        for (int i = tid; i < decoder_used_len; i += blockDim.x) {
-            block_table_now[encoder_block_len + i] =
-                free_list[ori_free_list_len - i - 1];
-        }
-        // 恢复input_ids
-        for (int i = tid; i < step_idx_now; i += blockDim.x) {
-            input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
-        }
     }
+    __syncthreads();
+    // 恢复block table
+    for (int i = tid; i < decoder_used_len; i += blockDim.x) {
+      block_table_now[encoder_block_len + i] =
+          free_list[ori_free_list_len - i - 1];
+    }
+    // 恢复input_ids
+    for (int i = tid; i < step_idx_now; i += blockDim.x) {
+      input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
+    }
+  }
 
-    if (bid == 0 && tid == 0) {
-        recover_len[0] = 0;
-    }
+  if (bid == 0 && tid == 0) {
+    recover_len[0] = 0;
+  }
 }
 
 void SpeculateStepPaddle(
@@ -337,82 +335,82 @@ void SpeculateStepPaddle(
     const int block_size,
     const int encoder_decoder_block_num,
     const int max_draft_tokens) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    const int length = input_ids.shape()[1];
-    const int pre_id_length = pre_ids.shape()[1];
-    constexpr int BlockSize = 256;  // bsz <= 256
-    const int max_decoder_block_num = length / block_size;
-    // const int max_decoder_block_num = 2048 / block_size -
-    // encoder_decoder_block_num;
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  const int length = input_ids.shape()[1];
+  const int pre_id_length = pre_ids.shape()[1];
+  constexpr int BlockSize = 256;  // bsz <= 256
+  const int max_decoder_block_num = length / block_size;
+  // const int max_decoder_block_num = 2048 / block_size -
+  // encoder_decoder_block_num;
 #ifdef DEBUG_STEP
-    printf(
-        "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
-        "%d\n",
+  printf(
+      "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
+      "%d\n",
+      bsz,
+      block_num_per_seq,
+      length,
+      max_decoder_block_num);
+#endif
+  speculate_free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
+      const_cast<bool *>(stop_flags.data<bool>()),
+      const_cast<int *>(seq_lens_this_time.data<int>()),
+      const_cast<int *>(seq_lens_decoder.data<int>()),
+      const_cast<int *>(block_tables.data<int>()),
+      const_cast<int *>(encoder_block_lens.data<int>()),
+      const_cast<bool *>(is_block_step.data<bool>()),
+      const_cast<int *>(step_block_list.data<int>()),
+      const_cast<int *>(step_lens.data<int>()),
+      const_cast<int *>(recover_block_list.data<int>()),
+      const_cast<int *>(recover_lens.data<int>()),
+      const_cast<int *>(need_block_list.data<int>()),
+      const_cast<int *>(need_block_len.data<int>()),
+      const_cast<int *>(used_list_len.data<int>()),
+      const_cast<int *>(free_list.data<int>()),
+      const_cast<int *>(free_list_len.data<int>()),
+      const_cast<int64_t *>(first_token_ids.data<int64_t>()),
+      const_cast<int *>(accept_num.data<int>()),
+      bsz,
+      block_size,
+      block_num_per_seq,
+      max_decoder_block_num,
+      max_draft_tokens);
+#ifdef DEBUG_STEP
+  cudaDeviceSynchronize();
+#endif
+  auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
+  const int grid_size = cpu_recover_lens.data<int>()[0];
+#ifdef DEBUG_STEP
+  printf("grid_size2 %d\n", grid_size);
+#endif
+  if (grid_size > 0) {
+    speculate_recover_block<<<grid_size, BlockSize, 0, cu_stream>>>(
+        const_cast<int *>(recover_block_list.data<int>()),
+        const_cast<int *>(recover_lens.data<int>()),
+        const_cast<bool *>(stop_flags.data<bool>()),
+        const_cast<int *>(seq_lens_this_time.data<int>()),
+        const_cast<int *>(ori_seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_decoder.data<int>()),
+        const_cast<int *>(block_tables.data<int>()),
+        const_cast<int *>(free_list.data<int>()),
+        const_cast<int *>(free_list_len.data<int>()),
+        const_cast<int64_t *>(input_ids.data<int64_t>()),
+        const_cast<int64_t *>(pre_ids.data<int64_t>()),
+        const_cast<int64_t *>(step_idx.data<int64_t>()),
+        const_cast<int *>(encoder_block_lens.data<int>()),
+        const_cast<int *>(used_list_len.data<int>()),
+        next_tokens.data<int64_t>(),
+        first_token_ids.data<int64_t>(),
         bsz,
         block_num_per_seq,
         length,
-        max_decoder_block_num);
-#endif
-    speculate_free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
-        const_cast<bool *>(stop_flags.data<bool>()),
-        const_cast<int *>(seq_lens_this_time.data<int>()),
-        const_cast<int *>(seq_lens_decoder.data<int>()),
-        const_cast<int *>(block_tables.data<int>()),
-        const_cast<int *>(encoder_block_lens.data<int>()),
-        const_cast<bool *>(is_block_step.data<bool>()),
-        const_cast<int *>(step_block_list.data<int>()),
-        const_cast<int *>(step_lens.data<int>()),
-        const_cast<int *>(recover_block_list.data<int>()),
-        const_cast<int *>(recover_lens.data<int>()),
-        const_cast<int *>(need_block_list.data<int>()),
-        const_cast<int *>(need_block_len.data<int>()),
-        const_cast<int *>(used_list_len.data<int>()),
-        const_cast<int *>(free_list.data<int>()),
-        const_cast<int *>(free_list_len.data<int>()),
-        const_cast<int64_t *>(first_token_ids.data<int64_t>()),
-        const_cast<int *>(accept_num.data<int>()),
-        bsz,
-        block_size,
-        block_num_per_seq,
-        max_decoder_block_num,
-        max_draft_tokens);
+        pre_id_length);
 #ifdef DEBUG_STEP
     cudaDeviceSynchronize();
 #endif
-    auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
-    const int grid_size = cpu_recover_lens.data<int>()[0];
-#ifdef DEBUG_STEP
-    printf("grid_size2 %d\n", grid_size);
-#endif
-    if (grid_size > 0) {
-        speculate_recover_block<<<grid_size, BlockSize, 0, cu_stream>>>(
-            const_cast<int *>(recover_block_list.data<int>()),
-            const_cast<int *>(recover_lens.data<int>()),
-            const_cast<bool *>(stop_flags.data<bool>()),
-            const_cast<int *>(seq_lens_this_time.data<int>()),
-            const_cast<int *>(ori_seq_lens_encoder.data<int>()),
-            const_cast<int *>(seq_lens_encoder.data<int>()),
-            const_cast<int *>(seq_lens_decoder.data<int>()),
-            const_cast<int *>(block_tables.data<int>()),
-            const_cast<int *>(free_list.data<int>()),
-            const_cast<int *>(free_list_len.data<int>()),
-            const_cast<int64_t *>(input_ids.data<int64_t>()),
-            const_cast<int64_t *>(pre_ids.data<int64_t>()),
-            const_cast<int64_t *>(step_idx.data<int64_t>()),
-            const_cast<int *>(encoder_block_lens.data<int>()),
-            const_cast<int *>(used_list_len.data<int>()),
-            next_tokens.data<int64_t>(),
-            first_token_ids.data<int64_t>(),
-            bsz,
-            block_num_per_seq,
-            length,
-            pre_id_length);
-#ifdef DEBUG_STEP
-        cudaDeviceSynchronize();
-#endif
-    }
+  }
 }
 
 PD_BUILD_STATIC_OP(speculate_step_paddle)
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_step_reschedule.cu b/custom_ops/gpu_ops/speculate_decoding/speculate_step_reschedule.cu
index 853f894e0b..6c520dcdc3 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_step_reschedule.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_step_reschedule.cu
@@ -19,308 +19,313 @@
 #define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
 #endif
 
-__device__ __forceinline__ bool in_need_block_list_schedule(const int &qid,
-                                            int *need_block_list,
-                                            const int &need_block_len) {
-    bool res = false;
-    for (int i = 0; i < need_block_len; i++) {
-        if (qid == need_block_list[i]) {
-            res = true;
-            need_block_list[i] = -1;
-            break;
-        }
+__device__ __forceinline__ bool in_need_block_list_schedule(
+    const int &qid, int *need_block_list, const int &need_block_len) {
+  bool res = false;
+  for (int i = 0; i < need_block_len; i++) {
+    if (qid == need_block_list[i]) {
+      res = true;
+      need_block_list[i] = -1;
+      break;
     }
-    return res;
+  }
+  return res;
 }
 
 __global__ void speculate_free_and_reschedule(bool *stop_flags,
-                                   int *seq_lens_this_time,
-                                   int *seq_lens_decoder,
-                                   int *block_tables,
-                                   int *encoder_block_lens,
-                                   bool *is_block_step,
-                                   int *step_block_list,  // [bsz]
-                                   int *step_len,
-                                   int *recover_block_list,
-                                   int *recover_len,
-                                   int *need_block_list,
-                                   int *need_block_len,
-                                   int *used_list_len,
-                                   int *free_list,
-                                   int *free_list_len,
-                                   int64_t *first_token_ids,
-                                   int* accept_num,
-                                   const int bsz,
-                                   const int block_size,
-                                   const int block_num_per_seq,
-                                   const int max_decoder_block_num,
-                                   const int max_draft_tokens) {
-    typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
-    __shared__ bool step_max_block_flag;
-    __shared__ int in_need_block_list_len;
-    const int tid = threadIdx.x;
-    if (tid < bsz) {
-        if (tid == 0) {
-            step_max_block_flag = false;
-            in_need_block_list_len = 0;
-        }
-        int *block_table_now = block_tables + tid * block_num_per_seq;
-        int max_possible_block_idx = (seq_lens_decoder[tid] + max_draft_tokens + 1 ) / block_size;
-        if (stop_flags[tid]) {
-            // 回收block块
-            first_token_ids[tid] = -1;
-            const int encoder_block_len = encoder_block_lens[tid];
-            const int decoder_used_len = used_list_len[tid];
-            if (decoder_used_len > 0) {
-                const int ori_free_list_len =
-                    atomicAdd(free_list_len, decoder_used_len);
-#ifdef DEBUG_STEP
-                printf(
-                    "free block seq_id: %d, free block num: %d, "
-                    "encoder_block_len: %d, ori_free_list_len: %d\n",
-                    tid,
-                    decoder_used_len,
-                    encoder_block_len,
-                    ori_free_list_len);
-#endif
-                for (int i = 0; i < decoder_used_len; i++) {
-                    free_list[ori_free_list_len + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                encoder_block_lens[tid] = 0;
-                used_list_len[tid] = 0;
-            }
-        } else if (seq_lens_this_time[tid] != 0 && max_possible_block_idx < block_num_per_seq &&
-                   block_table_now[(seq_lens_decoder[tid]  + max_draft_tokens  +
-                                    1) /
-                                   block_size] == -1) {
-            // 统计需要分配block的位置和总数
-#ifdef DEBUG_STEP
-            printf("step seq_id:%d, ##### pin 1 #####\n", tid);
-#endif
-            const int ori_need_block_len = atomicAdd(need_block_len, 1);
-            need_block_list[ori_need_block_len] = tid;
-#ifdef DEBUG_STEP
-            printf("seq_id: %d need block\n", tid);
-#endif
-        }
-    }
-#ifdef DEBUG_STEP
-    printf("step seq_id:%d, ##### pin 2 #####\n", tid);
-#endif
-    __syncthreads();
-
-    // 调度block，直到满足need_block_len
-    while (need_block_len[0] > free_list_len[0]) {
-        if (tid == 0) {
-            printf("need_block_len: %d, free_list_len: %d\n",
-                   need_block_len[0],
-                   free_list_len[0]);
-        }
-        // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
-        const int used_block_num =
-            tid < bsz ? used_list_len[tid] : 0;
-        cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
-        kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
-        if (tid == 0) {
-            if (kv_pair.value == 0) {
-                step_max_block_flag = true;
-            } else {
-                const int encoder_block_len = encoder_block_lens[kv_pair.key];
-                printf("step max_id: %d, max_num: %d, encoder_block_len: %d\n",
-                       kv_pair.key,
-                       kv_pair.value,
-                       encoder_block_len);
-                int *block_table_now =
-                    block_tables + kv_pair.key * block_num_per_seq;
-                // 回收调度位的block
-                for (int i = 0; i < kv_pair.value; i++) {
-                    free_list[free_list_len[0] + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                step_block_list[step_len[0]] = kv_pair.key;
-                // 如果调度位置本次也需要block，对应的处理
-                if (in_need_block_list_schedule(
-                        kv_pair.key,
-                        need_block_list,
-                        need_block_len[0] + in_need_block_list_len)) {
-                    need_block_len[0] -= 1;
-                    in_need_block_list_len += 1;
-                }
-                step_len[0] += 1;
-                free_list_len[0] += kv_pair.value;
-                stop_flags[kv_pair.key] = true;
-                seq_lens_this_time[kv_pair.key] = 0;
-                seq_lens_decoder[kv_pair.key] = 0;
-                encoder_block_lens[kv_pair.key] = 0;
-                used_list_len[kv_pair.key] = 0;
-                printf(
-                    "free block seq_id: %d, free block num: %d, "
-                    "now_free_list_len: %d\n",
-                    (int)kv_pair.key,
-                    (int)kv_pair.value,
-                    (int)free_list_len[0]);
-            }
-        }
-        __syncthreads();
-    }
-#ifdef DEBUG_STEP
-    printf("step seq_id:%d, ##### pin 3 #####\n", tid);
-#endif
-    // 为需要block的位置分配block，每个位置分配一个block
-    if (tid < need_block_len[0] + in_need_block_list_len) {
-        const int need_block_id = need_block_list[tid];
-        if (need_block_id != -1) {
-            if (!stop_flags[need_block_id]) {
-                // 如果需要的位置正好是上一步中被释放的位置，不做处理
-                used_list_len[need_block_id] += 1;
-                const int ori_free_list_len = atomicSub(free_list_len, 1);
-                int *block_table_now =
-                    block_tables + need_block_id * block_num_per_seq;
-                block_table_now[(seq_lens_decoder[need_block_id] +
-                                 max_draft_tokens + 1) /
-                                block_size] = free_list[ori_free_list_len - 1];
-            }
-            need_block_list[tid] = -1;
-        }
-    }
-    __syncthreads();
-    // reset need_block_len
+                                              int *seq_lens_this_time,
+                                              int *seq_lens_decoder,
+                                              int *block_tables,
+                                              int *encoder_block_lens,
+                                              bool *is_block_step,
+                                              int *step_block_list,  // [bsz]
+                                              int *step_len,
+                                              int *recover_block_list,
+                                              int *recover_len,
+                                              int *need_block_list,
+                                              int *need_block_len,
+                                              int *used_list_len,
+                                              int *free_list,
+                                              int *free_list_len,
+                                              int64_t *first_token_ids,
+                                              int *accept_num,
+                                              const int bsz,
+                                              const int block_size,
+                                              const int block_num_per_seq,
+                                              const int max_decoder_block_num,
+                                              const int max_draft_tokens) {
+  typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ bool step_max_block_flag;
+  __shared__ int in_need_block_list_len;
+  const int tid = threadIdx.x;
+  if (tid < bsz) {
     if (tid == 0) {
-        need_block_len[0] = 0;
+      step_max_block_flag = false;
+      in_need_block_list_len = 0;
     }
+    int *block_table_now = block_tables + tid * block_num_per_seq;
+    int max_possible_block_idx =
+        (seq_lens_decoder[tid] + max_draft_tokens + 1) / block_size;
+    if (stop_flags[tid]) {
+      // 回收block块
+      first_token_ids[tid] = -1;
+      const int encoder_block_len = encoder_block_lens[tid];
+      const int decoder_used_len = used_list_len[tid];
+      if (decoder_used_len > 0) {
+        const int ori_free_list_len =
+            atomicAdd(free_list_len, decoder_used_len);
+#ifdef DEBUG_STEP
+        printf(
+            "free block seq_id: %d, free block num: %d, "
+            "encoder_block_len: %d, ori_free_list_len: %d\n",
+            tid,
+            decoder_used_len,
+            encoder_block_len,
+            ori_free_list_len);
+#endif
+        for (int i = 0; i < decoder_used_len; i++) {
+          free_list[ori_free_list_len + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        encoder_block_lens[tid] = 0;
+        used_list_len[tid] = 0;
+      }
+    } else if (seq_lens_this_time[tid] != 0 &&
+               max_possible_block_idx < block_num_per_seq &&
+               block_table_now[(seq_lens_decoder[tid] + max_draft_tokens + 1) /
+                               block_size] == -1) {
+      // 统计需要分配block的位置和总数
+#ifdef DEBUG_STEP
+      printf("step seq_id:%d, ##### pin 1 #####\n", tid);
+#endif
+      const int ori_need_block_len = atomicAdd(need_block_len, 1);
+      need_block_list[ori_need_block_len] = tid;
+#ifdef DEBUG_STEP
+      printf("seq_id: %d need block\n", tid);
+#endif
+    }
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 2 #####\n", tid);
+#endif
+  __syncthreads();
+
+  // 调度block，直到满足need_block_len
+  while (need_block_len[0] > free_list_len[0]) {
+    if (tid == 0) {
+      printf("need_block_len: %d, free_list_len: %d\n",
+             need_block_len[0],
+             free_list_len[0]);
+    }
+    // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
+    const int used_block_num = tid < bsz ? used_list_len[tid] : 0;
+    cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
+    kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
+    if (tid == 0) {
+      if (kv_pair.value == 0) {
+        step_max_block_flag = true;
+      } else {
+        const int encoder_block_len = encoder_block_lens[kv_pair.key];
+        printf("step max_id: %d, max_num: %d, encoder_block_len: %d\n",
+               kv_pair.key,
+               kv_pair.value,
+               encoder_block_len);
+        int *block_table_now = block_tables + kv_pair.key * block_num_per_seq;
+        // 回收调度位的block
+        for (int i = 0; i < kv_pair.value; i++) {
+          free_list[free_list_len[0] + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        step_block_list[step_len[0]] = kv_pair.key;
+        // 如果调度位置本次也需要block，对应的处理
+        if (in_need_block_list_schedule(
+                kv_pair.key,
+                need_block_list,
+                need_block_len[0] + in_need_block_list_len)) {
+          need_block_len[0] -= 1;
+          in_need_block_list_len += 1;
+        }
+        step_len[0] += 1;
+        free_list_len[0] += kv_pair.value;
+        stop_flags[kv_pair.key] = true;
+        seq_lens_this_time[kv_pair.key] = 0;
+        seq_lens_decoder[kv_pair.key] = 0;
+        encoder_block_lens[kv_pair.key] = 0;
+        used_list_len[kv_pair.key] = 0;
+        printf(
+            "free block seq_id: %d, free block num: %d, "
+            "now_free_list_len: %d\n",
+            (int)kv_pair.key,
+            (int)kv_pair.value,
+            (int)free_list_len[0]);
+      }
+    }
+    __syncthreads();
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 3 #####\n", tid);
+#endif
+  // 为需要block的位置分配block，每个位置分配一个block
+  if (tid < need_block_len[0] + in_need_block_list_len) {
+    const int need_block_id = need_block_list[tid];
+    if (need_block_id != -1) {
+      if (!stop_flags[need_block_id]) {
+        // 如果需要的位置正好是上一步中被释放的位置，不做处理
+        used_list_len[need_block_id] += 1;
+        const int ori_free_list_len = atomicSub(free_list_len, 1);
+        int *block_table_now = block_tables + need_block_id * block_num_per_seq;
+        block_table_now[(seq_lens_decoder[need_block_id] + max_draft_tokens +
+                         1) /
+                        block_size] = free_list[ori_free_list_len - 1];
+      }
+      need_block_list[tid] = -1;
+    }
+  }
+  __syncthreads();
+  // reset need_block_len
+  if (tid == 0) {
+    need_block_len[0] = 0;
+  }
 }
 
 // 为不修改接口调用方式，入参暂不改变
-void SpeculateStepSchedule(const paddle::Tensor &stop_flags,
-              const paddle::Tensor &seq_lens_this_time,
-              const paddle::Tensor &ori_seq_lens_encoder,
-              const paddle::Tensor &seq_lens_encoder,
-              const paddle::Tensor &seq_lens_decoder,
-              const paddle::Tensor &block_tables,  // [bsz, block_num_per_seq]
-              const paddle::Tensor &encoder_block_lens,
-              const paddle::Tensor &is_block_step,
-              const paddle::Tensor &step_block_list,
-              const paddle::Tensor &step_lens,
-              const paddle::Tensor &recover_block_list,
-              const paddle::Tensor &recover_lens,
-              const paddle::Tensor &need_block_list,
-              const paddle::Tensor &need_block_len,
-              const paddle::Tensor &used_list_len,
-              const paddle::Tensor &free_list,
-              const paddle::Tensor &free_list_len,
-              const paddle::Tensor &input_ids,
-              const paddle::Tensor &pre_ids,
-              const paddle::Tensor &step_idx,
-              const paddle::Tensor &next_tokens,
-              const paddle::Tensor &first_token_ids,
-              const paddle::Tensor &accept_num,
-              const int block_size,
-              const int encoder_decoder_block_num,
-              const int max_draft_tokens) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    const int length = input_ids.shape()[1];
-    const int pre_id_length = pre_ids.shape()[1];
-    constexpr int BlockSize = 256;  // bsz <= 256
-    const int max_decoder_block_num = length / block_size - encoder_decoder_block_num; // 最大输出长度对应的block - 服务为解码分配的block数量
-    auto step_lens_inkernel = paddle::full({1}, 0, paddle::DataType::INT32, stop_flags.place());
-    auto step_bs_list = GetEmptyTensor({bsz}, paddle::DataType::INT32, stop_flags.place());
+void SpeculateStepSchedule(
+    const paddle::Tensor &stop_flags,
+    const paddle::Tensor &seq_lens_this_time,
+    const paddle::Tensor &ori_seq_lens_encoder,
+    const paddle::Tensor &seq_lens_encoder,
+    const paddle::Tensor &seq_lens_decoder,
+    const paddle::Tensor &block_tables,  // [bsz, block_num_per_seq]
+    const paddle::Tensor &encoder_block_lens,
+    const paddle::Tensor &is_block_step,
+    const paddle::Tensor &step_block_list,
+    const paddle::Tensor &step_lens,
+    const paddle::Tensor &recover_block_list,
+    const paddle::Tensor &recover_lens,
+    const paddle::Tensor &need_block_list,
+    const paddle::Tensor &need_block_len,
+    const paddle::Tensor &used_list_len,
+    const paddle::Tensor &free_list,
+    const paddle::Tensor &free_list_len,
+    const paddle::Tensor &input_ids,
+    const paddle::Tensor &pre_ids,
+    const paddle::Tensor &step_idx,
+    const paddle::Tensor &next_tokens,
+    const paddle::Tensor &first_token_ids,
+    const paddle::Tensor &accept_num,
+    const int block_size,
+    const int encoder_decoder_block_num,
+    const int max_draft_tokens) {
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  const int length = input_ids.shape()[1];
+  const int pre_id_length = pre_ids.shape()[1];
+  constexpr int BlockSize = 256;  // bsz <= 256
+  const int max_decoder_block_num =
+      length / block_size -
+      encoder_decoder_block_num;  // 最大输出长度对应的block -
+                                  // 服务为解码分配的block数量
+  auto step_lens_inkernel =
+      paddle::full({1}, 0, paddle::DataType::INT32, stop_flags.place());
+  auto step_bs_list =
+      GetEmptyTensor({bsz}, paddle::DataType::INT32, stop_flags.place());
 #ifdef DEBUG_STEP
-    printf(
-        "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
-        "%d\n",
-        bsz,
-        block_num_per_seq,
-        length,
-        max_decoder_block_num);
+  printf(
+      "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
+      "%d\n",
+      bsz,
+      block_num_per_seq,
+      length,
+      max_decoder_block_num);
 #endif
-    speculate_free_and_reschedule<<<1, BlockSize, 0, cu_stream>>>(
-        const_cast<bool *>(stop_flags.data<bool>()),
-        const_cast<int *>(seq_lens_this_time.data<int>()),
-        const_cast<int *>(seq_lens_decoder.data<int>()),
-        const_cast<int *>(block_tables.data<int>()),
-        const_cast<int *>(encoder_block_lens.data<int>()),
-        const_cast<bool *>(is_block_step.data<bool>()),
-        const_cast<int *>(step_bs_list.data<int>()),
-        const_cast<int *>(step_lens_inkernel.data<int>()),
-        const_cast<int *>(recover_block_list.data<int>()),
-        const_cast<int *>(recover_lens.data<int>()),
-        const_cast<int *>(need_block_list.data<int>()),
-        const_cast<int *>(need_block_len.data<int>()),
-        const_cast<int *>(used_list_len.data<int>()),
-        const_cast<int *>(free_list.data<int>()),
-        const_cast<int *>(free_list_len.data<int>()),
-        const_cast<int64_t *>(first_token_ids.data<int64_t>()),
-        const_cast<int *>(accept_num.data<int>()),
-        bsz,
-        block_size,
-        block_num_per_seq,
-        max_decoder_block_num,
-        max_draft_tokens);
+  speculate_free_and_reschedule<<<1, BlockSize, 0, cu_stream>>>(
+      const_cast<bool *>(stop_flags.data<bool>()),
+      const_cast<int *>(seq_lens_this_time.data<int>()),
+      const_cast<int *>(seq_lens_decoder.data<int>()),
+      const_cast<int *>(block_tables.data<int>()),
+      const_cast<int *>(encoder_block_lens.data<int>()),
+      const_cast<bool *>(is_block_step.data<bool>()),
+      const_cast<int *>(step_bs_list.data<int>()),
+      const_cast<int *>(step_lens_inkernel.data<int>()),
+      const_cast<int *>(recover_block_list.data<int>()),
+      const_cast<int *>(recover_lens.data<int>()),
+      const_cast<int *>(need_block_list.data<int>()),
+      const_cast<int *>(need_block_len.data<int>()),
+      const_cast<int *>(used_list_len.data<int>()),
+      const_cast<int *>(free_list.data<int>()),
+      const_cast<int *>(free_list_len.data<int>()),
+      const_cast<int64_t *>(first_token_ids.data<int64_t>()),
+      const_cast<int *>(accept_num.data<int>()),
+      bsz,
+      block_size,
+      block_num_per_seq,
+      max_decoder_block_num,
+      max_draft_tokens);
 #ifdef DEBUG_STEP
-    cudaDeviceSynchronize();
+  cudaDeviceSynchronize();
 #endif
-    // save output
-    auto step_lens_cpu = step_lens_inkernel.copy_to(paddle::CPUPlace(), false);
-    if (step_lens_cpu.data<int>()[0] > 0) {
-        auto step_bs_list_cpu = step_bs_list.copy_to(paddle::CPUPlace(), false);
-        auto next_tokens = paddle::full({bsz}, -1, paddle::DataType::INT64, paddle::CPUPlace());
-        for (int i = 0; i < step_lens_cpu.data<int>()[0]; i++) {
-            const int step_bid = step_bs_list_cpu.data<int>()[i];
-            next_tokens.data<int64_t>()[step_bid] = -3; // need reschedule
-        }
-        const int rank_id = static_cast<int>(stop_flags.place().GetDeviceId());
-        printf("reschedule rank_id: %d, step_lens: %d", rank_id, step_lens_cpu.data<int>()[0]);
-        const int64_t* x_data = next_tokens.data<int64_t>();
-        static struct speculate_msgdata msg_sed;
-        int msg_queue_id = rank_id;
-        if (const char* inference_msg_queue_id_env_p =
-                std::getenv("INFERENCE_MSG_QUEUE_ID")) {
-            std::string inference_msg_queue_id_env_str(
-                inference_msg_queue_id_env_p);
-            int inference_msg_queue_id_from_env =
-                std::stoi(inference_msg_queue_id_env_str);
-            msg_queue_id = inference_msg_queue_id_from_env;
-        } else {
-            std::cout << "Failed to got INFERENCE_MSG_QUEUE_ID at env, use default."
-                    << std::endl;
-        }
-        int inference_msg_id_from_env = 1;
-        if (const char* inference_msg_id_env_p = std::getenv("INFERENCE_MSG_ID")) {
-            std::string inference_msg_id_env_str(inference_msg_id_env_p);
-            inference_msg_id_from_env = std::stoi(inference_msg_id_env_str);
-            if (inference_msg_id_from_env == 2) {
-                // 2 and -2 is preserve for no-output indication.
-                throw std::runtime_error(
-                    " INFERENCE_MSG_ID cannot be 2, please use other number.");
-            }
-            if (inference_msg_id_from_env < 0) {
-                throw std::runtime_error(
-                    " INFERENCE_MSG_ID cannot be negative, please use other "
-                    "number.");
-            }
-
-        } else {
-        }
-        // static key_t key = ftok("/dev/shm", msg_queue_id);
-        static key_t key = ftok("./", msg_queue_id);
-
-        static int msgid = msgget(key, IPC_CREAT | 0666);
-        msg_sed.mtype = 1;
-        msg_sed.mtext[0] = inference_msg_id_from_env;
-        msg_sed.mtext[1] = bsz;
-        for (int i = 2; i < bsz + 2; i++) {
-            msg_sed.mtext[i] = (int)x_data[i - 2];
-        }
-        if ((msgsnd(msgid, &msg_sed, (MAX_BSZ + 2) * 4, 0)) == -1) {
-            printf("full msg buffer\n");
-        }
+  // save output
+  auto step_lens_cpu = step_lens_inkernel.copy_to(paddle::CPUPlace(), false);
+  if (step_lens_cpu.data<int>()[0] > 0) {
+    auto step_bs_list_cpu = step_bs_list.copy_to(paddle::CPUPlace(), false);
+    auto next_tokens =
+        paddle::full({bsz}, -1, paddle::DataType::INT64, paddle::CPUPlace());
+    for (int i = 0; i < step_lens_cpu.data<int>()[0]; i++) {
+      const int step_bid = step_bs_list_cpu.data<int>()[i];
+      next_tokens.data<int64_t>()[step_bid] = -3;  // need reschedule
     }
+    const int rank_id = static_cast<int>(stop_flags.place().GetDeviceId());
+    printf("reschedule rank_id: %d, step_lens: %d",
+           rank_id,
+           step_lens_cpu.data<int>()[0]);
+    const int64_t *x_data = next_tokens.data<int64_t>();
+    static struct speculate_msgdata msg_sed;
+    int msg_queue_id = rank_id;
+    if (const char *inference_msg_queue_id_env_p =
+            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+      std::string inference_msg_queue_id_env_str(inference_msg_queue_id_env_p);
+      int inference_msg_queue_id_from_env =
+          std::stoi(inference_msg_queue_id_env_str);
+      msg_queue_id = inference_msg_queue_id_from_env;
+    } else {
+      std::cout << "Failed to got INFERENCE_MSG_QUEUE_ID at env, use default."
+                << std::endl;
+    }
+    int inference_msg_id_from_env = 1;
+    if (const char *inference_msg_id_env_p = std::getenv("INFERENCE_MSG_ID")) {
+      std::string inference_msg_id_env_str(inference_msg_id_env_p);
+      inference_msg_id_from_env = std::stoi(inference_msg_id_env_str);
+      if (inference_msg_id_from_env == 2) {
+        // 2 and -2 is preserve for no-output indication.
+        throw std::runtime_error(
+            " INFERENCE_MSG_ID cannot be 2, please use other number.");
+      }
+      if (inference_msg_id_from_env < 0) {
+        throw std::runtime_error(
+            " INFERENCE_MSG_ID cannot be negative, please use other "
+            "number.");
+      }
+
+    } else {
+    }
+    // static key_t key = ftok("/dev/shm", msg_queue_id);
+    static key_t key = ftok("./", msg_queue_id);
+
+    static int msgid = msgget(key, IPC_CREAT | 0666);
+    msg_sed.mtype = 1;
+    msg_sed.mtext[0] = inference_msg_id_from_env;
+    msg_sed.mtext[1] = bsz;
+    for (int i = 2; i < bsz + 2; i++) {
+      msg_sed.mtext[i] = (int)x_data[i - 2];
+    }
+    if ((msgsnd(msgid, &msg_sed, (MAX_BSZ + 2) * 4, 0)) == -1) {
+      printf("full msg buffer\n");
+    }
+  }
 }
 
 PD_BUILD_STATIC_OP(speculate_step_reschedule)
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_step_system_cache.cu b/custom_ops/gpu_ops/speculate_decoding/speculate_step_system_cache.cu
index 5bdc2dbf81..c56fe418bf 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_step_system_cache.cu
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_step_system_cache.cu
@@ -39,59 +39,58 @@ __global__ void speculate_recover_block(int *recover_block_list,  // [bsz]
                                         const int block_num_per_seq,
                                         const int length,
                                         const int pre_id_length) {
-    const int bid = blockIdx.x;
-    const int tid = threadIdx.x;
-    __shared__ int ori_free_list_len;
-    if (bid < recover_len[0]) {
-        const int recover_id = recover_block_list[bid];
-        const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
-        const int step_idx_now = step_idx[recover_id];
-        const int seq_len = ori_seq_len_encoder + step_idx_now;
-        const int encoder_block_len = encoder_block_lens[recover_id];
-        const int decoder_used_len = used_list_len[recover_id];
-        int *block_table_now = block_tables + recover_id * block_num_per_seq;
-        int64_t *input_ids_now = input_ids + recover_id * length;
-        int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
-        if (tid == 0) {
-            seq_lens_this_time[recover_id] = seq_len;
-            seq_lens_encoder[recover_id] = seq_len;
-            seq_lens_decoder[recover_id] = ori_seq_lens_decoder[recover_id];
-            stop_flags[recover_id] = false;
-            // input_ids_now[ori_seq_len_encoder + step_idx_now - 1] =
-            //     next_tokens[recover_id];  // next tokens
-            input_ids_now[0] =
-                first_token_ids[recover_id];  // set first prompt token
-            const int ori_free_list_len_tid0 =
-                atomicSub(free_list_len, decoder_used_len);
-            ori_free_list_len = ori_free_list_len_tid0;
+  const int bid = blockIdx.x;
+  const int tid = threadIdx.x;
+  __shared__ int ori_free_list_len;
+  if (bid < recover_len[0]) {
+    const int recover_id = recover_block_list[bid];
+    const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
+    const int step_idx_now = step_idx[recover_id];
+    const int seq_len = ori_seq_len_encoder + step_idx_now;
+    const int encoder_block_len = encoder_block_lens[recover_id];
+    const int decoder_used_len = used_list_len[recover_id];
+    int *block_table_now = block_tables + recover_id * block_num_per_seq;
+    int64_t *input_ids_now = input_ids + recover_id * length;
+    int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
+    if (tid == 0) {
+      seq_lens_this_time[recover_id] = seq_len;
+      seq_lens_encoder[recover_id] = seq_len;
+      seq_lens_decoder[recover_id] = ori_seq_lens_decoder[recover_id];
+      stop_flags[recover_id] = false;
+      // input_ids_now[ori_seq_len_encoder + step_idx_now - 1] =
+      //     next_tokens[recover_id];  // next tokens
+      input_ids_now[0] = first_token_ids[recover_id];  // set first prompt token
+      const int ori_free_list_len_tid0 =
+          atomicSub(free_list_len, decoder_used_len);
+      ori_free_list_len = ori_free_list_len_tid0;
 #ifdef DEBUG_STEP
-            printf(
-                "seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, "
-                "seq_len: %d, ori_free_list_len_tid0: %d, "
-                "ori_free_list_len: %d\n",
-                recover_id,
-                ori_seq_len_encoder,
-                step_idx_now,
-                seq_len,
-                ori_free_list_len_tid0,
-                ori_free_list_len);
+      printf(
+          "seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, "
+          "seq_len: %d, ori_free_list_len_tid0: %d, "
+          "ori_free_list_len: %d\n",
+          recover_id,
+          ori_seq_len_encoder,
+          step_idx_now,
+          seq_len,
+          ori_free_list_len_tid0,
+          ori_free_list_len);
 #endif
-        }
-        __syncthreads();
-        // 恢复block table
-        for (int i = tid; i < decoder_used_len; i += blockDim.x) {
-            block_table_now[encoder_block_len + i] =
-                free_list[ori_free_list_len - i - 1];
-        }
-        // 恢复input_ids
-        for (int i = tid; i < step_idx_now; i += blockDim.x) {
-            input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
-        }
     }
+    __syncthreads();
+    // 恢复block table
+    for (int i = tid; i < decoder_used_len; i += blockDim.x) {
+      block_table_now[encoder_block_len + i] =
+          free_list[ori_free_list_len - i - 1];
+    }
+    // 恢复input_ids
+    for (int i = tid; i < step_idx_now; i += blockDim.x) {
+      input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
+    }
+  }
 
-    if (bid == 0 && tid == 0) {
-        recover_len[0] = 0;
-    }
+  if (bid == 0 && tid == 0) {
+    recover_len[0] = 0;
+  }
 }
 
 void SpeculateStepPaddle(
@@ -122,83 +121,83 @@ void SpeculateStepPaddle(
     const int block_size,
     const int encoder_decoder_block_num,
     const int max_draft_tokens) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    const int length = input_ids.shape()[1];
-    const int pre_id_length = pre_ids.shape()[1];
-    constexpr int BlockSize = 256;  // bsz <= 256
-    const int max_decoder_block_num = length / block_size;
-    // const int max_decoder_block_num = 2048 / block_size -
-    // encoder_decoder_block_num;
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  const int length = input_ids.shape()[1];
+  const int pre_id_length = pre_ids.shape()[1];
+  constexpr int BlockSize = 256;  // bsz <= 256
+  const int max_decoder_block_num = length / block_size;
+  // const int max_decoder_block_num = 2048 / block_size -
+  // encoder_decoder_block_num;
 #ifdef DEBUG_STEP
-    printf(
-        "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
-        "%d\n",
+  printf(
+      "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
+      "%d\n",
+      bsz,
+      block_num_per_seq,
+      length,
+      max_decoder_block_num);
+#endif
+  speculate_free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
+      const_cast<bool *>(stop_flags.data<bool>()),
+      const_cast<int *>(seq_lens_this_time.data<int>()),
+      const_cast<int *>(seq_lens_decoder.data<int>()),
+      const_cast<int *>(block_tables.data<int>()),
+      const_cast<int *>(encoder_block_lens.data<int>()),
+      const_cast<bool *>(is_block_step.data<bool>()),
+      const_cast<int *>(step_block_list.data<int>()),
+      const_cast<int *>(step_lens.data<int>()),
+      const_cast<int *>(recover_block_list.data<int>()),
+      const_cast<int *>(recover_lens.data<int>()),
+      const_cast<int *>(need_block_list.data<int>()),
+      const_cast<int *>(need_block_len.data<int>()),
+      const_cast<int *>(used_list_len.data<int>()),
+      const_cast<int *>(free_list.data<int>()),
+      const_cast<int *>(free_list_len.data<int>()),
+      const_cast<int64_t *>(first_token_ids.data<int64_t>()),
+      const_cast<int *>(accept_num.data<int>()),
+      bsz,
+      block_size,
+      block_num_per_seq,
+      max_decoder_block_num,
+      max_draft_tokens);
+#ifdef DEBUG_STEP
+  cudaDeviceSynchronize();
+#endif
+  auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
+  const int grid_size = cpu_recover_lens.data<int>()[0];
+#ifdef DEBUG_STEP
+  printf("grid_size2 %d\n", grid_size);
+#endif
+  if (grid_size > 0) {
+    speculate_recover_block<<<grid_size, BlockSize, 0, cu_stream>>>(
+        const_cast<int *>(recover_block_list.data<int>()),
+        const_cast<int *>(recover_lens.data<int>()),
+        const_cast<bool *>(stop_flags.data<bool>()),
+        const_cast<int *>(seq_lens_this_time.data<int>()),
+        const_cast<int *>(ori_seq_lens_encoder.data<int>()),
+        const_cast<int *>(ori_seq_lens_decoder.data<int>()),
+        const_cast<int *>(seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_decoder.data<int>()),
+        const_cast<int *>(block_tables.data<int>()),
+        const_cast<int *>(free_list.data<int>()),
+        const_cast<int *>(free_list_len.data<int>()),
+        const_cast<int64_t *>(input_ids.data<int64_t>()),
+        const_cast<int64_t *>(pre_ids.data<int64_t>()),
+        const_cast<int64_t *>(step_idx.data<int64_t>()),
+        const_cast<int *>(encoder_block_lens.data<int>()),
+        const_cast<int *>(used_list_len.data<int>()),
+        next_tokens.data<int64_t>(),
+        first_token_ids.data<int64_t>(),
         bsz,
         block_num_per_seq,
         length,
-        max_decoder_block_num);
-#endif
-    speculate_free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
-        const_cast<bool *>(stop_flags.data<bool>()),
-        const_cast<int *>(seq_lens_this_time.data<int>()),
-        const_cast<int *>(seq_lens_decoder.data<int>()),
-        const_cast<int *>(block_tables.data<int>()),
-        const_cast<int *>(encoder_block_lens.data<int>()),
-        const_cast<bool *>(is_block_step.data<bool>()),
-        const_cast<int *>(step_block_list.data<int>()),
-        const_cast<int *>(step_lens.data<int>()),
-        const_cast<int *>(recover_block_list.data<int>()),
-        const_cast<int *>(recover_lens.data<int>()),
-        const_cast<int *>(need_block_list.data<int>()),
-        const_cast<int *>(need_block_len.data<int>()),
-        const_cast<int *>(used_list_len.data<int>()),
-        const_cast<int *>(free_list.data<int>()),
-        const_cast<int *>(free_list_len.data<int>()),
-        const_cast<int64_t *>(first_token_ids.data<int64_t>()),
-        const_cast<int *>(accept_num.data<int>()),
-        bsz,
-        block_size,
-        block_num_per_seq,
-        max_decoder_block_num,
-        max_draft_tokens);
+        pre_id_length);
 #ifdef DEBUG_STEP
     cudaDeviceSynchronize();
 #endif
-    auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
-    const int grid_size = cpu_recover_lens.data<int>()[0];
-#ifdef DEBUG_STEP
-    printf("grid_size2 %d\n", grid_size);
-#endif
-    if (grid_size > 0) {
-        speculate_recover_block<<<grid_size, BlockSize, 0, cu_stream>>>(
-            const_cast<int *>(recover_block_list.data<int>()),
-            const_cast<int *>(recover_lens.data<int>()),
-            const_cast<bool *>(stop_flags.data<bool>()),
-            const_cast<int *>(seq_lens_this_time.data<int>()),
-            const_cast<int *>(ori_seq_lens_encoder.data<int>()),
-            const_cast<int *>(ori_seq_lens_decoder.data<int>()),
-            const_cast<int *>(seq_lens_encoder.data<int>()),
-            const_cast<int *>(seq_lens_decoder.data<int>()),
-            const_cast<int *>(block_tables.data<int>()),
-            const_cast<int *>(free_list.data<int>()),
-            const_cast<int *>(free_list_len.data<int>()),
-            const_cast<int64_t *>(input_ids.data<int64_t>()),
-            const_cast<int64_t *>(pre_ids.data<int64_t>()),
-            const_cast<int64_t *>(step_idx.data<int64_t>()),
-            const_cast<int *>(encoder_block_lens.data<int>()),
-            const_cast<int *>(used_list_len.data<int>()),
-            next_tokens.data<int64_t>(),
-            first_token_ids.data<int64_t>(),
-            bsz,
-            block_num_per_seq,
-            length,
-            pre_id_length);
-#ifdef DEBUG_STEP
-        cudaDeviceSynchronize();
-#endif
-    }
+  }
 }
 
 PD_BUILD_STATIC_OP(speculate_step_system_cache)
diff --git a/custom_ops/gpu_ops/speculate_decoding/speculate_update_input_ids_cpu.cc b/custom_ops/gpu_ops/speculate_decoding/speculate_update_input_ids_cpu.cc
index 1ada2ae332..a0437bdd2b 100644
--- a/custom_ops/gpu_ops/speculate_decoding/speculate_update_input_ids_cpu.cc
+++ b/custom_ops/gpu_ops/speculate_decoding/speculate_update_input_ids_cpu.cc
@@ -22,14 +22,14 @@ void UpdateInputIdsCPU(const paddle::Tensor& input_ids_cpu,
                        const std::vector<int64_t>& task_input_ids,
                        const int bid,
                        const int max_seq_len) {
-    int64_t* input_ids_cpu_data =
-        const_cast<int64_t*>(input_ids_cpu.data<int64_t>());
-    // printf("Input len is %d\n", task_input_ids.size());
+  int64_t* input_ids_cpu_data =
+      const_cast<int64_t*>(input_ids_cpu.data<int64_t>());
+  // printf("Input len is %d\n", task_input_ids.size());
 
-    for (int i = 0; i < task_input_ids.size(); i++) {
-        // printf("%lld\n", task_input_ids[i]);
-        input_ids_cpu_data[bid * max_seq_len + i] = task_input_ids[i];
-    }
+  for (int i = 0; i < task_input_ids.size(); i++) {
+    // printf("%lld\n", task_input_ids[i]);
+    input_ids_cpu_data[bid * max_seq_len + i] = task_input_ids[i];
+  }
 }
 
 PD_BUILD_STATIC_OP(speculate_update_input_ids_cpu)
diff --git a/custom_ops/gpu_ops/step.cu b/custom_ops/gpu_ops/step.cu
index 90b95c983d..8b5057a34b 100644
--- a/custom_ops/gpu_ops/step.cu
+++ b/custom_ops/gpu_ops/step.cu
@@ -17,15 +17,15 @@
 __device__ bool in_need_block_list(const int &qid,
                                    int *need_block_list,
                                    const int &need_block_len) {
-    bool res = false;
-    for (int i = 0; i < need_block_len; i++) {
-        if (qid == need_block_list[i]) {
-            res = true;
-            need_block_list[i] = -1;
-            break;
-        }
+  bool res = false;
+  for (int i = 0; i < need_block_len; i++) {
+    if (qid == need_block_list[i]) {
+      res = true;
+      need_block_list[i] = -1;
+      break;
     }
-    return res;
+  }
+  return res;
 }
 
 __global__ void free_and_dispatch_block(bool *stop_flags,
@@ -48,178 +48,177 @@ __global__ void free_and_dispatch_block(bool *stop_flags,
                                         const int block_size,
                                         const int block_num_per_seq,
                                         const int max_decoder_block_num) {
-    typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
-    __shared__ bool step_max_block_flag;
-    __shared__ int in_need_block_list_len;
-    const int tid = threadIdx.x;
-    if (tid < bsz) {
-        if (tid == 0) {
-            step_max_block_flag = false;
-            in_need_block_list_len = 0;
-        }
-        int *block_table_now = block_tables + tid * block_num_per_seq;
-        if (stop_flags[tid] && !is_block_step[tid]) {
-            // 回收block块
-            first_token_ids[tid] = -1;
-            const int encoder_block_len = encoder_block_lens[tid];
-            const int decoder_used_len = used_list_len[tid];
-            if (decoder_used_len > 0) {
-                const int ori_free_list_len =
-                    atomicAdd(free_list_len, decoder_used_len);
-#ifdef DEBUG_STEP
-                printf(
-                    "free block seq_id: %d, free block num: %d, "
-                    "encoder_block_len: %d, ori_free_list_len: %d\n",
-                    tid,
-                    decoder_used_len,
-                    encoder_block_len,
-                    ori_free_list_len);
-#endif
-                for (int i = 0; i < decoder_used_len; i++) {
-                    free_list[ori_free_list_len + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                encoder_block_lens[tid] = 0;
-                used_list_len[tid] = 0;
-            }
-        } else if (seq_lens_this_time[tid] != 0 &&
-                   block_table_now[seq_lens_decoder[tid] / block_size] == -1) {
-            // 统计需要分配block的位置和总数
-#ifdef DEBUG_STEP
-            printf("step seq_id:%d, ##### pin 1 #####\n", tid);
-#endif
-            const int ori_need_block_len = atomicAdd(need_block_len, 1);
-            need_block_list[ori_need_block_len] = tid;
-#ifdef DEBUG_STEP
-            printf("seq_id: %d need block\n", tid);
-#endif
-        }
-    }
-#ifdef DEBUG_STEP
-    printf("step seq_id:%d, ##### pin 2 #####\n", tid);
-#endif
-    __syncthreads();
-
-    while (need_block_len[0] > free_list_len[0]) {
-        if (tid == 0) {
-            printf("need_block_len: %d, free_list_len: %d\n",
-                   need_block_len[0],
-                   free_list_len[0]);
-        }
-        // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
-        const int used_block_num =
-            tid < bsz && !is_block_step[tid] &&
-                    (step_max_block_flag ||
-                     used_list_len[tid] != max_decoder_block_num)
-                ? used_list_len[tid]
-                : 0;
-        cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
-        kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
-
-        if (tid == 0) {
-            if (kv_pair.value == 0) {
-                step_max_block_flag = true;
-            } else {
-                const int encoder_block_len = encoder_block_lens[kv_pair.key];
-                printf("step max_id: %d, max_num: %d, encoder_block_len: %d\n",
-                       kv_pair.key,
-                       kv_pair.value,
-                       encoder_block_len);
-                int *block_table_now =
-                    block_tables + kv_pair.key * block_num_per_seq;
-                for (int i = 0; i < kv_pair.value; i++) {
-                    free_list[free_list_len[0] + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                step_block_list[step_len[0]] = kv_pair.key;
-                if (in_need_block_list(
-                        kv_pair.key,
-                        need_block_list,
-                        need_block_len[0] + in_need_block_list_len)) {
-                    need_block_len[0] -= 1;
-                    in_need_block_list_len += 1;
-                }
-                step_len[0] += 1;
-                free_list_len[0] += kv_pair.value;
-                stop_flags[kv_pair.key] = true;
-                is_block_step[kv_pair.key] = true;
-                seq_lens_this_time[kv_pair.key] = 0;
-                seq_lens_decoder[kv_pair.key] = 0;
-            }
-        }
-        __syncthreads();
-    }
-#ifdef DEBUG_STEP
-    printf("step seq_id:%d, ##### pin 3 #####\n", tid);
-#endif
-    // 为需要block的位置分配block，每个位置分配一个block
-    if (tid < need_block_len[0] + in_need_block_list_len) {
-        const int need_block_id = need_block_list[tid];
-        if (need_block_id != -1) {
-            if (!stop_flags[need_block_id]) {
-                // 如果需要的位置正好是上一步中被释放的位置，不做处理
-                used_list_len[need_block_id] += 1;
-                const int ori_free_list_len = atomicSub(free_list_len, 1);
-                int *block_table_now =
-                    block_tables + need_block_id * block_num_per_seq;
-                block_table_now[seq_lens_decoder[need_block_id] / block_size] =
-                    free_list[ori_free_list_len - 1];
-            }
-            need_block_list[tid] = -1;
-        }
-    }
-    __syncthreads();
-#ifdef DEBUG_STEP
-    printf("step seq_id:%d, ##### pin 4 #####\n", tid);
-#endif
-    // 计算可以复原的query id
+  typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ bool step_max_block_flag;
+  __shared__ int in_need_block_list_len;
+  const int tid = threadIdx.x;
+  if (tid < bsz) {
     if (tid == 0) {
-        int ori_free_list_len = free_list_len[0];
-        int ori_step_len = step_len[0];
-        if (ori_step_len > 0) {
-            int ori_step_block_id = step_block_list[ori_step_len - 1];
-            int tmp_used_len = used_list_len[ori_step_block_id];
-            // 比之前调度时多分配一个block，防止马上恢复刚调度的query(比如回收的seq_id在need_block_list中）
-            const int max_decoder_block_num_this_seq =
-                max_decoder_block_num - encoder_block_lens[ori_step_block_id];
-            int used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
-                            ? tmp_used_len + 1
-                            : max_decoder_block_num_this_seq;
-#ifdef DEBUG_STEP
-            printf("#### ori_step_len:%d, ori_free_list_len:%d, used_len:%d  \n",
-                    ori_step_len, ori_free_list_len, used_len);
-#endif
-            while (ori_step_len > 0 && ori_free_list_len >= used_len) {
-                printf(
-                    "recover seq_id: %d, free_list_len: %d, "
-                    "used_list_len: %d\n",
-                    ori_step_block_id,
-                    ori_free_list_len,
-                    used_len);
-                recover_block_list[recover_len[0]] = ori_step_block_id;
-                is_block_step[ori_step_block_id] = false;
-                used_list_len[ori_step_block_id] = used_len;
-                ori_free_list_len -= used_len;
-                step_block_list[ori_step_len - 1] = -1;
-                step_len[0] -= 1;
-                recover_len[0] += 1;
-                ori_step_len = step_len[0];
-                if (ori_step_len > 0) {
-                    ori_step_block_id = step_block_list[ori_step_len - 1];
-                    tmp_used_len = used_list_len[ori_step_block_id];
-                    used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
-                                ? tmp_used_len + 1
-                                : max_decoder_block_num_this_seq;
-                }
-            }
-        }
-        need_block_len[0] = 0;
+      step_max_block_flag = false;
+      in_need_block_list_len = 0;
     }
+    int *block_table_now = block_tables + tid * block_num_per_seq;
+    if (stop_flags[tid] && !is_block_step[tid]) {
+      // 回收block块
+      first_token_ids[tid] = -1;
+      const int encoder_block_len = encoder_block_lens[tid];
+      const int decoder_used_len = used_list_len[tid];
+      if (decoder_used_len > 0) {
+        const int ori_free_list_len =
+            atomicAdd(free_list_len, decoder_used_len);
 #ifdef DEBUG_STEP
-    printf("step seq_id:%d, ##### pin 5 #####\n", tid);
+        printf(
+            "free block seq_id: %d, free block num: %d, "
+            "encoder_block_len: %d, ori_free_list_len: %d\n",
+            tid,
+            decoder_used_len,
+            encoder_block_len,
+            ori_free_list_len);
+#endif
+        for (int i = 0; i < decoder_used_len; i++) {
+          free_list[ori_free_list_len + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        encoder_block_lens[tid] = 0;
+        used_list_len[tid] = 0;
+      }
+    } else if (seq_lens_this_time[tid] != 0 &&
+               block_table_now[seq_lens_decoder[tid] / block_size] == -1) {
+      // 统计需要分配block的位置和总数
+#ifdef DEBUG_STEP
+      printf("step seq_id:%d, ##### pin 1 #####\n", tid);
+#endif
+      const int ori_need_block_len = atomicAdd(need_block_len, 1);
+      need_block_list[ori_need_block_len] = tid;
+#ifdef DEBUG_STEP
+      printf("seq_id: %d need block\n", tid);
+#endif
+    }
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 2 #####\n", tid);
+#endif
+  __syncthreads();
+
+  while (need_block_len[0] > free_list_len[0]) {
+    if (tid == 0) {
+      printf("need_block_len: %d, free_list_len: %d\n",
+             need_block_len[0],
+             free_list_len[0]);
+    }
+    // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
+    const int used_block_num =
+        tid < bsz && !is_block_step[tid] &&
+                (step_max_block_flag ||
+                 used_list_len[tid] != max_decoder_block_num)
+            ? used_list_len[tid]
+            : 0;
+    cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
+    kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
+
+    if (tid == 0) {
+      if (kv_pair.value == 0) {
+        step_max_block_flag = true;
+      } else {
+        const int encoder_block_len = encoder_block_lens[kv_pair.key];
+        printf("step max_id: %d, max_num: %d, encoder_block_len: %d\n",
+               kv_pair.key,
+               kv_pair.value,
+               encoder_block_len);
+        int *block_table_now = block_tables + kv_pair.key * block_num_per_seq;
+        for (int i = 0; i < kv_pair.value; i++) {
+          free_list[free_list_len[0] + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        step_block_list[step_len[0]] = kv_pair.key;
+        if (in_need_block_list(kv_pair.key,
+                               need_block_list,
+                               need_block_len[0] + in_need_block_list_len)) {
+          need_block_len[0] -= 1;
+          in_need_block_list_len += 1;
+        }
+        step_len[0] += 1;
+        free_list_len[0] += kv_pair.value;
+        stop_flags[kv_pair.key] = true;
+        is_block_step[kv_pair.key] = true;
+        seq_lens_this_time[kv_pair.key] = 0;
+        seq_lens_decoder[kv_pair.key] = 0;
+      }
+    }
+    __syncthreads();
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 3 #####\n", tid);
+#endif
+  // 为需要block的位置分配block，每个位置分配一个block
+  if (tid < need_block_len[0] + in_need_block_list_len) {
+    const int need_block_id = need_block_list[tid];
+    if (need_block_id != -1) {
+      if (!stop_flags[need_block_id]) {
+        // 如果需要的位置正好是上一步中被释放的位置，不做处理
+        used_list_len[need_block_id] += 1;
+        const int ori_free_list_len = atomicSub(free_list_len, 1);
+        int *block_table_now = block_tables + need_block_id * block_num_per_seq;
+        block_table_now[seq_lens_decoder[need_block_id] / block_size] =
+            free_list[ori_free_list_len - 1];
+      }
+      need_block_list[tid] = -1;
+    }
+  }
+  __syncthreads();
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 4 #####\n", tid);
+#endif
+  // 计算可以复原的query id
+  if (tid == 0) {
+    int ori_free_list_len = free_list_len[0];
+    int ori_step_len = step_len[0];
+    if (ori_step_len > 0) {
+      int ori_step_block_id = step_block_list[ori_step_len - 1];
+      int tmp_used_len = used_list_len[ori_step_block_id];
+      // 比之前调度时多分配一个block，防止马上恢复刚调度的query(比如回收的seq_id在need_block_list中）
+      const int max_decoder_block_num_this_seq =
+          max_decoder_block_num - encoder_block_lens[ori_step_block_id];
+      int used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
+                         ? tmp_used_len + 1
+                         : max_decoder_block_num_this_seq;
+#ifdef DEBUG_STEP
+      printf("#### ori_step_len:%d, ori_free_list_len:%d, used_len:%d  \n",
+             ori_step_len,
+             ori_free_list_len,
+             used_len);
+#endif
+      while (ori_step_len > 0 && ori_free_list_len >= used_len) {
+        printf(
+            "recover seq_id: %d, free_list_len: %d, "
+            "used_list_len: %d\n",
+            ori_step_block_id,
+            ori_free_list_len,
+            used_len);
+        recover_block_list[recover_len[0]] = ori_step_block_id;
+        is_block_step[ori_step_block_id] = false;
+        used_list_len[ori_step_block_id] = used_len;
+        ori_free_list_len -= used_len;
+        step_block_list[ori_step_len - 1] = -1;
+        step_len[0] -= 1;
+        recover_len[0] += 1;
+        ori_step_len = step_len[0];
+        if (ori_step_len > 0) {
+          ori_step_block_id = step_block_list[ori_step_len - 1];
+          tmp_used_len = used_list_len[ori_step_block_id];
+          used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
+                         ? tmp_used_len + 1
+                         : max_decoder_block_num_this_seq;
+        }
+      }
+    }
+    need_block_len[0] = 0;
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 5 #####\n", tid);
 #endif
 }
 
@@ -245,58 +244,57 @@ __global__ void recover_block(int *recover_block_list,  // [bsz]
                               const int block_num_per_seq,
                               const int length,
                               const int pre_id_length) {
-    const int bid = blockIdx.x;
-    const int tid = threadIdx.x;
-    __shared__ int ori_free_list_len;
-    if (bid < recover_len[0]) {
-        const int recover_id = recover_block_list[bid];
-        const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
-        const int step_idx_now = step_idx[recover_id];
-        const int seq_len = ori_seq_len_encoder + step_idx_now;
-        const int encoder_block_len = encoder_block_lens[recover_id];
-        const int decoder_used_len = used_list_len[recover_id];
-        int *block_table_now = block_tables + recover_id * block_num_per_seq;
-        int64_t *input_ids_now = input_ids + recover_id * length;
-        int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
-        if (tid == 0) {
-            seq_lens_this_time[recover_id] = seq_len;
-            seq_lens_encoder[recover_id] = seq_len;
-            stop_flags[recover_id] = false;
-            input_ids_now[ori_seq_len_encoder + step_idx_now - 1] =
-                next_tokens[recover_id];  // next tokens
-            input_ids_now[0] =
-                first_token_ids[recover_id];  // set first prompt token
-            const int ori_free_list_len_tid0 =
-                atomicSub(free_list_len, decoder_used_len);
-            ori_free_list_len = ori_free_list_len_tid0;
+  const int bid = blockIdx.x;
+  const int tid = threadIdx.x;
+  __shared__ int ori_free_list_len;
+  if (bid < recover_len[0]) {
+    const int recover_id = recover_block_list[bid];
+    const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
+    const int step_idx_now = step_idx[recover_id];
+    const int seq_len = ori_seq_len_encoder + step_idx_now;
+    const int encoder_block_len = encoder_block_lens[recover_id];
+    const int decoder_used_len = used_list_len[recover_id];
+    int *block_table_now = block_tables + recover_id * block_num_per_seq;
+    int64_t *input_ids_now = input_ids + recover_id * length;
+    int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
+    if (tid == 0) {
+      seq_lens_this_time[recover_id] = seq_len;
+      seq_lens_encoder[recover_id] = seq_len;
+      stop_flags[recover_id] = false;
+      input_ids_now[ori_seq_len_encoder + step_idx_now - 1] =
+          next_tokens[recover_id];                     // next tokens
+      input_ids_now[0] = first_token_ids[recover_id];  // set first prompt token
+      const int ori_free_list_len_tid0 =
+          atomicSub(free_list_len, decoder_used_len);
+      ori_free_list_len = ori_free_list_len_tid0;
 #ifdef DEBUG_STEP
-            printf(
-                "seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, "
-                "seq_len: %d, ori_free_list_len_tid0: %d, "
-                "ori_free_list_len: %d\n",
-                recover_id,
-                ori_seq_len_encoder,
-                step_idx_now,
-                seq_len,
-                ori_free_list_len_tid0,
-                ori_free_list_len);
+      printf(
+          "seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, "
+          "seq_len: %d, ori_free_list_len_tid0: %d, "
+          "ori_free_list_len: %d\n",
+          recover_id,
+          ori_seq_len_encoder,
+          step_idx_now,
+          seq_len,
+          ori_free_list_len_tid0,
+          ori_free_list_len);
 #endif
-        }
-        __syncthreads();
-        // 恢复block table
-        for (int i = tid; i < decoder_used_len; i += blockDim.x) {
-            block_table_now[encoder_block_len + i] =
-                free_list[ori_free_list_len - i - 1];
-        }
-        // 恢复input_ids
-        for (int i = tid; i < step_idx_now - 1; i += blockDim.x) {
-            input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
-        }
     }
+    __syncthreads();
+    // 恢复block table
+    for (int i = tid; i < decoder_used_len; i += blockDim.x) {
+      block_table_now[encoder_block_len + i] =
+          free_list[ori_free_list_len - i - 1];
+    }
+    // 恢复input_ids
+    for (int i = tid; i < step_idx_now - 1; i += blockDim.x) {
+      input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
+    }
+  }
 
-    if (bid == 0 && tid == 0) {
-        recover_len[0] = 0;
-    }
+  if (bid == 0 && tid == 0) {
+    recover_len[0] = 0;
+  }
 }
 
 void StepPaddle(const paddle::Tensor &stop_flags,
@@ -324,84 +322,86 @@ void StepPaddle(const paddle::Tensor &stop_flags,
                 const int block_size,
                 const int encoder_decoder_block_num) {
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
-    auto dev_ctx = static_cast<const phi::CustomContext*>(paddle::experimental::DeviceContextPool::Instance().Get(seq_lens_this_time.place()));
-    auto cu_stream = dev_ctx->stream();
+  auto dev_ctx = static_cast<const phi::CustomContext *>(
+      paddle::experimental::DeviceContextPool::Instance().Get(
+          seq_lens_this_time.place()));
+  auto cu_stream = dev_ctx->stream();
 #else
-    auto cu_stream = seq_lens_this_time.stream();
+  auto cu_stream = seq_lens_this_time.stream();
 #endif
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    const int length = input_ids.shape()[1];
-    const int pre_id_length = pre_ids.shape()[1];
-    constexpr int BlockSize = 256;  // bsz <= 256
-    const int max_decoder_block_num = length / block_size;
-    // const int max_decoder_block_num = 2048 / block_size -
-    // encoder_decoder_block_num;
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  const int length = input_ids.shape()[1];
+  const int pre_id_length = pre_ids.shape()[1];
+  constexpr int BlockSize = 256;  // bsz <= 256
+  const int max_decoder_block_num = length / block_size;
+  // const int max_decoder_block_num = 2048 / block_size -
+  // encoder_decoder_block_num;
 #ifdef DEBUG_STEP
-    printf(
-        "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
-        "%d\n",
+  printf(
+      "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
+      "%d\n",
+      bsz,
+      block_num_per_seq,
+      length,
+      max_decoder_block_num);
+#endif
+  free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
+      const_cast<bool *>(stop_flags.data<bool>()),
+      const_cast<int *>(seq_lens_this_time.data<int>()),
+      const_cast<int *>(seq_lens_decoder.data<int>()),
+      const_cast<int *>(block_tables.data<int>()),
+      const_cast<int *>(encoder_block_lens.data<int>()),
+      const_cast<bool *>(is_block_step.data<bool>()),
+      const_cast<int *>(step_block_list.data<int>()),
+      const_cast<int *>(step_lens.data<int>()),
+      const_cast<int *>(recover_block_list.data<int>()),
+      const_cast<int *>(recover_lens.data<int>()),
+      const_cast<int *>(need_block_list.data<int>()),
+      const_cast<int *>(need_block_len.data<int>()),
+      const_cast<int *>(used_list_len.data<int>()),
+      const_cast<int *>(free_list.data<int>()),
+      const_cast<int *>(free_list_len.data<int>()),
+      const_cast<int64_t *>(first_token_ids.data<int64_t>()),
+      bsz,
+      block_size,
+      block_num_per_seq,
+      max_decoder_block_num);
+#ifdef DEBUG_STEP
+  cudaDeviceSynchronize();
+#endif
+  auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
+  const int grid_size = cpu_recover_lens.data<int>()[0];
+#ifdef DEBUG_STEP
+  printf("grid_size2 %d\n", grid_size);
+#endif
+  if (grid_size > 0) {
+    recover_block<<<grid_size, BlockSize, 0, cu_stream>>>(
+        const_cast<int *>(recover_block_list.data<int>()),
+        const_cast<int *>(recover_lens.data<int>()),
+        const_cast<bool *>(stop_flags.data<bool>()),
+        const_cast<int *>(seq_lens_this_time.data<int>()),
+        const_cast<int *>(ori_seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_decoder.data<int>()),
+        const_cast<int *>(block_tables.data<int>()),
+        const_cast<int *>(free_list.data<int>()),
+        const_cast<int *>(free_list_len.data<int>()),
+        const_cast<int64_t *>(input_ids.data<int64_t>()),
+        const_cast<int64_t *>(pre_ids.data<int64_t>()),
+        const_cast<int64_t *>(step_idx.data<int64_t>()),
+        const_cast<int *>(encoder_block_lens.data<int>()),
+        const_cast<int *>(used_list_len.data<int>()),
+        next_tokens.data<int64_t>(),
+        first_token_ids.data<int64_t>(),
         bsz,
         block_num_per_seq,
         length,
-        max_decoder_block_num);
-#endif
-    free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
-        const_cast<bool *>(stop_flags.data<bool>()),
-        const_cast<int *>(seq_lens_this_time.data<int>()),
-        const_cast<int *>(seq_lens_decoder.data<int>()),
-        const_cast<int *>(block_tables.data<int>()),
-        const_cast<int *>(encoder_block_lens.data<int>()),
-        const_cast<bool *>(is_block_step.data<bool>()),
-        const_cast<int *>(step_block_list.data<int>()),
-        const_cast<int *>(step_lens.data<int>()),
-        const_cast<int *>(recover_block_list.data<int>()),
-        const_cast<int *>(recover_lens.data<int>()),
-        const_cast<int *>(need_block_list.data<int>()),
-        const_cast<int *>(need_block_len.data<int>()),
-        const_cast<int *>(used_list_len.data<int>()),
-        const_cast<int *>(free_list.data<int>()),
-        const_cast<int *>(free_list_len.data<int>()),
-        const_cast<int64_t *>(first_token_ids.data<int64_t>()),
-        bsz,
-        block_size,
-        block_num_per_seq,
-        max_decoder_block_num);
+        pre_id_length);
 #ifdef DEBUG_STEP
     cudaDeviceSynchronize();
 #endif
-    auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
-    const int grid_size = cpu_recover_lens.data<int>()[0];
-#ifdef DEBUG_STEP
-    printf("grid_size2 %d\n", grid_size);
-#endif
-    if (grid_size > 0) {
-        recover_block<<<grid_size, BlockSize, 0, cu_stream>>>(
-            const_cast<int *>(recover_block_list.data<int>()),
-            const_cast<int *>(recover_lens.data<int>()),
-            const_cast<bool *>(stop_flags.data<bool>()),
-            const_cast<int *>(seq_lens_this_time.data<int>()),
-            const_cast<int *>(ori_seq_lens_encoder.data<int>()),
-            const_cast<int *>(seq_lens_encoder.data<int>()),
-            const_cast<int *>(seq_lens_decoder.data<int>()),
-            const_cast<int *>(block_tables.data<int>()),
-            const_cast<int *>(free_list.data<int>()),
-            const_cast<int *>(free_list_len.data<int>()),
-            const_cast<int64_t *>(input_ids.data<int64_t>()),
-            const_cast<int64_t *>(pre_ids.data<int64_t>()),
-            const_cast<int64_t *>(step_idx.data<int64_t>()),
-            const_cast<int *>(encoder_block_lens.data<int>()),
-            const_cast<int *>(used_list_len.data<int>()),
-            next_tokens.data<int64_t>(),
-            first_token_ids.data<int64_t>(),
-            bsz,
-            block_num_per_seq,
-            length,
-            pre_id_length);
-#ifdef DEBUG_STEP
-        cudaDeviceSynchronize();
-#endif
-    }
+  }
 }
 
 PD_BUILD_STATIC_OP(step_paddle)
diff --git a/custom_ops/gpu_ops/step_reschedule.cu b/custom_ops/gpu_ops/step_reschedule.cu
index 8475b317e5..1321b6dd25 100644
--- a/custom_ops/gpu_ops/step_reschedule.cu
+++ b/custom_ops/gpu_ops/step_reschedule.cu
@@ -20,170 +20,167 @@
 __device__ bool in_need_block_list_schedule(const int &qid,
                                             int *need_block_list,
                                             const int &need_block_len) {
-    bool res = false;
-    for (int i = 0; i < need_block_len; i++) {
-        if (qid == need_block_list[i]) {
-            res = true;
-            need_block_list[i] = -1;
-            break;
-        }
+  bool res = false;
+  for (int i = 0; i < need_block_len; i++) {
+    if (qid == need_block_list[i]) {
+      res = true;
+      need_block_list[i] = -1;
+      break;
     }
-    return res;
+  }
+  return res;
 }
 
 __global__ void free_and_reschedule(bool *stop_flags,
-                                   int *seq_lens_this_time,
-                                   int *seq_lens_decoder,
-                                   int *block_tables,
-                                   int *encoder_block_lens,
-                                   bool *is_block_step,
-                                   int *step_block_list,  // [bsz]
-                                   int *step_len,
-                                   int *recover_block_list,
-                                   int *recover_len,
-                                   int *need_block_list,
-                                   int *need_block_len,
-                                   int *used_list_len,
-                                   int *free_list,
-                                   int *free_list_len,
-                                   int64_t *first_token_ids,
-                                   const int bsz,
-                                   const int block_size,
-                                   const int block_num_per_seq,
-                                   const int max_decoder_block_num) {
-    typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
-    __shared__ bool step_max_block_flag;
-    __shared__ int in_need_block_list_len;
-    const int tid = threadIdx.x;
-    if (tid < bsz) {
-        if (tid == 0) {
-            step_max_block_flag = false;
-            in_need_block_list_len = 0;
-        }
-        int *block_table_now = block_tables + tid * block_num_per_seq;
-        if (stop_flags[tid]) {
-            // 回收block块
-            first_token_ids[tid] = -1;
-            const int encoder_block_len = encoder_block_lens[tid];
-            const int decoder_used_len = used_list_len[tid];
-            if (decoder_used_len > 0) {
-                const int ori_free_list_len =
-                    atomicAdd(free_list_len, decoder_used_len);
-#ifdef DEBUG_STEP
-                printf(
-                    "free block seq_id: %d, free block num: %d, "
-                    "encoder_block_len: %d, ori_free_list_len: %d\n",
-                    tid,
-                    decoder_used_len,
-                    encoder_block_len,
-                    ori_free_list_len);
-#endif
-                for (int i = 0; i < decoder_used_len; i++) {
-                    free_list[ori_free_list_len + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                encoder_block_lens[tid] = 0;
-                used_list_len[tid] = 0;
-            }
-        } else if (seq_lens_this_time[tid] != 0 &&
-                   block_table_now[seq_lens_decoder[tid] / block_size] == -1) {
-            // 统计需要分配block的位置和总数
-#ifdef DEBUG_STEP
-            printf("step seq_id:%d, ##### pin 1 #####\n", tid);
-#endif
-            const int ori_need_block_len = atomicAdd(need_block_len, 1);
-            need_block_list[ori_need_block_len] = tid;
-#ifdef DEBUG_STEP
-            printf("seq_id: %d need block\n", tid);
-#endif
-        }
-    }
-#ifdef DEBUG_STEP
-    printf("step seq_id:%d, ##### pin 2 #####\n", tid);
-#endif
-    __syncthreads();
-
-    // 调度block，直到满足need_block_len
-    while (need_block_len[0] > free_list_len[0]) {
-        if (tid == 0) {
-            printf("need_block_len: %d, free_list_len: %d\n",
-                   need_block_len[0],
-                   free_list_len[0]);
-        }
-        // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
-        const int used_block_num =
-            tid < bsz ? used_list_len[tid] : 0;
-        cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
-        kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
-        if (tid == 0) {
-            if (kv_pair.value == 0) {
-                step_max_block_flag = true;
-            } else {
-                const int encoder_block_len = encoder_block_lens[kv_pair.key];
-                printf("step max_id: %d, max_num: %d, encoder_block_len: %d\n",
-                       kv_pair.key,
-                       kv_pair.value,
-                       encoder_block_len);
-                int *block_table_now =
-                    block_tables + kv_pair.key * block_num_per_seq;
-                // 回收调度位的block
-                for (int i = 0; i < kv_pair.value; i++) {
-                    free_list[free_list_len[0] + i] =
-                        block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                step_block_list[step_len[0]] = kv_pair.key;
-                // 如果调度位置本次也需要block，对应的处理
-                if (in_need_block_list_schedule(
-                        kv_pair.key,
-                        need_block_list,
-                        need_block_len[0] + in_need_block_list_len)) {
-                    need_block_len[0] -= 1;
-                    in_need_block_list_len += 1;
-                }
-                step_len[0] += 1;
-                free_list_len[0] += kv_pair.value;
-                stop_flags[kv_pair.key] = true;
-                seq_lens_this_time[kv_pair.key] = 0;
-                seq_lens_decoder[kv_pair.key] = 0;
-                encoder_block_lens[kv_pair.key] = 0;
-                used_list_len[kv_pair.key] = 0;
-                printf(
-                    "free block seq_id: %d, free block num: %d, "
-                    "now_free_list_len: %d\n",
-                    (int)kv_pair.key,
-                    (int)kv_pair.value,
-                    (int)free_list_len[0]);
-            }
-        }
-        __syncthreads();
-    }
-#ifdef DEBUG_STEP
-    printf("step seq_id:%d, ##### pin 3 #####\n", tid);
-#endif
-    // 为需要block的位置分配block，每个位置分配一个block
-    if (tid < need_block_len[0] + in_need_block_list_len) {
-        const int need_block_id = need_block_list[tid];
-        if (need_block_id != -1) {
-            if (!stop_flags[need_block_id]) {
-                // 如果需要的位置正好是上一步中被释放的位置，不做处理
-                used_list_len[need_block_id] += 1;
-                const int ori_free_list_len = atomicSub(free_list_len, 1);
-                int *block_table_now =
-                    block_tables + need_block_id * block_num_per_seq;
-                block_table_now[seq_lens_decoder[need_block_id] / block_size] =
-                    free_list[ori_free_list_len - 1];
-            }
-            need_block_list[tid] = -1;
-        }
-    }
-    __syncthreads();
-    // reset need_block_len
+                                    int *seq_lens_this_time,
+                                    int *seq_lens_decoder,
+                                    int *block_tables,
+                                    int *encoder_block_lens,
+                                    bool *is_block_step,
+                                    int *step_block_list,  // [bsz]
+                                    int *step_len,
+                                    int *recover_block_list,
+                                    int *recover_len,
+                                    int *need_block_list,
+                                    int *need_block_len,
+                                    int *used_list_len,
+                                    int *free_list,
+                                    int *free_list_len,
+                                    int64_t *first_token_ids,
+                                    const int bsz,
+                                    const int block_size,
+                                    const int block_num_per_seq,
+                                    const int max_decoder_block_num) {
+  typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ bool step_max_block_flag;
+  __shared__ int in_need_block_list_len;
+  const int tid = threadIdx.x;
+  if (tid < bsz) {
     if (tid == 0) {
-        need_block_len[0] = 0;
+      step_max_block_flag = false;
+      in_need_block_list_len = 0;
     }
+    int *block_table_now = block_tables + tid * block_num_per_seq;
+    if (stop_flags[tid]) {
+      // 回收block块
+      first_token_ids[tid] = -1;
+      const int encoder_block_len = encoder_block_lens[tid];
+      const int decoder_used_len = used_list_len[tid];
+      if (decoder_used_len > 0) {
+        const int ori_free_list_len =
+            atomicAdd(free_list_len, decoder_used_len);
+#ifdef DEBUG_STEP
+        printf(
+            "free block seq_id: %d, free block num: %d, "
+            "encoder_block_len: %d, ori_free_list_len: %d\n",
+            tid,
+            decoder_used_len,
+            encoder_block_len,
+            ori_free_list_len);
+#endif
+        for (int i = 0; i < decoder_used_len; i++) {
+          free_list[ori_free_list_len + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        encoder_block_lens[tid] = 0;
+        used_list_len[tid] = 0;
+      }
+    } else if (seq_lens_this_time[tid] != 0 &&
+               block_table_now[seq_lens_decoder[tid] / block_size] == -1) {
+      // 统计需要分配block的位置和总数
+#ifdef DEBUG_STEP
+      printf("step seq_id:%d, ##### pin 1 #####\n", tid);
+#endif
+      const int ori_need_block_len = atomicAdd(need_block_len, 1);
+      need_block_list[ori_need_block_len] = tid;
+#ifdef DEBUG_STEP
+      printf("seq_id: %d need block\n", tid);
+#endif
+    }
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 2 #####\n", tid);
+#endif
+  __syncthreads();
+
+  // 调度block，直到满足need_block_len
+  while (need_block_len[0] > free_list_len[0]) {
+    if (tid == 0) {
+      printf("need_block_len: %d, free_list_len: %d\n",
+             need_block_len[0],
+             free_list_len[0]);
+    }
+    // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
+    const int used_block_num = tid < bsz ? used_list_len[tid] : 0;
+    cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
+    kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
+    if (tid == 0) {
+      if (kv_pair.value == 0) {
+        step_max_block_flag = true;
+      } else {
+        const int encoder_block_len = encoder_block_lens[kv_pair.key];
+        printf("step max_id: %d, max_num: %d, encoder_block_len: %d\n",
+               kv_pair.key,
+               kv_pair.value,
+               encoder_block_len);
+        int *block_table_now = block_tables + kv_pair.key * block_num_per_seq;
+        // 回收调度位的block
+        for (int i = 0; i < kv_pair.value; i++) {
+          free_list[free_list_len[0] + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        step_block_list[step_len[0]] = kv_pair.key;
+        // 如果调度位置本次也需要block，对应的处理
+        if (in_need_block_list_schedule(
+                kv_pair.key,
+                need_block_list,
+                need_block_len[0] + in_need_block_list_len)) {
+          need_block_len[0] -= 1;
+          in_need_block_list_len += 1;
+        }
+        step_len[0] += 1;
+        free_list_len[0] += kv_pair.value;
+        stop_flags[kv_pair.key] = true;
+        seq_lens_this_time[kv_pair.key] = 0;
+        seq_lens_decoder[kv_pair.key] = 0;
+        encoder_block_lens[kv_pair.key] = 0;
+        used_list_len[kv_pair.key] = 0;
+        printf(
+            "free block seq_id: %d, free block num: %d, "
+            "now_free_list_len: %d\n",
+            (int)kv_pair.key,
+            (int)kv_pair.value,
+            (int)free_list_len[0]);
+      }
+    }
+    __syncthreads();
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 3 #####\n", tid);
+#endif
+  // 为需要block的位置分配block，每个位置分配一个block
+  if (tid < need_block_len[0] + in_need_block_list_len) {
+    const int need_block_id = need_block_list[tid];
+    if (need_block_id != -1) {
+      if (!stop_flags[need_block_id]) {
+        // 如果需要的位置正好是上一步中被释放的位置，不做处理
+        used_list_len[need_block_id] += 1;
+        const int ori_free_list_len = atomicSub(free_list_len, 1);
+        int *block_table_now = block_tables + need_block_id * block_num_per_seq;
+        block_table_now[seq_lens_decoder[need_block_id] / block_size] =
+            free_list[ori_free_list_len - 1];
+      }
+      need_block_list[tid] = -1;
+    }
+  }
+  __syncthreads();
+  // reset need_block_len
+  if (tid == 0) {
+    need_block_len[0] = 0;
+  }
 }
 
 // 为不修改接口调用方式，入参暂不改变
@@ -211,101 +208,109 @@ void Schedule(const paddle::Tensor &stop_flags,
               const paddle::Tensor &first_token_ids,
               const int block_size,
               const int encoder_decoder_block_num) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    const int length = input_ids.shape()[1];
-    const int pre_id_length = pre_ids.shape()[1];
-    constexpr int BlockSize = 256;  // bsz <= 256
-    const int max_decoder_block_num = length / block_size - encoder_decoder_block_num; // 最大输出长度对应的block - 服务为解码分配的block数量
-    auto step_lens_inkernel = paddle::full({1}, 0, paddle::DataType::INT32, stop_flags.place());
-    auto step_bs_list = GetEmptyTensor({bsz}, paddle::DataType::INT32, stop_flags.place());
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  const int length = input_ids.shape()[1];
+  const int pre_id_length = pre_ids.shape()[1];
+  constexpr int BlockSize = 256;  // bsz <= 256
+  const int max_decoder_block_num =
+      length / block_size -
+      encoder_decoder_block_num;  // 最大输出长度对应的block -
+                                  // 服务为解码分配的block数量
+  auto step_lens_inkernel =
+      paddle::full({1}, 0, paddle::DataType::INT32, stop_flags.place());
+  auto step_bs_list =
+      GetEmptyTensor({bsz}, paddle::DataType::INT32, stop_flags.place());
 #ifdef DEBUG_STEP
-    printf(
-        "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
-        "%d\n",
-        bsz,
-        block_num_per_seq,
-        length,
-        max_decoder_block_num);
+  printf(
+      "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
+      "%d\n",
+      bsz,
+      block_num_per_seq,
+      length,
+      max_decoder_block_num);
 #endif
-    free_and_reschedule<<<1, BlockSize, 0, cu_stream>>>(
-        const_cast<bool *>(stop_flags.data<bool>()),
-        const_cast<int *>(seq_lens_this_time.data<int>()),
-        const_cast<int *>(seq_lens_decoder.data<int>()),
-        const_cast<int *>(block_tables.data<int>()),
-        const_cast<int *>(encoder_block_lens.data<int>()),
-        const_cast<bool *>(is_block_step.data<bool>()),
-        const_cast<int *>(step_bs_list.data<int>()),
-        const_cast<int *>(step_lens_inkernel.data<int>()),
-        const_cast<int *>(recover_block_list.data<int>()),
-        const_cast<int *>(recover_lens.data<int>()),
-        const_cast<int *>(need_block_list.data<int>()),
-        const_cast<int *>(need_block_len.data<int>()),
-        const_cast<int *>(used_list_len.data<int>()),
-        const_cast<int *>(free_list.data<int>()),
-        const_cast<int *>(free_list_len.data<int>()),
-        const_cast<int64_t *>(first_token_ids.data<int64_t>()),
-        bsz,
-        block_size,
-        block_num_per_seq,
-        max_decoder_block_num);
+  free_and_reschedule<<<1, BlockSize, 0, cu_stream>>>(
+      const_cast<bool *>(stop_flags.data<bool>()),
+      const_cast<int *>(seq_lens_this_time.data<int>()),
+      const_cast<int *>(seq_lens_decoder.data<int>()),
+      const_cast<int *>(block_tables.data<int>()),
+      const_cast<int *>(encoder_block_lens.data<int>()),
+      const_cast<bool *>(is_block_step.data<bool>()),
+      const_cast<int *>(step_bs_list.data<int>()),
+      const_cast<int *>(step_lens_inkernel.data<int>()),
+      const_cast<int *>(recover_block_list.data<int>()),
+      const_cast<int *>(recover_lens.data<int>()),
+      const_cast<int *>(need_block_list.data<int>()),
+      const_cast<int *>(need_block_len.data<int>()),
+      const_cast<int *>(used_list_len.data<int>()),
+      const_cast<int *>(free_list.data<int>()),
+      const_cast<int *>(free_list_len.data<int>()),
+      const_cast<int64_t *>(first_token_ids.data<int64_t>()),
+      bsz,
+      block_size,
+      block_num_per_seq,
+      max_decoder_block_num);
 #ifdef DEBUG_STEP
-    cudaDeviceSynchronize();
+  cudaDeviceSynchronize();
 #endif
-    // save output
-    auto step_lens_cpu = step_lens_inkernel.copy_to(paddle::CPUPlace(), false);
-    if (step_lens_cpu.data<int>()[0] > 0) {
-        auto step_bs_list_cpu = step_bs_list.copy_to(paddle::CPUPlace(), false);
-        auto next_tokens = paddle::full({bsz}, -1, paddle::DataType::INT64, paddle::CPUPlace());
-        for (int i = 0; i < step_lens_cpu.data<int>()[0]; i++) {
-            const int step_bid = step_bs_list_cpu.data<int>()[i];
-            next_tokens.data<int64_t>()[step_bid] = RECOVERY_STOP_SIGNAL; // need reschedule
-        }
-        const int rank_id = static_cast<int>(stop_flags.place().GetDeviceId());
-        printf("reschedule rank_id: %d, step_lens: %d", rank_id, step_lens_cpu.data<int>()[0]);
-        const int64_t* x_data = next_tokens.data<int64_t>();
-        static struct msgdata msg_sed;
-        int msg_queue_id = rank_id;
-        if (const char* inference_msg_queue_id_env_p =
-                std::getenv("INFERENCE_MSG_QUEUE_ID")) {
-            std::string inference_msg_queue_id_env_str(
-                inference_msg_queue_id_env_p);
-            int inference_msg_queue_id_from_env =
-                std::stoi(inference_msg_queue_id_env_str);
-            msg_queue_id = inference_msg_queue_id_from_env;
-        } else {
-            std::cout << "Failed to got INFERENCE_MSG_QUEUE_ID at env, use default."
-                    << std::endl;
-        }
-        int inference_msg_id_from_env = 1;
-        if (const char* inference_msg_id_env_p = std::getenv("INFERENCE_MSG_ID")) {
-            std::string inference_msg_id_env_str(inference_msg_id_env_p);
-            inference_msg_id_from_env = std::stoi(inference_msg_id_env_str);
-            if (inference_msg_id_from_env == 2) {
-                // 2 and -2 is preserve for no-output indication.
-                throw std::runtime_error(
-                    " INFERENCE_MSG_ID cannot be 2, please use other number.");
-            }
-            if (inference_msg_id_from_env < 0) {
-                throw std::runtime_error(
-                    " INFERENCE_MSG_ID cannot be negative, please use other "
-                    "number.");
-            }
-        }
-        static key_t key = ftok("/dev/shm", msg_queue_id);
-
-        static int msgid = msgget(key, IPC_CREAT | 0666);
-        msg_sed.mtype = 1;
-        msg_sed.mtext[0] = inference_msg_id_from_env;
-        msg_sed.mtext[1] = bsz;
-        for (int i = 2; i < bsz + 2; i++) {
-            msg_sed.mtext[i] = (int)x_data[i - 2];
-        }
-        if ((msgsnd(msgid, &msg_sed, (MAX_BSZ + 2) * 4, 0)) == -1) {
-            printf("full msg buffer\n");
-        }
+  // save output
+  auto step_lens_cpu = step_lens_inkernel.copy_to(paddle::CPUPlace(), false);
+  if (step_lens_cpu.data<int>()[0] > 0) {
+    auto step_bs_list_cpu = step_bs_list.copy_to(paddle::CPUPlace(), false);
+    auto next_tokens =
+        paddle::full({bsz}, -1, paddle::DataType::INT64, paddle::CPUPlace());
+    for (int i = 0; i < step_lens_cpu.data<int>()[0]; i++) {
+      const int step_bid = step_bs_list_cpu.data<int>()[i];
+      next_tokens.data<int64_t>()[step_bid] =
+          RECOVERY_STOP_SIGNAL;  // need reschedule
     }
+    const int rank_id = static_cast<int>(stop_flags.place().GetDeviceId());
+    printf("reschedule rank_id: %d, step_lens: %d",
+           rank_id,
+           step_lens_cpu.data<int>()[0]);
+    const int64_t *x_data = next_tokens.data<int64_t>();
+    static struct msgdata msg_sed;
+    int msg_queue_id = rank_id;
+    if (const char *inference_msg_queue_id_env_p =
+            std::getenv("INFERENCE_MSG_QUEUE_ID")) {
+      std::string inference_msg_queue_id_env_str(inference_msg_queue_id_env_p);
+      int inference_msg_queue_id_from_env =
+          std::stoi(inference_msg_queue_id_env_str);
+      msg_queue_id = inference_msg_queue_id_from_env;
+    } else {
+      std::cout << "Failed to got INFERENCE_MSG_QUEUE_ID at env, use default."
+                << std::endl;
+    }
+    int inference_msg_id_from_env = 1;
+    if (const char *inference_msg_id_env_p = std::getenv("INFERENCE_MSG_ID")) {
+      std::string inference_msg_id_env_str(inference_msg_id_env_p);
+      inference_msg_id_from_env = std::stoi(inference_msg_id_env_str);
+      if (inference_msg_id_from_env == 2) {
+        // 2 and -2 is preserve for no-output indication.
+        throw std::runtime_error(
+            " INFERENCE_MSG_ID cannot be 2, please use other number.");
+      }
+      if (inference_msg_id_from_env < 0) {
+        throw std::runtime_error(
+            " INFERENCE_MSG_ID cannot be negative, please use other "
+            "number.");
+      }
+    }
+    static key_t key = ftok("/dev/shm", msg_queue_id);
+
+    static int msgid = msgget(key, IPC_CREAT | 0666);
+    msg_sed.mtype = 1;
+    msg_sed.mtext[0] = inference_msg_id_from_env;
+    msg_sed.mtext[1] = bsz;
+    for (int i = 2; i < bsz + 2; i++) {
+      msg_sed.mtext[i] = (int)x_data[i - 2];
+    }
+    if ((msgsnd(msgid, &msg_sed, (MAX_BSZ + 2) * 4, 0)) == -1) {
+      printf("full msg buffer\n");
+    }
+  }
 }
 
 PD_BUILD_STATIC_OP(step_reschedule)
diff --git a/custom_ops/gpu_ops/step_system_cache.cu b/custom_ops/gpu_ops/step_system_cache.cu
index 4b236bd80a..6ce9741ee2 100644
--- a/custom_ops/gpu_ops/step_system_cache.cu
+++ b/custom_ops/gpu_ops/step_system_cache.cu
@@ -15,166 +15,181 @@
 #include "helper.h"
 
 // 根据上一步计算出的可以复原的query_id进行状态恢复
-__global__ void recover_block_system_cache(int *recover_block_list, // [bsz]
-                                           int *recover_len,
-                                           bool *stop_flags,
-                                           int *seq_lens_this_time,
-                                           int *ori_seq_lens_encoder,
-                                           int *ori_seq_lens_decoder,
-                                           int *seq_lens_encoder,
-                                           int *seq_lens_decoder,
-                                           int *block_tables,
-                                           int *free_list,
-                                           int *free_list_len,
-                                           int64_t *input_ids,
-                                           int64_t *pre_ids,
-                                           int64_t *step_idx,
-                                           int *encoder_block_lens,
-                                           int *used_list_len,
-                                           const int64_t *next_tokens,
-                                           const int64_t *first_token_ids,
+__global__ void recover_block_system_cache(int* recover_block_list,  // [bsz]
+                                           int* recover_len,
+                                           bool* stop_flags,
+                                           int* seq_lens_this_time,
+                                           int* ori_seq_lens_encoder,
+                                           int* ori_seq_lens_decoder,
+                                           int* seq_lens_encoder,
+                                           int* seq_lens_decoder,
+                                           int* block_tables,
+                                           int* free_list,
+                                           int* free_list_len,
+                                           int64_t* input_ids,
+                                           int64_t* pre_ids,
+                                           int64_t* step_idx,
+                                           int* encoder_block_lens,
+                                           int* used_list_len,
+                                           const int64_t* next_tokens,
+                                           const int64_t* first_token_ids,
                                            const int bsz,
                                            const int block_num_per_seq,
                                            const int length,
                                            const int pre_id_length) {
-    const int bid = blockIdx.x;
-    const int tid = threadIdx.x;
-    __shared__ int ori_free_list_len;
-    if (bid < recover_len[0]) {
-        const int recover_id = recover_block_list[bid];
-        const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
-        const int step_idx_now = step_idx[recover_id];
-        const int seq_len = ori_seq_len_encoder + step_idx_now;
-        const int encoder_block_len = encoder_block_lens[recover_id];
-        const int decoder_used_len = used_list_len[recover_id];
-        int *block_table_now = block_tables + recover_id * block_num_per_seq;
-        int64_t *input_ids_now = input_ids + recover_id * length;
-        int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
-        if (tid == 0) {
-            seq_lens_this_time[recover_id] = seq_len;
-            seq_lens_encoder[recover_id] = seq_len;
-            seq_lens_decoder[recover_id] = ori_seq_lens_decoder[recover_id];
-            stop_flags[recover_id] = false;
-            input_ids_now[ori_seq_len_encoder + step_idx_now - 1] = next_tokens[recover_id]; // next tokens
-            input_ids_now[0] = first_token_ids[recover_id]; // set first prompt token
-            const int ori_free_list_len_tid0 = atomicSub(free_list_len, decoder_used_len);
-            ori_free_list_len = ori_free_list_len_tid0;
+  const int bid = blockIdx.x;
+  const int tid = threadIdx.x;
+  __shared__ int ori_free_list_len;
+  if (bid < recover_len[0]) {
+    const int recover_id = recover_block_list[bid];
+    const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
+    const int step_idx_now = step_idx[recover_id];
+    const int seq_len = ori_seq_len_encoder + step_idx_now;
+    const int encoder_block_len = encoder_block_lens[recover_id];
+    const int decoder_used_len = used_list_len[recover_id];
+    int* block_table_now = block_tables + recover_id * block_num_per_seq;
+    int64_t* input_ids_now = input_ids + recover_id * length;
+    int64_t* pre_ids_now = pre_ids + recover_id * pre_id_length;
+    if (tid == 0) {
+      seq_lens_this_time[recover_id] = seq_len;
+      seq_lens_encoder[recover_id] = seq_len;
+      seq_lens_decoder[recover_id] = ori_seq_lens_decoder[recover_id];
+      stop_flags[recover_id] = false;
+      input_ids_now[ori_seq_len_encoder + step_idx_now - 1] =
+          next_tokens[recover_id];                     // next tokens
+      input_ids_now[0] = first_token_ids[recover_id];  // set first prompt token
+      const int ori_free_list_len_tid0 =
+          atomicSub(free_list_len, decoder_used_len);
+      ori_free_list_len = ori_free_list_len_tid0;
 #ifdef DEBUG_STEP
-            printf("seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, seq_len: %d, ori_free_list_len_tid0: %d, ori_free_list_len: %d\n",
-                    recover_id, ori_seq_len_encoder, step_idx_now, seq_len, ori_free_list_len_tid0, ori_free_list_len);
+      printf(
+          "seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, seq_len: %d, "
+          "ori_free_list_len_tid0: %d, ori_free_list_len: %d\n",
+          recover_id,
+          ori_seq_len_encoder,
+          step_idx_now,
+          seq_len,
+          ori_free_list_len_tid0,
+          ori_free_list_len);
 #endif
-        }
-        __syncthreads();
-        // 恢复block table
-        for (int i = tid; i < decoder_used_len; i += blockDim.x) {
-            block_table_now[encoder_block_len + i] = free_list[ori_free_list_len - i - 1];
-        }
-        // 恢复input_ids
-        for (int i = tid; i < step_idx_now - 1; i += blockDim.x) {
-            input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
-        }
     }
+    __syncthreads();
+    // 恢复block table
+    for (int i = tid; i < decoder_used_len; i += blockDim.x) {
+      block_table_now[encoder_block_len + i] =
+          free_list[ori_free_list_len - i - 1];
+    }
+    // 恢复input_ids
+    for (int i = tid; i < step_idx_now - 1; i += blockDim.x) {
+      input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
+    }
+  }
 
-    if (bid == 0 && tid == 0) {
-        recover_len[0] = 0;
-    }
+  if (bid == 0 && tid == 0) {
+    recover_len[0] = 0;
+  }
 }
 
-void StepSystemCache(const paddle::Tensor& stop_flags,
-                     const paddle::Tensor& seq_lens_this_time,
-                     const paddle::Tensor& ori_seq_lens_encoder,
-                     const paddle::Tensor& ori_seq_lens_decoder,
-                     const paddle::Tensor& seq_lens_encoder,
-                     const paddle::Tensor& seq_lens_decoder,
-                     const paddle::Tensor& block_tables, // [bsz, block_num_per_seq]
-                     const paddle::Tensor& encoder_block_lens,
-                     const paddle::Tensor& is_block_step,
-                     const paddle::Tensor& step_block_list,
-                     const paddle::Tensor& step_lens,
-                     const paddle::Tensor& recover_block_list,
-                     const paddle::Tensor& recover_lens,
-                     const paddle::Tensor& need_block_list,
-                     const paddle::Tensor& need_block_len,
-                     const paddle::Tensor& used_list_len,
-                     const paddle::Tensor& free_list,
-                     const paddle::Tensor& free_list_len,
-                     const paddle::Tensor& input_ids,
-                     const paddle::Tensor& pre_ids,
-                     const paddle::Tensor& step_idx,
-                     const paddle::Tensor& next_tokens,
-                     const paddle::Tensor& first_token_ids,
-                     const int block_size,
-                     const int encoder_decoder_block_num) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    const int length = input_ids.shape()[1];
-    const int pre_id_length = pre_ids.shape()[1];
-    constexpr int BlockSize = 256; // bsz <= 256
-    const int max_decoder_block_num = length / block_size;
-    // const int max_decoder_block_num = 2048 / block_size - encoder_decoder_block_num;
+void StepSystemCache(
+    const paddle::Tensor& stop_flags,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& ori_seq_lens_encoder,
+    const paddle::Tensor& ori_seq_lens_decoder,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& block_tables,  // [bsz, block_num_per_seq]
+    const paddle::Tensor& encoder_block_lens,
+    const paddle::Tensor& is_block_step,
+    const paddle::Tensor& step_block_list,
+    const paddle::Tensor& step_lens,
+    const paddle::Tensor& recover_block_list,
+    const paddle::Tensor& recover_lens,
+    const paddle::Tensor& need_block_list,
+    const paddle::Tensor& need_block_len,
+    const paddle::Tensor& used_list_len,
+    const paddle::Tensor& free_list,
+    const paddle::Tensor& free_list_len,
+    const paddle::Tensor& input_ids,
+    const paddle::Tensor& pre_ids,
+    const paddle::Tensor& step_idx,
+    const paddle::Tensor& next_tokens,
+    const paddle::Tensor& first_token_ids,
+    const int block_size,
+    const int encoder_decoder_block_num) {
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  const int length = input_ids.shape()[1];
+  const int pre_id_length = pre_ids.shape()[1];
+  constexpr int BlockSize = 256;  // bsz <= 256
+  const int max_decoder_block_num = length / block_size;
+  // const int max_decoder_block_num = 2048 / block_size -
+  // encoder_decoder_block_num;
 #ifdef DEBUG_STEP
-    printf("bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: %d\n", bsz, block_num_per_seq, length, max_decoder_block_num);
+  printf(
+      "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: %d\n",
+      bsz,
+      block_num_per_seq,
+      length,
+      max_decoder_block_num);
 #endif
-    free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
-        const_cast<bool*>(stop_flags.data<bool>()),
-        const_cast<int*>(seq_lens_this_time.data<int>()),
-        const_cast<int*>(seq_lens_decoder.data<int>()),
-        const_cast<int*>(block_tables.data<int>()),
-        const_cast<int*>(encoder_block_lens.data<int>()),
-        const_cast<bool*>(is_block_step.data<bool>()),
-        const_cast<int*>(step_block_list.data<int>()),
-        const_cast<int*>(step_lens.data<int>()),
+  free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
+      const_cast<bool*>(stop_flags.data<bool>()),
+      const_cast<int*>(seq_lens_this_time.data<int>()),
+      const_cast<int*>(seq_lens_decoder.data<int>()),
+      const_cast<int*>(block_tables.data<int>()),
+      const_cast<int*>(encoder_block_lens.data<int>()),
+      const_cast<bool*>(is_block_step.data<bool>()),
+      const_cast<int*>(step_block_list.data<int>()),
+      const_cast<int*>(step_lens.data<int>()),
+      const_cast<int*>(recover_block_list.data<int>()),
+      const_cast<int*>(recover_lens.data<int>()),
+      const_cast<int*>(need_block_list.data<int>()),
+      const_cast<int*>(need_block_len.data<int>()),
+      const_cast<int*>(used_list_len.data<int>()),
+      const_cast<int*>(free_list.data<int>()),
+      const_cast<int*>(free_list_len.data<int>()),
+      const_cast<int64_t*>(first_token_ids.data<int64_t>()),
+      bsz,
+      block_size,
+      block_num_per_seq,
+      max_decoder_block_num);
+#ifdef DEBUG_STEP
+  cudaDeviceSynchronize();
+#endif
+  auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
+  const int grid_size = cpu_recover_lens.data<int>()[0];
+#ifdef DEBUG_STEP
+  printf("grid_size2 %d\n", grid_size);
+#endif
+  if (grid_size > 0) {
+    recover_block_system_cache<<<grid_size, BlockSize, 0, cu_stream>>>(
         const_cast<int*>(recover_block_list.data<int>()),
         const_cast<int*>(recover_lens.data<int>()),
-        const_cast<int*>(need_block_list.data<int>()),
-        const_cast<int*>(need_block_len.data<int>()),
-        const_cast<int*>(used_list_len.data<int>()),
+        const_cast<bool*>(stop_flags.data<bool>()),
+        const_cast<int*>(seq_lens_this_time.data<int>()),
+        const_cast<int*>(ori_seq_lens_encoder.data<int>()),
+        const_cast<int*>(ori_seq_lens_decoder.data<int>()),
+        const_cast<int*>(seq_lens_encoder.data<int>()),
+        const_cast<int*>(seq_lens_decoder.data<int>()),
+        const_cast<int*>(block_tables.data<int>()),
         const_cast<int*>(free_list.data<int>()),
         const_cast<int*>(free_list_len.data<int>()),
-        const_cast<int64_t*>(first_token_ids.data<int64_t>()),
+        const_cast<int64_t*>(input_ids.data<int64_t>()),
+        const_cast<int64_t*>(pre_ids.data<int64_t>()),
+        const_cast<int64_t*>(step_idx.data<int64_t>()),
+        const_cast<int*>(encoder_block_lens.data<int>()),
+        const_cast<int*>(used_list_len.data<int>()),
+        next_tokens.data<int64_t>(),
+        first_token_ids.data<int64_t>(),
         bsz,
-        block_size,
         block_num_per_seq,
-        max_decoder_block_num
-    );
+        length,
+        pre_id_length);
 #ifdef DEBUG_STEP
     cudaDeviceSynchronize();
 #endif
-    auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
-    const int grid_size = cpu_recover_lens.data<int>()[0];
-#ifdef DEBUG_STEP
-    printf("grid_size2 %d\n", grid_size);
-#endif
-    if (grid_size > 0) {
-        recover_block_system_cache<<<grid_size, BlockSize, 0, cu_stream>>>(
-            const_cast<int*>(recover_block_list.data<int>()),
-            const_cast<int*>(recover_lens.data<int>()),
-            const_cast<bool*>(stop_flags.data<bool>()),
-            const_cast<int*>(seq_lens_this_time.data<int>()),
-            const_cast<int*>(ori_seq_lens_encoder.data<int>()),
-            const_cast<int*>(ori_seq_lens_decoder.data<int>()),
-            const_cast<int*>(seq_lens_encoder.data<int>()),
-            const_cast<int*>(seq_lens_decoder.data<int>()),
-            const_cast<int*>(block_tables.data<int>()),
-            const_cast<int*>(free_list.data<int>()),
-            const_cast<int*>(free_list_len.data<int>()),
-            const_cast<int64_t*>(input_ids.data<int64_t>()),
-            const_cast<int64_t*>(pre_ids.data<int64_t>()),
-            const_cast<int64_t*>(step_idx.data<int64_t>()),
-            const_cast<int*>(encoder_block_lens.data<int>()),
-            const_cast<int*>(used_list_len.data<int>()),
-            next_tokens.data<int64_t>(),
-            first_token_ids.data<int64_t>(),
-            bsz,
-            block_num_per_seq,
-            length,
-            pre_id_length
-        );
-#ifdef DEBUG_STEP
-        cudaDeviceSynchronize();
-#endif
-    }
+  }
 }
 
 PD_BUILD_STATIC_OP(step_system_cache)
@@ -201,8 +216,7 @@ PD_BUILD_STATIC_OP(step_system_cache)
              "step_idx",
              "next_tokens",
              "first_token_ids"})
-    .Attrs({"block_size: int",
-            "encoder_decoder_block_num: int"})
+    .Attrs({"block_size: int", "encoder_decoder_block_num: int"})
     .Outputs({"stop_flags_out",
               "seq_lens_this_time_out",
               "seq_lens_encoder_out",
diff --git a/custom_ops/gpu_ops/stop_generation.cu b/custom_ops/gpu_ops/stop_generation.cu
index d7ed5c3ae5..1dcb3c3aae 100644
--- a/custom_ops/gpu_ops/stop_generation.cu
+++ b/custom_ops/gpu_ops/stop_generation.cu
@@ -27,13 +27,13 @@
 #endif
 
 void set_flags(char *str_flags, bool *res, int length) {
-    for (int i = 0; i < length; i++) {
-        if (str_flags[i] == '0') {
-            res[i] = false;
-        } else {
-            res[i] = true;
-        }
+  for (int i = 0; i < length; i++) {
+    if (str_flags[i] == '0') {
+      res[i] = false;
+    } else {
+      res[i] = true;
     }
+  }
 }
 
 __global__ void set_value_by_flags(const bool *stop_flags,
@@ -41,14 +41,14 @@ __global__ void set_value_by_flags(const bool *stop_flags,
                                    bool *stop_flags_out,
                                    const int bs,
                                    const int64_t end_id) {
-    int tid = threadIdx.x;
-    if (tid < bs) {
-        topk_ids[tid] = stop_flags[tid] ? end_id : topk_ids[tid];
-        __syncthreads();
-        if (topk_ids[tid] == end_id) {
-            stop_flags_out[tid] = true;
-        }
+  int tid = threadIdx.x;
+  if (tid < bs) {
+    topk_ids[tid] = stop_flags[tid] ? end_id : topk_ids[tid];
+    __syncthreads();
+    if (topk_ids[tid] == end_id) {
+      stop_flags_out[tid] = true;
     }
+  }
 }
 
 __global__ void set_value_by_flags_both(const bool *flags,
@@ -57,93 +57,92 @@ __global__ void set_value_by_flags_both(const bool *flags,
                                         bool *stop_flags_out,
                                         const int bs,
                                         const int64_t end_id) {
-    int tid = threadIdx.x;
-    if (tid < bs) {
-        topk_ids[tid] = flags[tid] || stop_flags[tid] ? end_id : topk_ids[tid];
-        __syncthreads();
-        if (topk_ids[tid] == end_id) {
-            stop_flags_out[tid] = true;
-        }
+  int tid = threadIdx.x;
+  if (tid < bs) {
+    topk_ids[tid] = flags[tid] || stop_flags[tid] ? end_id : topk_ids[tid];
+    __syncthreads();
+    if (topk_ids[tid] == end_id) {
+      stop_flags_out[tid] = true;
     }
+  }
 }
 
 std::vector<paddle::Tensor> GetStopFlags(const paddle::Tensor &topk_ids,
                                          const paddle::Tensor &stop_flags,
                                          int64_t end_id,
                                          int64_t mode) {
-    // mode = 0, stop_generation and stop_flags
-    // mode = 1, just stop_generation
-    // mode = 2, just stop_flags
-    PD_CHECK(mode <= 2);
-    PD_CHECK(topk_ids.dtype() == paddle::DataType::INT64);
-    PD_CHECK(stop_flags.dtype() == paddle::DataType::BOOL);
-    auto cu_stream = topk_ids.stream();
-    std::vector<int64_t> shape = topk_ids.shape();
-    int64_t bs_now = shape[0];
-    auto topk_ids_out =
-        topk_ids.copy_to(topk_ids.place(), false);  // gpu -> gpu
-    auto stop_flags_out =
-        stop_flags.copy_to(stop_flags.place(), false);  // gpu -> gpu
-    if (mode == 0 || mode == 1) {
-        constexpr char *path =
-            "/root/paddlejob/workspace/env_run/lzy/ERNIE_ALL/"
-            "ERNIE3.0-fused-fp16/ops/test";
-        auto flags = paddle::full(
-            {bs_now, 1}, 1, paddle::DataType::BOOL, paddle::CPUPlace());
-        int fd = -1;
-        int ret = -1;
-        void *addr = nullptr;
-        fd = open(path, O_RDWR);
-        if (-1 == fd) {
-            perror("open error");
-        }
-        addr = mmap(NULL, bs_now, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
-        if (addr == MAP_FAILED) {
-            perror("mmap error");
-        }
-        close(fd);
-        set_flags((char *)addr, flags.data<bool>(), bs_now);
-        munmap(addr, bs_now);
-        auto flags_gpu = flags.copy_to(topk_ids.place(), false);  // cpu -> gpu
-        int block_size = (bs_now + 32 - 1) / 32 * 32;
-        if (mode == 0) {
-            set_value_by_flags_both<<<1, block_size, 0, cu_stream>>>(
-                flags_gpu.data<bool>(),
-                stop_flags.data<bool>(),
-                topk_ids_out.data<int64_t>(),
-                stop_flags_out.data<bool>(),
-                bs_now,
-                end_id);
-        } else {
-            set_value_by_flags<<<1, block_size, 0, cu_stream>>>(
-                flags_gpu.data<bool>(),
-                topk_ids_out.data<int64_t>(),
-                stop_flags_out.data<bool>(),
-                bs_now,
-                end_id);
-        }
-    } else if (mode == 2) {
-        int block_size = (bs_now + 32 - 1) / 32 * 32;
-        set_value_by_flags<<<1, block_size, 0, cu_stream>>>(
-            stop_flags.data<bool>(),
-            topk_ids_out.data<int64_t>(),
-            stop_flags_out.data<bool>(),
-            bs_now,
-            end_id);
+  // mode = 0, stop_generation and stop_flags
+  // mode = 1, just stop_generation
+  // mode = 2, just stop_flags
+  PD_CHECK(mode <= 2);
+  PD_CHECK(topk_ids.dtype() == paddle::DataType::INT64);
+  PD_CHECK(stop_flags.dtype() == paddle::DataType::BOOL);
+  auto cu_stream = topk_ids.stream();
+  std::vector<int64_t> shape = topk_ids.shape();
+  int64_t bs_now = shape[0];
+  auto topk_ids_out = topk_ids.copy_to(topk_ids.place(), false);  // gpu -> gpu
+  auto stop_flags_out =
+      stop_flags.copy_to(stop_flags.place(), false);  // gpu -> gpu
+  if (mode == 0 || mode == 1) {
+    constexpr char *path =
+        "/root/paddlejob/workspace/env_run/lzy/ERNIE_ALL/"
+        "ERNIE3.0-fused-fp16/ops/test";
+    auto flags = paddle::full(
+        {bs_now, 1}, 1, paddle::DataType::BOOL, paddle::CPUPlace());
+    int fd = -1;
+    int ret = -1;
+    void *addr = nullptr;
+    fd = open(path, O_RDWR);
+    if (-1 == fd) {
+      perror("open error");
     }
-    return {topk_ids_out, stop_flags_out};
+    addr = mmap(NULL, bs_now, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+    if (addr == MAP_FAILED) {
+      perror("mmap error");
+    }
+    close(fd);
+    set_flags((char *)addr, flags.data<bool>(), bs_now);
+    munmap(addr, bs_now);
+    auto flags_gpu = flags.copy_to(topk_ids.place(), false);  // cpu -> gpu
+    int block_size = (bs_now + 32 - 1) / 32 * 32;
+    if (mode == 0) {
+      set_value_by_flags_both<<<1, block_size, 0, cu_stream>>>(
+          flags_gpu.data<bool>(),
+          stop_flags.data<bool>(),
+          topk_ids_out.data<int64_t>(),
+          stop_flags_out.data<bool>(),
+          bs_now,
+          end_id);
+    } else {
+      set_value_by_flags<<<1, block_size, 0, cu_stream>>>(
+          flags_gpu.data<bool>(),
+          topk_ids_out.data<int64_t>(),
+          stop_flags_out.data<bool>(),
+          bs_now,
+          end_id);
+    }
+  } else if (mode == 2) {
+    int block_size = (bs_now + 32 - 1) / 32 * 32;
+    set_value_by_flags<<<1, block_size, 0, cu_stream>>>(
+        stop_flags.data<bool>(),
+        topk_ids_out.data<int64_t>(),
+        stop_flags_out.data<bool>(),
+        bs_now,
+        end_id);
+  }
+  return {topk_ids_out, stop_flags_out};
 }
 
 std::vector<std::vector<int64_t>> GetStopFlagsInferShape(
     const std::vector<int64_t> &topk_ids_shape,
     const std::vector<int64_t> &stop_flags_shape) {
-    return {topk_ids_shape, stop_flags_shape};
+  return {topk_ids_shape, stop_flags_shape};
 }
 
 std::vector<paddle::DataType> GetStopFlagsInferDtype(
     const paddle::DataType &topk_ids_dtype,
     const paddle::DataType &stop_flags_dtype) {
-    return {topk_ids_dtype, stop_flags_dtype};
+  return {topk_ids_dtype, stop_flags_dtype};
 }
 
 PD_BUILD_STATIC_OP(set_stop_value)
diff --git a/custom_ops/gpu_ops/swap_cache.cu b/custom_ops/gpu_ops/swap_cache.cu
index a25d08886e..9a4bed316d 100644
--- a/custom_ops/gpu_ops/swap_cache.cu
+++ b/custom_ops/gpu_ops/swap_cache.cu
@@ -16,93 +16,104 @@
 #include "paddle/extension.h"
 
 template <paddle::DataType D>
-void SwapCacheImpl(const paddle::Tensor& cache_gpu, // gpu
-                   const int64_t& cache_cpu_pointer, // cpu
+void SwapCacheImpl(const paddle::Tensor& cache_gpu,   // gpu
+                   const int64_t& cache_cpu_pointer,  // cpu
                    const int64_t& max_block_num_cpu,
                    const std::vector<int64_t>& swap_block_ids_gpu,
                    const std::vector<int64_t>& swap_block_ids_cpu,
-                //    const paddle::Tensor& swap_block_ids_dst, // cpu
-                //    const paddle::Tensor& swap_block_ids_src, // cpu
+                   //    const paddle::Tensor& swap_block_ids_dst, // cpu
+                   //    const paddle::Tensor& swap_block_ids_src, // cpu
                    int mode) {
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
-    data_t* cache_gpu_ptr = const_cast<data_t*>(cache_gpu.data<data_t>());
-    auto* cache_cpu_ptr = reinterpret_cast<data_t*>(cache_cpu_pointer);
-    auto cache_shape = cache_gpu.shape();
-    // auto* swap_block_ids_dst_ptr = swap_block_ids_dst.data<int32_t>();
-    // auto* swap_block_ids_src_ptr = swap_block_ids_src.data<int32_t>();
-    // const int swap_block_length = swap_block_ids_dst.shape()[0];
-    const int64_t max_block_num_gpu = cache_shape[0];
-    const int num_heads = cache_shape[1];
-    const int block_size = cache_shape[2];
-    const int head_dim = cache_shape[3];
-    const int64_t cache_stride = num_heads * block_size * head_dim;
-    auto stream = cache_gpu.stream();
-    for (int i = 0; i < swap_block_ids_gpu.size(); ++i) {
-        int64_t gpu_block_id = swap_block_ids_gpu[i];
-        int64_t cpu_block_id = swap_block_ids_cpu[i];
-        assert(gpu_block_id >= 0 && gpu_block_id < max_block_num_gpu);
-        assert(cpu_block_id >= 0 && cpu_block_id < max_block_num_cpu);
-        auto *cache_gpu_ptr_now = cache_gpu_ptr + gpu_block_id * cache_stride;
-        auto *cache_cpu_ptr_now = cache_cpu_ptr + cpu_block_id * cache_stride;
-        if (mode == 0) { // copy from device to host
-            cudaMemcpyAsync(cache_cpu_ptr_now, cache_gpu_ptr_now, cache_stride * sizeof(DataType_), cudaMemcpyDeviceToHost, stream);
-            // cudaMemcpy(cache_dst_ptr_now, cache_src_ptr_now, cache_stride * sizeof(DataType_), cudaMemcpyDeviceToHost);
-        } else { // copy from host to device
-            cudaMemcpyAsync(cache_gpu_ptr_now, cache_cpu_ptr_now, cache_stride * sizeof(DataType_), cudaMemcpyHostToDevice, stream);
-            // cudaMemcpy(cache_dst_ptr_now, cache_src_ptr_now, cache_stride * sizeof(DataType_), cudaMemcpyHostToDevice);
-        }
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+  data_t* cache_gpu_ptr = const_cast<data_t*>(cache_gpu.data<data_t>());
+  auto* cache_cpu_ptr = reinterpret_cast<data_t*>(cache_cpu_pointer);
+  auto cache_shape = cache_gpu.shape();
+  // auto* swap_block_ids_dst_ptr = swap_block_ids_dst.data<int32_t>();
+  // auto* swap_block_ids_src_ptr = swap_block_ids_src.data<int32_t>();
+  // const int swap_block_length = swap_block_ids_dst.shape()[0];
+  const int64_t max_block_num_gpu = cache_shape[0];
+  const int num_heads = cache_shape[1];
+  const int block_size = cache_shape[2];
+  const int head_dim = cache_shape[3];
+  const int64_t cache_stride = num_heads * block_size * head_dim;
+  auto stream = cache_gpu.stream();
+  for (int i = 0; i < swap_block_ids_gpu.size(); ++i) {
+    int64_t gpu_block_id = swap_block_ids_gpu[i];
+    int64_t cpu_block_id = swap_block_ids_cpu[i];
+    assert(gpu_block_id >= 0 && gpu_block_id < max_block_num_gpu);
+    assert(cpu_block_id >= 0 && cpu_block_id < max_block_num_cpu);
+    auto* cache_gpu_ptr_now = cache_gpu_ptr + gpu_block_id * cache_stride;
+    auto* cache_cpu_ptr_now = cache_cpu_ptr + cpu_block_id * cache_stride;
+    if (mode == 0) {  // copy from device to host
+      cudaMemcpyAsync(cache_cpu_ptr_now,
+                      cache_gpu_ptr_now,
+                      cache_stride * sizeof(DataType_),
+                      cudaMemcpyDeviceToHost,
+                      stream);
+      // cudaMemcpy(cache_dst_ptr_now, cache_src_ptr_now, cache_stride *
+      // sizeof(DataType_), cudaMemcpyDeviceToHost);
+    } else {  // copy from host to device
+      cudaMemcpyAsync(cache_gpu_ptr_now,
+                      cache_cpu_ptr_now,
+                      cache_stride * sizeof(DataType_),
+                      cudaMemcpyHostToDevice,
+                      stream);
+      // cudaMemcpy(cache_dst_ptr_now, cache_src_ptr_now, cache_stride *
+      // sizeof(DataType_), cudaMemcpyHostToDevice);
     }
-    cudaStreamSynchronize(stream);
+  }
+  cudaStreamSynchronize(stream);
 }
 
-void SwapCache(const paddle::Tensor& cache_gpu, // gpu
-               int64_t cache_cpu_ptr, // cpu memory pointer
-               int64_t max_block_num_cpu, // cpu max block num
+void SwapCache(const paddle::Tensor& cache_gpu,  // gpu
+               int64_t cache_cpu_ptr,            // cpu memory pointer
+               int64_t max_block_num_cpu,        // cpu max block num
                const std::vector<int64_t>& swap_block_ids_gpu,
                const std::vector<int64_t>& swap_block_ids_cpu,
                int rank,
                int mode) {
-    cudaSetDevice(rank); // used for distributed launch
-    switch (cache_gpu.dtype()) {
-        case paddle::DataType::BFLOAT16:
-            return SwapCacheImpl<paddle::DataType::BFLOAT16>(
-                        cache_gpu,
-                        cache_cpu_ptr,
-                        max_block_num_cpu,
-                        swap_block_ids_gpu,
-                        swap_block_ids_cpu,
-                        mode);
-        case paddle::DataType::FLOAT16:
-            return SwapCacheImpl<paddle::DataType::FLOAT16>(
-                        cache_gpu,
-                        cache_cpu_ptr,
-                        max_block_num_cpu,
-                        swap_block_ids_gpu,
-                        swap_block_ids_cpu,
-                        mode);
-        case paddle::DataType::UINT8:
-            return SwapCacheImpl<paddle::DataType::UINT8>(
-                        cache_gpu,
-                        cache_cpu_ptr,
-                        max_block_num_cpu,
-                        swap_block_ids_gpu,
-                        swap_block_ids_cpu,
-                        mode);
-        default:
-            PD_THROW("Unsupported data type.");
-    }
+  cudaSetDevice(rank);  // used for distributed launch
+  switch (cache_gpu.dtype()) {
+    case paddle::DataType::BFLOAT16:
+      return SwapCacheImpl<paddle::DataType::BFLOAT16>(cache_gpu,
+                                                       cache_cpu_ptr,
+                                                       max_block_num_cpu,
+                                                       swap_block_ids_gpu,
+                                                       swap_block_ids_cpu,
+                                                       mode);
+    case paddle::DataType::FLOAT16:
+      return SwapCacheImpl<paddle::DataType::FLOAT16>(cache_gpu,
+                                                      cache_cpu_ptr,
+                                                      max_block_num_cpu,
+                                                      swap_block_ids_gpu,
+                                                      swap_block_ids_cpu,
+                                                      mode);
+    case paddle::DataType::UINT8:
+      return SwapCacheImpl<paddle::DataType::UINT8>(cache_gpu,
+                                                    cache_cpu_ptr,
+                                                    max_block_num_cpu,
+                                                    swap_block_ids_gpu,
+                                                    swap_block_ids_cpu,
+                                                    mode);
+    default:
+      PD_THROW("Unsupported data type.");
+  }
 }
 
 PD_BUILD_STATIC_OP(swap_cache)
-    .Inputs({"cache_gpu",})
-    .Attrs({"cache_cpu_ptr: int64_t",
-            "max_block_num_cpu: int64_t",
-            "swap_block_ids_gpu: std::vector<int64_t>",
-            "swap_block_ids_cpu: std::vector<int64_t>",
-            "rank: int",
-            "mode: int",})
+    .Inputs({
+        "cache_gpu",
+    })
+    .Attrs({
+        "cache_cpu_ptr: int64_t",
+        "max_block_num_cpu: int64_t",
+        "swap_block_ids_gpu: std::vector<int64_t>",
+        "swap_block_ids_cpu: std::vector<int64_t>",
+        "rank: int",
+        "mode: int",
+    })
     .Outputs({"cache_dst_out"})
     .SetInplaceMap({{"cache_gpu", "cache_dst_out"}})
     .SetKernelFn(PD_KERNEL(SwapCache));
diff --git a/custom_ops/gpu_ops/system2group.cu b/custom_ops/gpu_ops/system2group.cu
index 09c4ca63d2..1615271b91 100644
--- a/custom_ops/gpu_ops/system2group.cu
+++ b/custom_ops/gpu_ops/system2group.cu
@@ -27,73 +27,71 @@ __global__ void System2GroupKernel(int* group_ids,
                                    const int* system_ids,
                                    const int bsz,
                                    const int max_bsz) {
-    const int ti = threadIdx.x;
-    if (ti < bsz) {
-        if (seq_lens_this_time[ti] <=
-            0) {  // 终止位置不参与分组，encoder需要是一个特定的system
-                  // id，在seqs2seqs里处理
-            return;
-        }
-        int group_id = system_ids[ti];
-        int group_len_now = atomicAdd(&group_lens[group_id], 1);
-        if (seq_lens_encoder[ti] <= 0) {  // is decoder
-            atomicAdd(dec_group_num, 1);
-            atomicAdd(&group_lens_without_encoder[group_id], 1);
-        }
-        group_ids[group_id * max_bsz + group_len_now] = ti;
+  const int ti = threadIdx.x;
+  if (ti < bsz) {
+    if (seq_lens_this_time[ti] <=
+        0) {  // 终止位置不参与分组，encoder需要是一个特定的system
+              // id，在seqs2seqs里处理
+      return;
     }
+    int group_id = system_ids[ti];
+    int group_len_now = atomicAdd(&group_lens[group_id], 1);
+    if (seq_lens_encoder[ti] <= 0) {  // is decoder
+      atomicAdd(dec_group_num, 1);
+      atomicAdd(&group_lens_without_encoder[group_id], 1);
+    }
+    group_ids[group_id * max_bsz + group_len_now] = ti;
+  }
 }
 
 std::vector<paddle::Tensor> System2Group(
     const paddle::Tensor& system_ids,
     const paddle::Tensor& seq_lens_this_time,
     const paddle::Tensor& seq_lens_encoder) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int max_bsz = seq_lens_encoder.shape()[0];
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int max_bsz = seq_lens_encoder.shape()[0];
 
-    auto group_ids = paddle::full({bsz, max_bsz},
-                                  -1,
-                                  paddle::DataType::INT32,
-                                  seq_lens_this_time.place());
-    auto group_lens = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto group_lens_without_encoder = paddle::full(
-        {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
-    auto dec_group_num = paddle::full(
-        {1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto group_ids = paddle::full(
+      {bsz, max_bsz}, -1, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto group_lens = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto group_lens_without_encoder = paddle::full(
+      {bsz, 1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
+  auto dec_group_num =
+      paddle::full({1}, 0, paddle::DataType::INT32, seq_lens_this_time.place());
 
-    const int blockSize = (bsz + 32 - 1) / 32 * 32;
-    System2GroupKernel<<<1, blockSize, 0, cu_stream>>>(
-        group_ids.data<int>(),
-        group_lens.data<int>(),
-        group_lens_without_encoder.data<int>(),
-        dec_group_num.data<int>(),
-        seq_lens_this_time.data<int>(),
-        seq_lens_encoder.data<int>(),
-        system_ids.data<int>(),
-        bsz,
-        max_bsz);
-    return {group_ids, group_lens, group_lens_without_encoder, dec_group_num};
+  const int blockSize = (bsz + 32 - 1) / 32 * 32;
+  System2GroupKernel<<<1, blockSize, 0, cu_stream>>>(
+      group_ids.data<int>(),
+      group_lens.data<int>(),
+      group_lens_without_encoder.data<int>(),
+      dec_group_num.data<int>(),
+      seq_lens_this_time.data<int>(),
+      seq_lens_encoder.data<int>(),
+      system_ids.data<int>(),
+      bsz,
+      max_bsz);
+  return {group_ids, group_lens, group_lens_without_encoder, dec_group_num};
 }
 
 std::vector<std::vector<int64_t>> System2GroupInferShape(
     const std::vector<int64_t>& system_ids_shape,
     const std::vector<int64_t>& seq_lens_this_time_shape,
     const std::vector<int64_t>& seq_lens_encoder_shape) {
-    int64_t bsz = seq_lens_this_time_shape[0];
-    int64_t max_bsz = seq_lens_encoder_shape[0];
-    return {{bsz, max_bsz}, {bsz, 1}, {bsz, 1}, {1}};
+  int64_t bsz = seq_lens_this_time_shape[0];
+  int64_t max_bsz = seq_lens_encoder_shape[0];
+  return {{bsz, max_bsz}, {bsz, 1}, {bsz, 1}, {1}};
 }
 
 std::vector<paddle::DataType> System2GroupInferDtype(
     const paddle::DataType& system_ids_dtype,
     const paddle::DataType& seq_lens_this_time_dtype,
     const paddle::DataType& seq_lens_encoder_dtype) {
-    return {seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype,
-            seq_lens_this_time_dtype};
+  return {seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype,
+          seq_lens_this_time_dtype};
 }
 
 PD_BUILD_STATIC_OP(system2group)
diff --git a/custom_ops/gpu_ops/text_image_index_out.cu b/custom_ops/gpu_ops/text_image_index_out.cu
index 7a5b44472e..c2a2911695 100644
--- a/custom_ops/gpu_ops/text_image_index_out.cu
+++ b/custom_ops/gpu_ops/text_image_index_out.cu
@@ -15,54 +15,45 @@
 #include "helper.h"
 
 template <int VecSize>
-__global__ void text_image_index_out_kernel(
-    const int32_t* token_type_ids,
-    int32_t* text_index,
-    int32_t* image_index,
-    const int64_t token_num
-) {
-    int global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
-    if (global_thread_idx >= 1) return;
-    int text_count = 0;
-    int images_count = 0;
+__global__ void text_image_index_out_kernel(const int32_t* token_type_ids,
+                                            int32_t* text_index,
+                                            int32_t* image_index,
+                                            const int64_t token_num) {
+  int global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
+  if (global_thread_idx >= 1) return;
+  int text_count = 0;
+  int images_count = 0;
 
-    for (int i = 0; i < token_num; ++i) {
-        // printf(" %d %d  %d %d \n", text_index[i], text_count, images_count, i);
-        if (token_type_ids[i] == 0) {
-            text_index[i] = text_count;
-            text_count += 1;
-        } else if (token_type_ids[i] == 1) {
-            image_index[i] = images_count;
-            images_count += 1;
-        } else {
-            // skip cuda graph padding value
-            continue;
-        }
+  for (int i = 0; i < token_num; ++i) {
+    // printf(" %d %d  %d %d \n", text_index[i], text_count, images_count, i);
+    if (token_type_ids[i] == 0) {
+      text_index[i] = text_count;
+      text_count += 1;
+    } else if (token_type_ids[i] == 1) {
+      image_index[i] = images_count;
+      images_count += 1;
+    } else {
+      // skip cuda graph padding value
+      continue;
     }
+  }
 }
 
-void TextImageIndexOut(
-            const paddle::Tensor& token_type_ids,
-             paddle::Tensor& text_index,
-             paddle::Tensor& image_index) {
-
-    const int64_t token_num = token_type_ids.shape()[0];
-    auto stream = token_type_ids.stream();
-    text_image_index_out_kernel<1><<<1, 1, 0, stream>>>(
-        token_type_ids.data<int32_t>(),
-        text_index.data<int32_t>(),
-        image_index.data<int32_t>(),
-        token_num
-    );
+void TextImageIndexOut(const paddle::Tensor& token_type_ids,
+                       paddle::Tensor& text_index,
+                       paddle::Tensor& image_index) {
+  const int64_t token_num = token_type_ids.shape()[0];
+  auto stream = token_type_ids.stream();
+  text_image_index_out_kernel<1>
+      <<<1, 1, 0, stream>>>(token_type_ids.data<int32_t>(),
+                            text_index.data<int32_t>(),
+                            image_index.data<int32_t>(),
+                            token_num);
 }
 
-
 PD_BUILD_STATIC_OP(text_image_index_out)
-    .Inputs({"token_type_ids",
-             "text_index",
-             "image_index"})
-    .Outputs({"text_index_out",
-              "image_index_out"})
+    .Inputs({"token_type_ids", "text_index", "image_index"})
+    .Outputs({"text_index_out", "image_index_out"})
     .SetInplaceMap({{"text_index", "text_index_out"},
                     {"image_index", "image_index_out"}})
     .SetKernelFn(PD_KERNEL(TextImageIndexOut));
diff --git a/custom_ops/gpu_ops/token_penalty_only_once.cu b/custom_ops/gpu_ops/token_penalty_only_once.cu
index d7c06574b0..1b579f2fff 100644
--- a/custom_ops/gpu_ops/token_penalty_only_once.cu
+++ b/custom_ops/gpu_ops/token_penalty_only_once.cu
@@ -33,27 +33,27 @@ class PDTraits;
 
 template <>
 class PDTraits<paddle::DataType::FLOAT32> {
-    public:
-    typedef float DataType;
-    typedef float data_t;
+ public:
+  typedef float DataType;
+  typedef float data_t;
 };
 
 template <>
 class PDTraits<paddle::DataType::FLOAT16> {
-    public:
-    typedef half DataType;
-    typedef paddle::float16 data_t;
+ public:
+  typedef half DataType;
+  typedef paddle::float16 data_t;
 };
 
 template <>
 class PDTraits<paddle::DataType::BFLOAT16> {
-    public:
+ public:
 #ifdef PADDLE_WITH_HIP
-    typedef hip_bfloat16 DataType;
+  typedef hip_bfloat16 DataType;
 #else
-    typedef __nv_bfloat16 DataType;
+  typedef __nv_bfloat16 DataType;
 #endif
-    typedef paddle::bfloat16 data_t;
+  typedef paddle::bfloat16 data_t;
 };
 
 __global__ void update_id_flags(const int64_t *pre_ids,
@@ -61,15 +61,15 @@ __global__ void update_id_flags(const int64_t *pre_ids,
                                 const int64_t bs,
                                 const int64_t length,
                                 const int64_t length_id) {
-    int bi = blockIdx.x;
-    int tid = threadIdx.x;
-    const int64_t *pre_ids_now = pre_ids + bi * length_id;
-    bool *id_flags_now = id_flags + bi * length;
-    for (int i = tid; i < length_id; i += blockDim.x) {
-        int64_t id = pre_ids_now[i];
-        if (id < 0) break;
-        id_flags_now[id] = true;
-    }
+  int bi = blockIdx.x;
+  int tid = threadIdx.x;
+  const int64_t *pre_ids_now = pre_ids + bi * length_id;
+  bool *id_flags_now = id_flags + bi * length;
+  for (int i = tid; i < length_id; i += blockDim.x) {
+    int64_t id = pre_ids_now[i];
+    if (id < 0) break;
+    id_flags_now[id] = true;
+  }
 }
 
 template <typename T>
@@ -78,20 +78,20 @@ __global__ void update_value_by_id_flags(const bool *id_flags,
                                          T *logits,
                                          const int64_t bs,
                                          const int64_t length) {
-    int bi = blockIdx.x;
-    int tid = threadIdx.x;
-    T *logits_now = logits + bi * length;
-    const bool *id_flags_now = id_flags + bi * length;
-    float alpha = static_cast<float>(penalty_scores[bi]);
-    for (int i = tid; i < length; i += blockDim.x) {
-        bool flag = id_flags_now[i];
-        if (!flag) continue;
-        float logit_now = static_cast<float>(logits_now[i]);
-        logits_now[i] = static_cast<T>(logit_now < 0 ? logit_now * alpha
-                                                     : logit_now / alpha);
-        // printf("bi: %d, i: %d, length: %d, logit: %f, alpha: %f, res: %f\n",
-        // bi, i, length, logit_now, alpha, (float)logits_now[id]);
-    }
+  int bi = blockIdx.x;
+  int tid = threadIdx.x;
+  T *logits_now = logits + bi * length;
+  const bool *id_flags_now = id_flags + bi * length;
+  float alpha = static_cast<float>(penalty_scores[bi]);
+  for (int i = tid; i < length; i += blockDim.x) {
+    bool flag = id_flags_now[i];
+    if (!flag) continue;
+    float logit_now = static_cast<float>(logits_now[i]);
+    logits_now[i] =
+        static_cast<T>(logit_now < 0 ? logit_now * alpha : logit_now / alpha);
+    // printf("bi: %d, i: %d, length: %d, logit: %f, alpha: %f, res: %f\n",
+    // bi, i, length, logit_now, alpha, (float)logits_now[id]);
+  }
 }
 
 template <paddle::DataType D>
@@ -99,78 +99,78 @@ std::vector<paddle::Tensor> token_penalty_once_kernel(
     const paddle::Tensor &pre_ids,
     const paddle::Tensor &logits,
     const paddle::Tensor &penalty_scores) {
-    // print_shape(pre_ids, "pre_ids");
-    // print_shape(logits, "logits");
-    // print_shape(penalty_scores, "penalty_scores");
-    typedef PDTraits<D> traits_;
-    typedef typename traits_::DataType DataType_;
-    typedef typename traits_::data_t data_t;
-    auto cu_stream = logits.stream();
-    std::vector<int64_t> shape = logits.shape();
-    auto id_flags =
-        paddle::full(shape, false, paddle::DataType::BOOL, pre_ids.place());
-    int64_t bs = shape[0];
-    int64_t length = shape[1];
-    int64_t length_id = pre_ids.shape()[1];
-    auto logits_out = logits.copy_to(logits.place(), false);  // gpu -> gpu
-    int block_size_1 = (length_id + 32 - 1) / 32 * 32;
-    block_size_1 = min(block_size_1, 512);
-    update_id_flags<<<bs, block_size_1, 0, cu_stream>>>(
-        pre_ids.data<int64_t>(), id_flags.data<bool>(), bs, length, length_id);
-    int block_size_2 = (length + 32 - 1) / 32 * 32;
-    block_size_2 = min(block_size_2, 512);
-    update_value_by_id_flags<DataType_><<<bs, block_size_2, 0, cu_stream>>>(
-        id_flags.data<bool>(),
-        reinterpret_cast<DataType_ *>(
-            const_cast<data_t *>(penalty_scores.data<data_t>())),
-        reinterpret_cast<DataType_ *>(
-            const_cast<data_t *>(logits_out.data<data_t>())),
-        bs,
-        length);
-    return {logits_out};
+  // print_shape(pre_ids, "pre_ids");
+  // print_shape(logits, "logits");
+  // print_shape(penalty_scores, "penalty_scores");
+  typedef PDTraits<D> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+  auto cu_stream = logits.stream();
+  std::vector<int64_t> shape = logits.shape();
+  auto id_flags =
+      paddle::full(shape, false, paddle::DataType::BOOL, pre_ids.place());
+  int64_t bs = shape[0];
+  int64_t length = shape[1];
+  int64_t length_id = pre_ids.shape()[1];
+  auto logits_out = logits.copy_to(logits.place(), false);  // gpu -> gpu
+  int block_size_1 = (length_id + 32 - 1) / 32 * 32;
+  block_size_1 = min(block_size_1, 512);
+  update_id_flags<<<bs, block_size_1, 0, cu_stream>>>(
+      pre_ids.data<int64_t>(), id_flags.data<bool>(), bs, length, length_id);
+  int block_size_2 = (length + 32 - 1) / 32 * 32;
+  block_size_2 = min(block_size_2, 512);
+  update_value_by_id_flags<DataType_><<<bs, block_size_2, 0, cu_stream>>>(
+      id_flags.data<bool>(),
+      reinterpret_cast<DataType_ *>(
+          const_cast<data_t *>(penalty_scores.data<data_t>())),
+      reinterpret_cast<DataType_ *>(
+          const_cast<data_t *>(logits_out.data<data_t>())),
+      bs,
+      length);
+  return {logits_out};
 }
 
 std::vector<paddle::Tensor> TokenPenaltyOnce(
     const paddle::Tensor &pre_ids,
     const paddle::Tensor &logits,
     const paddle::Tensor &penalty_scores) {
-    switch (logits.type()) {
-        case paddle::DataType::BFLOAT16: {
-            // printf("bf16\n");
-            return token_penalty_once_kernel<paddle::DataType::BFLOAT16>(
-                pre_ids, logits, penalty_scores);
-        }
-        case paddle::DataType::FLOAT16: {
-            // printf("fp16\n");
-            return token_penalty_once_kernel<paddle::DataType::FLOAT16>(
-                pre_ids, logits, penalty_scores);
-        }
-        case paddle::DataType::FLOAT32: {
-            // printf("fp32\n");
-            return token_penalty_once_kernel<paddle::DataType::FLOAT32>(
-                pre_ids, logits, penalty_scores);
-        }
-        default: {
-            PD_THROW(
-                "NOT supported data type. "
-                "Only float16 and float32 are supported. ");
-            break;
-        }
+  switch (logits.type()) {
+    case paddle::DataType::BFLOAT16: {
+      // printf("bf16\n");
+      return token_penalty_once_kernel<paddle::DataType::BFLOAT16>(
+          pre_ids, logits, penalty_scores);
     }
+    case paddle::DataType::FLOAT16: {
+      // printf("fp16\n");
+      return token_penalty_once_kernel<paddle::DataType::FLOAT16>(
+          pre_ids, logits, penalty_scores);
+    }
+    case paddle::DataType::FLOAT32: {
+      // printf("fp32\n");
+      return token_penalty_once_kernel<paddle::DataType::FLOAT32>(
+          pre_ids, logits, penalty_scores);
+    }
+    default: {
+      PD_THROW(
+          "NOT supported data type. "
+          "Only float16 and float32 are supported. ");
+      break;
+    }
+  }
 }
 
 std::vector<std::vector<int64_t>> TokenPenaltyOnceInferShape(
     const std::vector<int64_t> &pre_ids_shape,
     const std::vector<int64_t> &logits_shape,
     const std::vector<int64_t> &penalty_scores_shape) {
-    return {logits_shape};
+  return {logits_shape};
 }
 
 std::vector<paddle::DataType> TokenPenaltyOnceInferDtype(
     const paddle::DataType &pre_ids_dtype,
     const paddle::DataType &logits_dtype,
     const paddle::DataType &penalty_scores_dtype) {
-    return {logits_dtype};
+  return {logits_dtype};
 }
 
 PD_BUILD_STATIC_OP(get_token_penalty_once)
diff --git a/custom_ops/gpu_ops/token_transfer.hpp b/custom_ops/gpu_ops/token_transfer.hpp
index edfa07e997..20b5740ff7 100644
--- a/custom_ops/gpu_ops/token_transfer.hpp
+++ b/custom_ops/gpu_ops/token_transfer.hpp
@@ -30,79 +30,79 @@ namespace transfer {
 using server_callback_fn = void(std::vector<int64_t>, void *);
 
 struct BatchResult {
-    BatchResult(int64_t cur_batch_size, std::vector<int64_t> &cur_tokens)
-        : batch_size(cur_batch_size), tokens(cur_tokens) {}
-    int64_t batch_size;
-    std::vector<int64_t> tokens;
+  BatchResult(int64_t cur_batch_size, std::vector<int64_t> &cur_tokens)
+      : batch_size(cur_batch_size), tokens(cur_tokens) {}
+  int64_t batch_size;
+  std::vector<int64_t> tokens;
 };
 
 class TokenTransfer {
-    public:
-    TokenTransfer(const TokenTransfer &o) = delete;
-    const TokenTransfer &operator=(const TokenTransfer &o) = delete;
-    ~TokenTransfer() {}
+ public:
+  TokenTransfer(const TokenTransfer &o) = delete;
+  const TokenTransfer &operator=(const TokenTransfer &o) = delete;
+  ~TokenTransfer() {}
 
-    static TokenTransfer &Instance() {
-        static TokenTransfer instance;
-        return instance;
+  static TokenTransfer &Instance() {
+    static TokenTransfer instance;
+    return instance;
+  }
+
+  void RegisterCallback(server_callback_fn *cb_fn, void *cb_data) {
+    stream_cb_fn_ = cb_fn;
+    stream_cb_data_ = cb_data;
+  }
+
+  void UnRegisterCallback() {
+    stream_cb_fn_ = nullptr;
+    stream_cb_data_ = nullptr;
+  }
+
+  // once copy: cpu --> cpu
+  // array length should be (1 + MAX_BATCH)
+  bool GetBatchToken(int64_t *array) {
+    if (Empty()) {
+      return false;
+    } else {
+      assert(array != nullptr);
+      std::lock_guard<std::mutex> mtx(mtx_);
+      array[0] = q_.front().batch_size;
+      if (array[0] != 0) {
+        memmove(reinterpret_cast<void *>(array + 1),
+                reinterpret_cast<void *>(q_.front().tokens.data()),
+                sizeof(int64_t) * array[0]);
+      }
+      q_.pop();
+      return true;
     }
+  }
 
-    void RegisterCallback(server_callback_fn *cb_fn, void *cb_data) {
-        stream_cb_fn_ = cb_fn;
-        stream_cb_data_ = cb_data;
+  void PushBatchToken(int64_t cur_batch_size, int64_t *cur_tokens) {
+    std::lock_guard<std::mutex> mtx(mtx_);
+    if (q_.size() > MAX_CACHE_LENGTH) {
+      std::cout << "Warning: The queue that stores the results "
+                << "has exceeded MAX_CACHE_LENGTH and will be forcefully "
+                   "cleared."
+                << std::endl;
+      std::queue<BatchResult> empty;
+      std::swap(q_, empty);
     }
+    std::vector<int64_t> tmp(cur_tokens, cur_tokens + cur_batch_size);
+    q_.emplace(cur_batch_size, tmp);
+  }
 
-    void UnRegisterCallback() {
-        stream_cb_fn_ = nullptr;
-        stream_cb_data_ = nullptr;
-    }
+  bool Empty() {
+    std::lock_guard<std::mutex> mtx(mtx_);
+    return q_.empty();
+  }
 
-    // once copy: cpu --> cpu
-    // array length should be (1 + MAX_BATCH)
-    bool GetBatchToken(int64_t *array) {
-        if (Empty()) {
-            return false;
-        } else {
-            assert(array != nullptr);
-            std::lock_guard<std::mutex> mtx(mtx_);
-            array[0] = q_.front().batch_size;
-            if (array[0] != 0) {
-                memmove(reinterpret_cast<void *>(array + 1),
-                        reinterpret_cast<void *>(q_.front().tokens.data()),
-                        sizeof(int64_t) * array[0]);
-            }
-            q_.pop();
-            return true;
-        }
-    }
+  server_callback_fn *stream_cb_fn_ = nullptr;
+  void *stream_cb_data_ = nullptr;
 
-    void PushBatchToken(int64_t cur_batch_size, int64_t *cur_tokens) {
-        std::lock_guard<std::mutex> mtx(mtx_);
-        if (q_.size() > MAX_CACHE_LENGTH) {
-            std::cout << "Warning: The queue that stores the results "
-                      << "has exceeded MAX_CACHE_LENGTH and will be forcefully "
-                         "cleared."
-                      << std::endl;
-            std::queue<BatchResult> empty;
-            std::swap(q_, empty);
-        }
-        std::vector<int64_t> tmp(cur_tokens, cur_tokens + cur_batch_size);
-        q_.emplace(cur_batch_size, tmp);
-    }
+ private:
+  TokenTransfer() {}
 
-    bool Empty() {
-        std::lock_guard<std::mutex> mtx(mtx_);
-        return q_.empty();
-    }
-
-    server_callback_fn *stream_cb_fn_ = nullptr;
-    void *stream_cb_data_ = nullptr;
-
-    private:
-    TokenTransfer() {}
-
-    std::mutex mtx_;
-    std::queue<BatchResult> q_;
+  std::mutex mtx_;
+  std::queue<BatchResult> q_;
 };
 
 }  // namespace transfer
diff --git a/custom_ops/gpu_ops/transfer_output.cc b/custom_ops/gpu_ops/transfer_output.cc
index 0221cab0fc..bec4d46252 100644
--- a/custom_ops/gpu_ops/transfer_output.cc
+++ b/custom_ops/gpu_ops/transfer_output.cc
@@ -35,66 +35,66 @@
 
 #ifdef DEBUG_TRANSFER_OUTPUT
 void PrintVec(std::vector<int64_t> &vec) {
-    std::cout << "std::vector vec_size: " << vec.size();
-    for (int i{0}; i < vec.size(); i++) {
-        std::cout << " " << vec[i];
-    }
-    std::cout << std::endl;
+  std::cout << "std::vector vec_size: " << vec.size();
+  for (int i{0}; i < vec.size(); i++) {
+    std::cout << " " << vec[i];
+  }
+  std::cout << std::endl;
 }
 
 void PrintVec(int64_t *arr) {
-    std::cout << "READ vec_size: " << arr[0];
-    for (int i{1}; i < arr[0] + 1; i++) {
-        std::cout << " " << arr[i];
-    }
-    std::cout << std::endl;
+  std::cout << "READ vec_size: " << arr[0];
+  for (int i{1}; i < arr[0] + 1; i++) {
+    std::cout << " " << arr[i];
+  }
+  std::cout << std::endl;
 }
 
 void PrintVec(int64_t bs, int64_t *arr) {
-    std::cout << "WRITE vec_size: " << bs;
-    for (int i{0}; i < bs; i++) {
-        std::cout << " " << arr[i];
-    }
-    std::cout << std::endl;
+  std::cout << "WRITE vec_size: " << bs;
+  for (int i{0}; i < bs; i++) {
+    std::cout << " " << arr[i];
+  }
+  std::cout << std::endl;
 }
 #endif
 
 std::vector<paddle::Tensor> TransferOutput(const paddle::Tensor &x,
                                            int64_t rank_id) {
-    using namespace paddle::inference::transfer;
+  using namespace paddle::inference::transfer;
 
-    auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
-    if (rank_id != 0) {
-        return {x_cpu};
-    }
-    std::vector<int64_t> x_shape = x_cpu.shape();
-    int64_t token_num = x_cpu.numel();
-    // only support int64_t
-    assert(x_cpu.type() == paddle::DataType::INT64);
+  auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
+  if (rank_id != 0) {
+    return {x_cpu};
+  }
+  std::vector<int64_t> x_shape = x_cpu.shape();
+  int64_t token_num = x_cpu.numel();
+  // only support int64_t
+  assert(x_cpu.type() == paddle::DataType::INT64);
 
-    auto &token_transfer = TokenTransfer::Instance();
-    if (token_transfer.stream_cb_fn_) {
-        auto data_ptr = x_cpu.data<int64_t>();
-        std::vector<int64_t> tokens(data_ptr, data_ptr + token_num);
-        token_transfer.stream_cb_fn_(tokens, token_transfer.stream_cb_data_);
-    }
+  auto &token_transfer = TokenTransfer::Instance();
+  if (token_transfer.stream_cb_fn_) {
+    auto data_ptr = x_cpu.data<int64_t>();
+    std::vector<int64_t> tokens(data_ptr, data_ptr + token_num);
+    token_transfer.stream_cb_fn_(tokens, token_transfer.stream_cb_data_);
+  }
 #ifdef DEBUG_TRANSFER_OUTPUT
-    else {
-        token_transfer.PushBatchToken(token_num, x_cpu.data<int64_t>());
-    }
+  else {
+    token_transfer.PushBatchToken(token_num, x_cpu.data<int64_t>());
+  }
 #endif
 
-    return {x_cpu};
+  return {x_cpu};
 }
 
 std::vector<std::vector<int64_t>> TransferOutputInferShape(
     const std::vector<int64_t> &x_shape) {
-    return {x_shape};
+  return {x_shape};
 }
 
 std::vector<paddle::DataType> TransferOutputInferDtype(
     const paddle::DataType &x_dtype) {
-    return {x_dtype};
+  return {x_dtype};
 }
 
 PD_BUILD_STATIC_OP(transfer_output)
diff --git a/custom_ops/gpu_ops/tune_cublaslt_gemm.cu b/custom_ops/gpu_ops/tune_cublaslt_gemm.cu
index 428d563641..f3b1e3aca2 100644
--- a/custom_ops/gpu_ops/tune_cublaslt_gemm.cu
+++ b/custom_ops/gpu_ops/tune_cublaslt_gemm.cu
@@ -180,7 +180,8 @@ static void TestMatmulRun(cublasLtHandle_t ltHandle,
     std::cerr << "Not enough workspace! Required "
               << static_cast<double>(heurResult.workspaceSize) / 1024.0 /
                      1024.0 / 1024.0
-              << " GiB" << ", But remaining "
+              << " GiB"
+              << ", But remaining "
               << static_cast<double>(remainingMemorySize) / 1024.0 / 1024.0 /
                      1024.0
               << " GiB" << std::endl;
@@ -483,7 +484,7 @@ class DevContext {};
 class CPUContext : public DevContext {};
 
 class CUBLASLTContext : public DevContext {
-public:
+ public:
   CUBLASLTContext() { CUDA_CHECK(cublasLtCreate(&handle)); }
 
   cublasLtHandle_t handle;
diff --git a/custom_ops/gpu_ops/unset_data_ipc.cu b/custom_ops/gpu_ops/unset_data_ipc.cu
index d265e93af8..6a22963279 100644
--- a/custom_ops/gpu_ops/unset_data_ipc.cu
+++ b/custom_ops/gpu_ops/unset_data_ipc.cu
@@ -46,10 +46,11 @@ static inline int sharedMemoryUnlinkByName(const char* name) {
 }
 
 void UnsetDataIpc(const paddle::Tensor& tmp_input,
-                         const std::string& shm_name,
-                         bool close_ipc,
-                         bool unlink_shm) {
-  // 1) 关闭消费者导入的 IPC 映射（仅当 close_ipc=true 且该指针确为 OpenMemHandle 得来）
+                  const std::string& shm_name,
+                  bool close_ipc,
+                  bool unlink_shm) {
+  // 1) 关闭消费者导入的 IPC 映射（仅当 close_ipc=true 且该指针确为
+  // OpenMemHandle 得来）
   if (close_ipc) {
     void* ptr = const_cast<void*>(tmp_input.data());
     checkCudaErrors(cudaIpcCloseMemHandle(ptr));
@@ -59,8 +60,8 @@ void UnsetDataIpc(const paddle::Tensor& tmp_input,
   if (unlink_shm) {
     int rc = sharedMemoryUnlinkByName(shm_name.c_str());
     if (rc != 0) {
-      PD_THROW("Unlink shared memory failed: name=%s, err=%d",
-               shm_name.c_str(), rc);
+      PD_THROW(
+          "Unlink shared memory failed: name=%s, err=%d", shm_name.c_str(), rc);
     }
   }
 }
diff --git a/custom_ops/gpu_ops/update_split_fuse_input.cu b/custom_ops/gpu_ops/update_split_fuse_input.cu
index 17bacb0132..a43b5903ca 100644
--- a/custom_ops/gpu_ops/update_split_fuse_input.cu
+++ b/custom_ops/gpu_ops/update_split_fuse_input.cu
@@ -29,39 +29,38 @@ update_split_fuse_inputs_kernel(int* split_fuse_seq_lens,
                                 int64_t* step_idx,
                                 const int split_fuse_size,
                                 const int max_seq_len) {
-    const int bi = blockIdx.x;
-    const int tidx = threadIdx.x;
-    if (split_fuse_seq_lens[bi] <= 0) {
-        return;
+  const int bi = blockIdx.x;
+  const int tidx = threadIdx.x;
+  if (split_fuse_seq_lens[bi] <= 0) {
+    return;
+  }
+  if (split_fuse_cur_seq_lens[bi] < split_fuse_seq_lens[bi]) {
+    const int cur_add_tokens = min(
+        split_fuse_seq_lens[bi] - split_fuse_cur_seq_lens[bi], split_fuse_size);
+    int64_t* split_fuse_all_input_ids_cur_batch = split_fuse_all_input_ids +
+                                                  bi * max_seq_len +
+                                                  split_fuse_cur_seq_lens[bi];
+    int64_t* input_ids_cur_batch = input_ids + bi * max_seq_len;
+    for (int i = tidx; i < cur_add_tokens; i += blockDim.x) {
+      input_ids_cur_batch[i] = split_fuse_all_input_ids_cur_batch[i];
     }
-    if (split_fuse_cur_seq_lens[bi] < split_fuse_seq_lens[bi]) {
-        const int cur_add_tokens =
-            min(split_fuse_seq_lens[bi] - split_fuse_cur_seq_lens[bi],
-                split_fuse_size);
-        int64_t* split_fuse_all_input_ids_cur_batch =
-            split_fuse_all_input_ids + bi * max_seq_len +
-            split_fuse_cur_seq_lens[bi];
-        int64_t* input_ids_cur_batch = input_ids + bi * max_seq_len;
-        for (int i = tidx; i < cur_add_tokens; i += blockDim.x) {
-            input_ids_cur_batch[i] = split_fuse_all_input_ids_cur_batch[i];
-        }
-        if (threadIdx.x == 0) {
-            seq_lens_this_time[bi] = cur_add_tokens;
-            seq_lens_encoder[bi] = cur_add_tokens;
-            seq_lens_decoder[bi] = split_fuse_cur_seq_lens[bi];
-            step_idx[bi] = 0;
-            split_fuse_cur_seq_lens[bi] += cur_add_tokens;
-        }
-    } else if (split_fuse_cur_seq_lens[bi] >= split_fuse_seq_lens[bi]) {
-        if (threadIdx.x == 0) {
-            seq_lens_decoder[bi] = split_fuse_cur_seq_lens[bi];
-            seq_lens_this_time[bi] = 1;
-            step_idx[bi] = 1;
-            seq_lens_encoder[bi] = 0;
-            split_fuse_cur_seq_lens[bi] = 0;
-            split_fuse_seq_lens[bi] = 0;
-        }
+    if (threadIdx.x == 0) {
+      seq_lens_this_time[bi] = cur_add_tokens;
+      seq_lens_encoder[bi] = cur_add_tokens;
+      seq_lens_decoder[bi] = split_fuse_cur_seq_lens[bi];
+      step_idx[bi] = 0;
+      split_fuse_cur_seq_lens[bi] += cur_add_tokens;
     }
+  } else if (split_fuse_cur_seq_lens[bi] >= split_fuse_seq_lens[bi]) {
+    if (threadIdx.x == 0) {
+      seq_lens_decoder[bi] = split_fuse_cur_seq_lens[bi];
+      seq_lens_this_time[bi] = 1;
+      step_idx[bi] = 1;
+      seq_lens_encoder[bi] = 0;
+      split_fuse_cur_seq_lens[bi] = 0;
+      split_fuse_seq_lens[bi] = 0;
+    }
+  }
 }
 
 void UpdateSplitFuseInputes(const paddle::Tensor& split_fuse_seq_lens,
@@ -75,23 +74,20 @@ void UpdateSplitFuseInputes(const paddle::Tensor& split_fuse_seq_lens,
                             const int max_seq_len,
                             const int max_batch_size,
                             const int split_fuse_size) {
-    dim3 grids;
-    grids.x = max_batch_size;
-    const int block_size = 128;
-    update_split_fuse_inputs_kernel<<<grids,
-                                      block_size,
-                                      0,
-                                      input_ids.stream()>>>(
-        const_cast<int*>(split_fuse_seq_lens.data<int>()),
-        const_cast<int*>(split_fuse_cur_seq_lens.data<int>()),
-        const_cast<int64_t*>(split_fuse_all_input_ids.data<int64_t>()),
-        const_cast<int64_t*>(input_ids.data<int64_t>()),
-        const_cast<int*>(seq_lens_this_time.data<int>()),
-        const_cast<int*>(seq_lens_encoder.data<int>()),
-        const_cast<int*>(seq_lens_decoder.data<int>()),
-        const_cast<int64_t*>(step_idx.data<int64_t>()),
-        split_fuse_size,
-        max_seq_len);
+  dim3 grids;
+  grids.x = max_batch_size;
+  const int block_size = 128;
+  update_split_fuse_inputs_kernel<<<grids, block_size, 0, input_ids.stream()>>>(
+      const_cast<int*>(split_fuse_seq_lens.data<int>()),
+      const_cast<int*>(split_fuse_cur_seq_lens.data<int>()),
+      const_cast<int64_t*>(split_fuse_all_input_ids.data<int64_t>()),
+      const_cast<int64_t*>(input_ids.data<int64_t>()),
+      const_cast<int*>(seq_lens_this_time.data<int>()),
+      const_cast<int*>(seq_lens_encoder.data<int>()),
+      const_cast<int*>(seq_lens_decoder.data<int>()),
+      const_cast<int64_t*>(step_idx.data<int64_t>()),
+      split_fuse_size,
+      max_seq_len);
 }
 
 PD_BUILD_STATIC_OP(update_split_fuse_inputs)
diff --git a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/kernel_traits.h b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/kernel_traits.h
index db4e86a2aa..cd92a6b6e5 100644
--- a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/kernel_traits.h
+++ b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/kernel_traits.h
@@ -23,129 +23,160 @@
 
 using namespace cute;
 
-template <int kStages, class GemmType, class OutputType, class SmemLayoutA,
+template <int kStages,
+          class GemmType,
+          class OutputType,
+          class SmemLayoutA,
           class SmemLayoutE,
-          class SmemLayoutB, class SmemLayoutC>
+          class SmemLayoutB,
+          class SmemLayoutC>
 struct SharedStorage {
-    union {
-        struct {
-            cute::array_aligned<GemmType, cute::cosize_v<SmemLayoutA>> smem_a;
-            cute::array_aligned<uint32_t, cute::cosize_v<SmemLayoutE>> smem_e;
-            cute::array_aligned<GemmType, cute::cosize_v<SmemLayoutB>> smem_b;
-        };
-        cute::array_aligned<OutputType, cute::cosize_v<SmemLayoutC>> smem_c;
+  union {
+    struct {
+      cute::array_aligned<GemmType, cute::cosize_v<SmemLayoutA>> smem_a;
+      cute::array_aligned<uint32_t, cute::cosize_v<SmemLayoutE>> smem_e;
+      cute::array_aligned<GemmType, cute::cosize_v<SmemLayoutB>> smem_b;
     };
+    cute::array_aligned<OutputType, cute::cosize_v<SmemLayoutC>> smem_c;
+  };
 
   struct {
     typename cutlass::PipelineTmaAsync<kStages>::SharedStorage pipeline;
   };
 };
 
-template<int kBlockM_, int kBlockN_, int kBlockK_,
-        int kNWarps_, int kStages_,
-        int kTiles_, int M_,
-        int TokenPackSize_,
-        int TAIL_N_ = 0,
-        int kClusterM_ = 1,
-        typename elem_type=cutlass::float_e4m3_t,
-        typename OutputType = cutlass::bfloat16_t>
+template <int kBlockM_,
+          int kBlockN_,
+          int kBlockK_,
+          int kNWarps_,
+          int kStages_,
+          int kTiles_,
+          int M_,
+          int TokenPackSize_,
+          int TAIL_N_ = 0,
+          int kClusterM_ = 1,
+          typename elem_type = cutlass::float_e4m3_t,
+          typename OutputType = cutlass::bfloat16_t>
 struct Kernel_traits {
-    using Element = elem_type;
-    using ElementAccum = float;
-    using ElementOutput = OutputType;
-    static_assert(cutlass::sizeof_bits_v<Element> == 8);
+  using Element = elem_type;
+  using ElementAccum = float;
+  using ElementOutput = OutputType;
+  static_assert(cutlass::sizeof_bits_v<Element> == 8);
 
-    static constexpr int kNWarps = kNWarps_;
-    static constexpr int kNThreads = kNWarps * cutlass::NumThreadsPerWarp;
-    static constexpr int NumProducerThreads = cutlass::NumThreadsPerWarpGroup;
-    static constexpr int NumMmaThreads = kNThreads - NumProducerThreads;
+  static constexpr int kNWarps = kNWarps_;
+  static constexpr int kNThreads = kNWarps * cutlass::NumThreadsPerWarp;
+  static constexpr int NumProducerThreads = cutlass::NumThreadsPerWarpGroup;
+  static constexpr int NumMmaThreads = kNThreads - NumProducerThreads;
 
-    static_assert(kNWarps_ == 12);
+  static_assert(kNWarps_ == 12);
 
-    static constexpr int kBlockM = kBlockM_;
-    static constexpr int kBlockN = kBlockN_;
-    static constexpr int kBlockK = kBlockK_;
-    static constexpr int kTiles = kTiles_;
-    static constexpr int TokenPackSize = TokenPackSize_;
-    static constexpr int TAIL_N = TAIL_N_;
-    static constexpr int M = M_;
+  static constexpr int kBlockM = kBlockM_;
+  static constexpr int kBlockN = kBlockN_;
+  static constexpr int kBlockK = kBlockK_;
+  static constexpr int kTiles = kTiles_;
+  static constexpr int TokenPackSize = TokenPackSize_;
+  static constexpr int TAIL_N = TAIL_N_;
+  static constexpr int M = M_;
 
-    using TileShape_MNK = Shape<Int<kBlockM>, Int<kBlockN>, Int<kBlockK>>;
-    using TileShape_MNK_TAIL = Shape<Int<kBlockM>, Int<TAIL_N>, Int<kBlockK>>;
-    static constexpr int kClusterM = kClusterM_;
-    using ClusterShape_MNK = Shape<Int<kClusterM>, _1, _1>;
+  using TileShape_MNK = Shape<Int<kBlockM>, Int<kBlockN>, Int<kBlockK>>;
+  using TileShape_MNK_TAIL = Shape<Int<kBlockM>, Int<TAIL_N>, Int<kBlockK>>;
+  static constexpr int kClusterM = kClusterM_;
+  using ClusterShape_MNK = Shape<Int<kClusterM>, _1, _1>;
 
-    static constexpr int kStages = kStages_;
-    static_assert(kStages > 1);
+  static constexpr int kStages = kStages_;
+  static_assert(kStages > 1);
 
-    using AtomLayoutMNK = Layout<Shape<Int<kBlockM / 64>, _1, _1>>;
+  using AtomLayoutMNK = Layout<Shape<Int<kBlockM / 64>, _1, _1>>;
 
-    using TiledMma = decltype(cute::make_tiled_mma(
-        cute::GMMA::ss_op_selector<Element, Element, ElementAccum, TileShape_MNK>(),
-        AtomLayoutMNK{}));
+  using TiledMma = decltype(cute::make_tiled_mma(
+      cute::GMMA::
+          ss_op_selector<Element, Element, ElementAccum, TileShape_MNK>(),
+      AtomLayoutMNK{}));
 
-    using Mma = decltype(cute::GMMA::ss_op_selector_sparse<Element, Element, ElementAccum, TileShape_MNK>());
+  using Mma = decltype(cute::GMMA::ss_op_selector_sparse<Element,
+                                                         Element,
+                                                         ElementAccum,
+                                                         TileShape_MNK>());
 
-    using Mma_TAIL = decltype(cute::GMMA::ss_op_selector_sparse<Element, Element, ElementAccum, TileShape_MNK_TAIL>());
+  using Mma_TAIL =
+      decltype(cute::GMMA::ss_op_selector_sparse<Element,
+                                                 Element,
+                                                 ElementAccum,
+                                                 TileShape_MNK_TAIL>());
 
-    using SmemLayoutAtomA = decltype(
-        cutlass::gemm::collective::detail::rs_smem_selector<
-            GMMA::Major::K, Element, Int<kBlockM / 2>, Int<kBlockK>>());
+  using SmemLayoutAtomA =
+      decltype(cutlass::gemm::collective::detail::rs_smem_selector<
+               GMMA::Major::K,
+               Element,
+               Int<kBlockM / 2>,
+               Int<kBlockK>>());
 
-    using SmemLayoutA = decltype(
-        tile_to_shape(SmemLayoutAtomA{},
-            make_shape(Int<kBlockM / 2>{}, Int<kBlockK>{}, Int<kStages>{})));
+  using SmemLayoutA = decltype(tile_to_shape(
+      SmemLayoutAtomA{},
+      make_shape(Int<kBlockM / 2>{}, Int<kBlockK>{}, Int<kStages>{})));
 
-    using SmemLayoutAtomB = decltype(
-        cutlass::gemm::collective::detail::ss_smem_selector<
-            GMMA::Major::K, Element, decltype(cute::get<1>(TileShape_MNK{})),
-            decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutAtomB =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               Element,
+               decltype(cute::get<1>(TileShape_MNK{})),
+               decltype(cute::get<2>(TileShape_MNK{}))>());
 
-    using SmemLayoutB = decltype(
-        tile_to_shape(SmemLayoutAtomB{},
-            make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{}), Int<kStages>{})));
+  using SmemLayoutB =
+      decltype(tile_to_shape(SmemLayoutAtomB{},
+                             make_shape(shape<1>(TileShape_MNK{}),
+                                        shape<2>(TileShape_MNK{}),
+                                        Int<kStages>{})));
 
-    using SmemLayoutAtomB_TAIL = decltype(
-        cutlass::gemm::collective::detail::rs_smem_selector<
-            GMMA::Major::K, Element, decltype(cute::get<1>(TileShape_MNK_TAIL{})),
-            decltype(cute::get<2>(TileShape_MNK_TAIL{}))>());
+  using SmemLayoutAtomB_TAIL =
+      decltype(cutlass::gemm::collective::detail::rs_smem_selector<
+               GMMA::Major::K,
+               Element,
+               decltype(cute::get<1>(TileShape_MNK_TAIL{})),
+               decltype(cute::get<2>(TileShape_MNK_TAIL{}))>());
 
-    using SmemLayoutB_TAIL = decltype(
-        tile_to_shape(SmemLayoutAtomB_TAIL{},
-            make_shape(
-                shape<1>(TileShape_MNK_TAIL{}),
-                shape<2>(TileShape_MNK_TAIL{}),
-                Int<kStages>{})
-            ));
-    using SmemLayoutAtomC = decltype(
-        cutlass::gemm::collective::detail::ss_smem_selector<
-        GMMA::Major::K, ElementOutput,
-        decltype(cute::get<0>(TileShape_MNK{})),
-        decltype(cute::get<1>(TileShape_MNK{}))>());
+  using SmemLayoutB_TAIL =
+      decltype(tile_to_shape(SmemLayoutAtomB_TAIL{},
+                             make_shape(shape<1>(TileShape_MNK_TAIL{}),
+                                        shape<2>(TileShape_MNK_TAIL{}),
+                                        Int<kStages>{})));
+  using SmemLayoutAtomC =
+      decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+               GMMA::Major::K,
+               ElementOutput,
+               decltype(cute::get<0>(TileShape_MNK{})),
+               decltype(cute::get<1>(TileShape_MNK{}))>());
 
-    using SmemLayoutC = decltype(tile_to_shape(SmemLayoutAtomC{}, select<0, 1>(TileShape_MNK{})));
+  using SmemLayoutC =
+      decltype(tile_to_shape(SmemLayoutAtomC{}, select<0, 1>(TileShape_MNK{})));
 
-    using SmemLayoutE = Layout<Shape<Int<NumMmaThreads>, Int<kBlockK / 64>, Int<kStages>>>;
+  using SmemLayoutE =
+      Layout<Shape<Int<NumMmaThreads>, Int<kBlockK / 64>, Int<kStages>>>;
 
-    using SharedStorage = SharedStorage<
-        kStages, Element, ElementOutput, SmemLayoutA, SmemLayoutE, SmemLayoutB, SmemLayoutC>;
+  using SharedStorage = SharedStorage<kStages,
+                                      Element,
+                                      ElementOutput,
+                                      SmemLayoutA,
+                                      SmemLayoutE,
+                                      SmemLayoutB,
+                                      SmemLayoutC>;
 
-    using MainloopPipeline = typename cutlass::PipelineTmaAsync<kStages>;
-    using PipelineState = typename cutlass::PipelineState<kStages>;
+  using MainloopPipeline = typename cutlass::PipelineTmaAsync<kStages>;
+  using PipelineState = typename cutlass::PipelineState<kStages>;
 
-    static constexpr int kNumVecElem = ceil_div(128, sizeof_bits_v<OutputType>);
-    static constexpr int kNumThreadsPerRow = kBlockN / kNumVecElem;
-    static constexpr int kNumRows = NumMmaThreads / kNumThreadsPerRow;
-    using TiledCopyCAtom = cute::Copy_Atom<cute::UniversalCopy<cutlass::uint128_t>, OutputType>;
-    using TiledCopyCThrLayout = decltype(cute::make_layout(
-        cute::make_shape(Int<kNumRows>{}, Int<kNumThreadsPerRow>{}),
-        LayoutRight{}));
-    using TiledCopyCValLayout = decltype(cute::make_layout(
-        cute::make_shape(_1{}, Int<kNumVecElem>{}),
-        LayoutRight{}));
-    using TiledCopyC = decltype(make_tiled_copy(
-        TiledCopyCAtom{},
-        TiledCopyCThrLayout{}, // Thr layout
-        TiledCopyCValLayout{} // Val layout
-    ));
+  static constexpr int kNumVecElem = ceil_div(128, sizeof_bits_v<OutputType>);
+  static constexpr int kNumThreadsPerRow = kBlockN / kNumVecElem;
+  static constexpr int kNumRows = NumMmaThreads / kNumThreadsPerRow;
+  using TiledCopyCAtom =
+      cute::Copy_Atom<cute::UniversalCopy<cutlass::uint128_t>, OutputType>;
+  using TiledCopyCThrLayout = decltype(cute::make_layout(
+      cute::make_shape(Int<kNumRows>{}, Int<kNumThreadsPerRow>{}),
+      LayoutRight{}));
+  using TiledCopyCValLayout = decltype(cute::make_layout(
+      cute::make_shape(_1{}, Int<kNumVecElem>{}), LayoutRight{}));
+  using TiledCopyC =
+      decltype(make_tiled_copy(TiledCopyCAtom{},
+                               TiledCopyCThrLayout{},  // Thr layout
+                               TiledCopyCValLayout{}   // Val layout
+                               ));
 };
diff --git a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/mainloop_fwd.h b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/mainloop_fwd.h
index 10f86d53b9..5e0d7a6347 100644
--- a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/mainloop_fwd.h
+++ b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/mainloop_fwd.h
@@ -29,438 +29,483 @@
 using namespace cute;
 template <typename Ktraits>
 struct CollectiveMainloopFwd {
+  using Element = typename Ktraits::Element;
+  using ElementOutput = typename Ktraits::ElementOutput;
+  using TileShape_MNK = typename Ktraits::TileShape_MNK;
+  using ClusterShape = typename Ktraits::ClusterShape_MNK;
+  using ElementAccum = typename Ktraits::ElementAccum;
 
-    using Element = typename Ktraits::Element;
-    using ElementOutput = typename Ktraits::ElementOutput;
-    using TileShape_MNK = typename Ktraits::TileShape_MNK;
-    using ClusterShape = typename Ktraits::ClusterShape_MNK;
-    using ElementAccum = typename Ktraits::ElementAccum;
+  static constexpr int kStages = Ktraits::kStages;
+  static constexpr int kBlockM = Ktraits::kBlockM;
+  static constexpr int kBlockN = Ktraits::kBlockN;
+  static constexpr int kBlockK = Ktraits::kBlockK;
+  static constexpr int NumCopyThreads = cutlass::NumThreadsPerWarpGroup;
+  static constexpr int kTiles = Ktraits::kTiles;
+  static constexpr int NumMmaThreads = size(typename Ktraits::TiledMma{});
+  static constexpr int TokenPackSize = Ktraits::TokenPackSize;
+  static constexpr int M = Ktraits::M;
 
-    static constexpr int kStages = Ktraits::kStages;
-    static constexpr int kBlockM = Ktraits::kBlockM;
-    static constexpr int kBlockN = Ktraits::kBlockN;
-    static constexpr int kBlockK = Ktraits::kBlockK;
-    static constexpr int NumCopyThreads = cutlass::NumThreadsPerWarpGroup;
-    static constexpr int kTiles = Ktraits::kTiles;
-    static constexpr int NumMmaThreads = size(typename Ktraits::TiledMma{});
-    static constexpr int TokenPackSize = Ktraits::TokenPackSize;
-    static constexpr int M = Ktraits::M;
+  using GmemTiledCopy = cute::SM90_TMA_LOAD;
+  using GmemTiledCopyStore = cute::SM90_TMA_STORE;
 
+  using SmemLayoutA = typename Ktraits::SmemLayoutA;
+  using SmemLayoutB = typename Ktraits::SmemLayoutB;
+  using SmemLayoutC = typename Ktraits::SmemLayoutC;
+  using SmemLayoutE = typename Ktraits::SmemLayoutE;
+  using SmemLayoutB_TAIL = typename Ktraits::SmemLayoutB_TAIL;
 
-    using GmemTiledCopy = cute::SM90_TMA_LOAD;
-    using GmemTiledCopyStore = cute::SM90_TMA_STORE;
+  using ShapeT = cute::Shape<int64_t, int64_t, int64_t>;
+  using StrideT = cute::Shape<int64_t, _1, int64_t>;
+  using LayoutT = cute::Layout<ShapeT, StrideT>;
 
-    using SmemLayoutA = typename Ktraits::SmemLayoutA;
-    using SmemLayoutB = typename Ktraits::SmemLayoutB;
-    using SmemLayoutC = typename Ktraits::SmemLayoutC;
-    using SmemLayoutE = typename Ktraits::SmemLayoutE;
-    using SmemLayoutB_TAIL = typename Ktraits::SmemLayoutB_TAIL;
+  using WShapeT = cute::Shape<int64_t, int64_t, int64_t, int64_t, int64_t>;
+  using WStrideT = cute::Shape<int64_t, _1, int64_t, int64_t, int64_t>;
+  using WLayoutT = cute::Layout<WShapeT, WStrideT>;
 
-    using ShapeT = cute::Shape<int64_t, int64_t, int64_t>;
-    using StrideT = cute::Shape<int64_t, _1, int64_t>;
-    using LayoutT = cute::Layout<ShapeT, StrideT>;
+  using EShapeT = cute::Shape<int64_t, int64_t, int64_t, int64_t, int64_t>;
+  using EStrideT = cute::Shape<_1, int64_t, int64_t, int64_t, int64_t>;
+  using ELayoutT = cute::Layout<EShapeT, EStrideT>;
 
-    using WShapeT = cute::Shape<int64_t, int64_t, int64_t, int64_t, int64_t>;
-    using WStrideT = cute::Shape<int64_t, _1, int64_t, int64_t, int64_t>;
-    using WLayoutT = cute::Layout<WShapeT, WStrideT>;
+  using TMA_A = decltype(make_tma_copy(
+      GmemTiledCopy{},
+      make_tensor(make_gmem_ptr(static_cast<Element const*>(nullptr)),
+                  WShapeT{},
+                  WStrideT{}),
+      SmemLayoutA{}(_, _, _0{}),
+      select<0, 1>(Shape<Int<kBlockM / 2>, Int<kBlockK>>{}),
+      size<0>(ClusterShape{})));
 
-    using EShapeT = cute::Shape<int64_t, int64_t, int64_t, int64_t, int64_t>;
-    using EStrideT = cute::Shape<_1, int64_t, int64_t, int64_t, int64_t>;
-    using ELayoutT = cute::Layout<EShapeT, EStrideT>;
+  using TMA_B = decltype(make_tma_copy(
+      GmemTiledCopy{},
+      make_tensor(make_gmem_ptr(static_cast<Element const*>(nullptr)),
+                  ShapeT{},
+                  StrideT{}),
+      take<0, 2>(SmemLayoutB{}),
+      select<1, 2>(TileShape_MNK{}),
+      size<0>(ClusterShape{})));
 
-    using TMA_A = decltype(make_tma_copy(
+  using TMA_E = decltype(make_tma_copy(
+      GmemTiledCopy{},
+      make_tensor(make_gmem_ptr(static_cast<uint32_t const*>(nullptr)),
+                  EShapeT{},
+                  EStrideT{}),
+      SmemLayoutE{}(_, _, _0{}),
+      select<0, 1>(Shape<Int<NumMmaThreads>, Int<kBlockK / 64>>{}),
+      size<0>(ClusterShape{})));
+
+  using MainloopPipeline = typename Ktraits::MainloopPipeline;
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState = typename MainloopPipeline::PipelineState;
+
+  static constexpr uint32_t TmaTransactionBytesA = static_cast<uint32_t>(
+      size(take<0, 2>(SmemLayoutA{})) * cutlass::sizeof_bits_v<Element> / 8);
+  static constexpr uint32_t TmaTransactionBytesB = static_cast<uint32_t>(
+      size(take<0, 2>(SmemLayoutB{})) * cutlass::sizeof_bits_v<Element> / 8);
+  static constexpr uint32_t TmaTransactionBytesE = static_cast<uint32_t>(
+      size(take<0, 2>(SmemLayoutE{})) * cutlass::sizeof_bits_v<int> / 8);
+
+  struct Arguments {
+    Element const* ptr_A;
+    WLayoutT layout_A;
+    uint32_t const* ptr_E;
+    ELayoutT layout_E;
+    Element const* ptr_B;
+    LayoutT layout_B;
+    ElementOutput* ptr_C;
+    LayoutT layout_C;
+    const int* tokens;
+    const float* weight_scale;
+  };
+
+  struct Params {
+    WLayoutT layout_A;
+    ELayoutT layout_E;
+    LayoutT layout_B;
+    TMA_A tma_load_A;
+    TMA_E tma_load_E;
+    TMA_B tma_load_B;
+    const int* tokens;
+    const float* weight_scale;
+    ElementOutput* ptr_C;
+  };
+
+  Params static to_underlying_arguments(Arguments const& args) {
+    Tensor mA = make_tensor(make_gmem_ptr(args.ptr_A), args.layout_A);
+    TMA_A tma_load_A =
+        make_tma_copy(GmemTiledCopy{},
+                      mA,
+                      SmemLayoutA{}(_, _, _0{}),
+                      select<0, 1>(Shape<Int<kBlockM / 2>, Int<kBlockK>>{}),
+                      size<0>(ClusterShape{}));
+    Tensor mE = make_tensor(make_gmem_ptr(args.ptr_E), args.layout_E);
+    TMA_E tma_load_E = make_tma_copy(
         GmemTiledCopy{},
-        make_tensor(
-            make_gmem_ptr(static_cast<Element const*>(nullptr)),
-            WShapeT{},
-            WStrideT{}
-        ),
-        SmemLayoutA{}(_, _, _0{}),
-        select<0, 1>(Shape<Int<kBlockM / 2>, Int<kBlockK>>{}),
-        size<0>(ClusterShape{})));
-
-    using TMA_B = decltype(make_tma_copy(
-        GmemTiledCopy{},
-        make_tensor(
-            make_gmem_ptr(static_cast<Element const*>(nullptr)),
-            ShapeT{},
-            StrideT{}
-        ),
-        take<0, 2>(SmemLayoutB{}),
-        select<1, 2>(TileShape_MNK{}),
-        size<0>(ClusterShape{})));
-
-    using TMA_E = decltype(make_tma_copy(
-        GmemTiledCopy{},
-        make_tensor(
-            make_gmem_ptr(static_cast<uint32_t const*>(nullptr)),
-            EShapeT{},
-            EStrideT{}
-        ),
+        mE,
         SmemLayoutE{}(_, _, _0{}),
         select<0, 1>(Shape<Int<NumMmaThreads>, Int<kBlockK / 64>>{}),
-        size<0>(ClusterShape{})));
+        size<0>(ClusterShape{}));
+    Tensor mB = make_tensor(make_gmem_ptr(args.ptr_B), args.layout_B);
+    TMA_B tma_load_B = make_tma_copy(GmemTiledCopy{},
+                                     mB,
+                                     SmemLayoutB{}(_, _, _0{}),
+                                     select<1, 2>(TileShape_MNK{}),
+                                     size<0>(ClusterShape{}));
 
-    using MainloopPipeline = typename Ktraits::MainloopPipeline;
-    using PipelineParams = typename MainloopPipeline::Params;
-    using PipelineState = typename MainloopPipeline::PipelineState;
+    return {args.layout_A,
+            args.layout_E,
+            args.layout_B,
+            tma_load_A,
+            tma_load_E,
+            tma_load_B,
+            args.tokens,
+            args.weight_scale,
+            args.ptr_C};
+  }
 
-    static constexpr uint32_t TmaTransactionBytesA = static_cast<uint32_t>(size(take<0, 2>(SmemLayoutA{})) * cutlass::sizeof_bits_v<Element> / 8);
-    static constexpr uint32_t TmaTransactionBytesB = static_cast<uint32_t>(size(take<0, 2>(SmemLayoutB{})) * cutlass::sizeof_bits_v<Element> / 8);
-    static constexpr uint32_t TmaTransactionBytesE = static_cast<uint32_t>(size(take<0, 2>(SmemLayoutE{})) * cutlass::sizeof_bits_v<int> / 8);
+  CUTLASS_DEVICE
+  static void prefetch_tma_descriptors(Params const& mainloop_params) {
+    cute::prefetch_tma_descriptor(
+        mainloop_params.tma_load_A.get_tma_descriptor());
+    cute::prefetch_tma_descriptor(
+        mainloop_params.tma_load_B.get_tma_descriptor());
+    cute::prefetch_tma_descriptor(
+        mainloop_params.tma_load_E.get_tma_descriptor());
+  }
 
-    struct Arguments {
-        Element const* ptr_A;
-        WLayoutT layout_A;
-        uint32_t const* ptr_E;
-        ELayoutT layout_E;
-        Element const* ptr_B;
-        LayoutT layout_B;
-        ElementOutput * ptr_C;
-        LayoutT layout_C;
-        const int *tokens;
-        const float *weight_scale;
-    };
+  template <int CUR_N, typename SharedStorage>
+  CUTLASS_DEVICE void store(Params const& mainloop_params,
+                            float* acc_s,
+                            SharedStorage& shared_storage,
+                            const int pre_fix_tokens,
+                            const int tokens,
+                            const float* weight_scale,
+                            const int bidm,
+                            const int bidn,
+                            const int bidb,
+                            const int tidx) {
+    typename Ktraits::TiledMma tiled_mma;
+    using packHalf = typename PackedHalf<ElementOutput>::Type;
+    Tensor tOrO_out = make_tensor<ElementOutput>(
+        partition_fragment_C(tiled_mma, select<0, 1>(TileShape_MNK{}))
+            .layout());
 
-    struct Params {
-        WLayoutT layout_A;
-        ELayoutT layout_E;
-        LayoutT layout_B;
-        TMA_A tma_load_A;
-        TMA_E tma_load_E;
-        TMA_B tma_load_B;
-        const int *tokens;
-        const float *weight_scale;
-        ElementOutput * ptr_C;
-    };
-
-
-    Params static
-    to_underlying_arguments(Arguments const& args) {
-        Tensor mA = make_tensor(make_gmem_ptr(args.ptr_A), args.layout_A);
-        TMA_A tma_load_A = make_tma_copy(
-            GmemTiledCopy{},
-            mA,
-            SmemLayoutA{}(_, _, _0{}),
-            select<0, 1>(Shape<Int<kBlockM / 2>, Int<kBlockK>>{}),
-            size<0>(ClusterShape{}));
-        Tensor mE = make_tensor(make_gmem_ptr(args.ptr_E), args.layout_E);
-        TMA_E tma_load_E = make_tma_copy(
-            GmemTiledCopy{},
-            mE,
-            SmemLayoutE{}(_, _, _0{}),
-            select<0, 1>(Shape<Int<NumMmaThreads>, Int<kBlockK / 64>>{}),
-            size<0>(ClusterShape{}));
-        Tensor mB = make_tensor(make_gmem_ptr(args.ptr_B), args.layout_B);
-        TMA_B tma_load_B = make_tma_copy(
-            GmemTiledCopy{},
-            mB,
-            SmemLayoutB{}(_, _, _0{}),
-            select<1, 2>(TileShape_MNK{}),
-            size<0>(ClusterShape{}));
-
-        return {args.layout_A, args.layout_E, args.layout_B,
-            tma_load_A, tma_load_E, tma_load_B,
-            args.tokens, args.weight_scale, args.ptr_C};
+#pragma unroll
+    for (int i = 0; i < size(tOrO_out); i += 4) {
+      acc_s[i] *= weight_scale[0];
+      acc_s[i + 1] *= weight_scale[0];
+      acc_s[i + 2] *= weight_scale[1];
+      acc_s[i + 3] *= weight_scale[1];
+      *reinterpret_cast<packHalf*>(&tOrO_out[i]) =
+          packHalf(acc_s[i], acc_s[i + 2]);
+      *reinterpret_cast<packHalf*>(&tOrO_out[i + 2]) =
+          packHalf(acc_s[i + 1], acc_s[i + 3]);
     }
 
-    CUTLASS_DEVICE
-    static void prefetch_tma_descriptors(Params const& mainloop_params) {
-        cute::prefetch_tma_descriptor(mainloop_params.tma_load_A.get_tma_descriptor());
-        cute::prefetch_tma_descriptor(mainloop_params.tma_load_B.get_tma_descriptor());
-        cute::prefetch_tma_descriptor(mainloop_params.tma_load_E.get_tma_descriptor());
+    uint16_t* smem_c =
+        reinterpret_cast<uint16_t*>(shared_storage.smem_c.data());
+
+    uint32_t* reg_data = reinterpret_cast<uint32_t*>(tOrO_out.data());
+
+    cutlass::arch::NamedBarrier::sync(NumMmaThreads, 0);
+
+    constexpr int k_copy_times = CUR_N / 16;
+
+#pragma unroll
+    for (int i = 0; i < k_copy_times; i++) {
+      uint32_t smem_ptr = cast_smem_ptr_to_uint(
+          reinterpret_cast<uint128_t*>(smem_c + i * 16 * 128) + tidx);
+#if defined(CUTE_ARCH_STSM_SM90_ENABLED)
+      asm volatile(
+          "stmatrix.sync.aligned.x4.trans.m8n8.shared.b16 [%0], {%1, %2, %3, "
+          "%4};\n" ::"r"(smem_ptr),
+          "r"(reg_data[4 * i + 0]),
+          "r"(reg_data[4 * i + 2]),
+          "r"(reg_data[4 * i + 1]),
+          "r"(reg_data[4 * i + 3]));
+#endif
     }
 
-    template <int CUR_N, typename SharedStorage>
-    CUTLASS_DEVICE void
-    store(Params const& mainloop_params,
-        float * acc_s,
-        SharedStorage& shared_storage,
-        const int pre_fix_tokens,
-        const int tokens,
-        const float * weight_scale,
-        const int bidm,
-        const int bidn,
-        const int bidb,
-        const int tidx) {
-        typename Ktraits::TiledMma tiled_mma;
-        using packHalf = typename PackedHalf<ElementOutput>::Type;
-        Tensor tOrO_out = make_tensor<ElementOutput>(partition_fragment_C(tiled_mma, select<0, 1>(TileShape_MNK{})).layout());
+    cutlass::arch::NamedBarrier::sync(NumMmaThreads, 0);
+    const int batch_idx =
+        TokenPackSize == 0 ? pre_fix_tokens * M : bidb * M * TokenPackSize;
+    ElementOutput* store_c = mainloop_params.ptr_C + batch_idx +
+                             bidn * (M * kBlockN) + bidm * kBlockM;
 
-        #pragma unroll
-        for (int i = 0; i < size(tOrO_out); i+=4) {
-            acc_s[i] *= weight_scale[0];
-            acc_s[i + 1] *= weight_scale[0];
-            acc_s[i + 2] *= weight_scale[1];
-            acc_s[i + 3] *= weight_scale[1];
-            *reinterpret_cast<packHalf*>(&tOrO_out[i]) = packHalf(acc_s[i], acc_s[i + 2]);
-            *reinterpret_cast<packHalf*>(&tOrO_out[i + 2]) = packHalf(acc_s[i + 1], acc_s[i + 3]);
-        }
+    const int reamin_tokens = tokens - bidn * kBlockN;
 
-        uint16_t *smem_c = reinterpret_cast<uint16_t *>(shared_storage.smem_c.data());
+    const int col = tidx % 2;
 
-        uint32_t * reg_data = reinterpret_cast<uint32_t*>(tOrO_out.data());
-
-        cutlass::arch::NamedBarrier::sync(NumMmaThreads, 0);
-
-        constexpr int k_copy_times = CUR_N / 16;
-
-        #pragma unroll
-        for (int i = 0; i < k_copy_times; i++) {
-            uint32_t smem_ptr = cast_smem_ptr_to_uint(reinterpret_cast<uint128_t*>(smem_c + i * 16 * 128) + tidx);
-            #if defined(CUTE_ARCH_STSM_SM90_ENABLED)
-            asm volatile (
-                "stmatrix.sync.aligned.x4.trans.m8n8.shared.b16 [%0], {%1, %2, %3, %4};\n"
-                :: "r"(smem_ptr), "r"(reg_data[4 * i + 0]), "r"(reg_data[4 * i + 2]), "r"(reg_data[4 * i + 1]), "r"(reg_data[4 * i + 3]));
-            #endif
-        }
-
-        cutlass::arch::NamedBarrier::sync(NumMmaThreads, 0);
-        const int batch_idx = TokenPackSize == 0 ? pre_fix_tokens * M : bidb * M * TokenPackSize;
-        ElementOutput * store_c = mainloop_params.ptr_C + batch_idx + bidn * (M * kBlockN) + bidm * kBlockM;
-
-        const int reamin_tokens = tokens - bidn * kBlockN;
-
-        const int col = tidx % 2;
-
-        constexpr int kPackSize = 16 / sizeof(ElementOutput);
-        constexpr int kNumVecElem = kBlockM / kPackSize;
-        constexpr int copy_len = CUR_N * kNumVecElem;
-        #pragma unroll
-        for (int idx = tidx; idx < copy_len; idx += NumMmaThreads) {
-            const int idx_div2 = idx / 2;
-            const int store_idx = idx_div2 / 128 * 128 + idx_div2 % 8 * 16 + idx_div2 % 128 / 16 + idx_div2 % 16 / 8 * 8;
-            const int store_global_idx = store_idx * 2 + col;
-            const int row = store_global_idx / kNumVecElem;
-            const int col = store_global_idx % kNumVecElem;
-            if (row >= reamin_tokens) {
-                continue;
-            }
-            const int offset = row * (M / kPackSize) + col;
-            reinterpret_cast<uint4*>(store_c)[offset] = reinterpret_cast<uint4*>(smem_c)[idx];
-        }
+    constexpr int kPackSize = 16 / sizeof(ElementOutput);
+    constexpr int kNumVecElem = kBlockM / kPackSize;
+    constexpr int copy_len = CUR_N * kNumVecElem;
+#pragma unroll
+    for (int idx = tidx; idx < copy_len; idx += NumMmaThreads) {
+      const int idx_div2 = idx / 2;
+      const int store_idx = idx_div2 / 128 * 128 + idx_div2 % 8 * 16 +
+                            idx_div2 % 128 / 16 + idx_div2 % 16 / 8 * 8;
+      const int store_global_idx = store_idx * 2 + col;
+      const int row = store_global_idx / kNumVecElem;
+      const int col = store_global_idx % kNumVecElem;
+      if (row >= reamin_tokens) {
+        continue;
+      }
+      const int offset = row * (M / kPackSize) + col;
+      reinterpret_cast<uint4*>(store_c)[offset] =
+          reinterpret_cast<uint4*>(smem_c)[idx];
     }
+  }
 
-    template <typename MTensor>
-    CUTLASS_DEVICE auto get_local_packed_tensor(
-        const MTensor &mB,
-        const int tokens,
-        const int bidn) const {
+  template <typename MTensor>
+  CUTLASS_DEVICE auto get_local_packed_tensor(const MTensor& mB,
+                                              const int tokens,
+                                              const int bidn) const {
+    auto mB_this_batch = make_tensor(
+        mB.data(),
+        make_layout(cute::make_shape(tokens, size<1>(mB)), mB.stride()));
+    return local_tile(
+        mB_this_batch, select<1, 2>(TileShape_MNK{}), make_coord(bidn, _));
+  }
 
-        auto mB_this_batch = make_tensor(
-            mB.data(),
-            make_layout(
-                cute::make_shape(tokens, size<1>(mB)),
-                mB.stride()
-            ));
-        return local_tile(mB_this_batch, select<1, 2>(TileShape_MNK{}), make_coord(bidn, _));
-    }
+  template <typename MTensor>
+  CUTLASS_DEVICE auto get_local_no_packed_tensor(const MTensor& mB,
+                                                 const int pre_fix_token,
+                                                 const int actual_token,
+                                                 const int bidn) const {
+    auto g_offset = local_tile(mB(_, _, 0),
+                               cute::make_shape(1, size<1>(mB)),
+                               make_coord(pre_fix_token, _0{}));
 
-    template <typename MTensor>
-    CUTLASS_DEVICE auto get_local_no_packed_tensor(
-        const MTensor &mB,
-        const int pre_fix_token,
-        const int actual_token,
-        const int bidn) const {
+    auto g_tensor =
+        make_tensor(g_offset.data(),
+                    make_layout(cute::make_shape(actual_token, size<1>(mB)),
+                                g_offset.stride()));
 
-        auto g_offset = local_tile(
-            mB(_, _, 0),
-            cute::make_shape(1, size<1>(mB)),
-            make_coord(pre_fix_token, _0{}));
+    Tensor gB = local_tile(
+        g_tensor, select<1, 2>(TileShape_MNK{}), make_coord(bidn, _));
 
-        auto g_tensor = make_tensor(
-            g_offset.data(),
-            make_layout(
-                cute::make_shape(actual_token, size<1>(mB)),
-                g_offset.stride()
-            ));
+    return gB;
+  }
 
-        Tensor gB = local_tile(g_tensor, select<1, 2>(TileShape_MNK{}), make_coord(bidn, _));
+  template <typename SharedStorage>
+  CUTLASS_DEVICE void load(Params const& mainloop_params,
+                           MainloopPipeline pipeline,
+                           PipelineState& smem_pipe_write,
+                           SharedStorage& shared_storage,
+                           const int pre_fix_tokens,
+                           const int tokens,
+                           const int bidm,
+                           const int bidn,
+                           const int bidb,
+                           const int tidx) {
+    Tensor sA =
+        make_tensor(make_smem_ptr(shared_storage.smem_a.data()), SmemLayoutA{});
+    Tensor sB =
+        make_tensor(make_smem_ptr(shared_storage.smem_b.data()), SmemLayoutB{});
+    Tensor sE =
+        make_tensor(make_smem_ptr(shared_storage.smem_e.data()), SmemLayoutE{});
 
-        return gB;
-    }
+    Tensor mA = mainloop_params.tma_load_A.get_tma_tensor(
+        mainloop_params.layout_A.shape());
+    Tensor mB = mainloop_params.tma_load_B.get_tma_tensor(
+        mainloop_params.layout_B.shape());
+    Tensor mE = mainloop_params.tma_load_E.get_tma_tensor(
+        mainloop_params.layout_E.shape());
 
+    Tensor gA =
+        local_tile(mA(_, _, _, bidm, bidb),
+                   select<0, 1>(Shape<Int<kBlockM / 2>, Int<kBlockK>>{}),
+                   make_coord(0, 0, _));
 
-    template <typename SharedStorage>
-    CUTLASS_DEVICE void
-    load(Params const& mainloop_params,
-         MainloopPipeline pipeline,
-         PipelineState& smem_pipe_write,
-         SharedStorage &shared_storage,
-         const int pre_fix_tokens,
-         const int tokens,
-         const int bidm,
-         const int bidn,
-         const int bidb,
-         const int tidx) {
+    Tensor gE =
+        local_tile(mE(_, _, _, bidm, bidb),
+                   select<0, 1>(Shape<Int<NumMmaThreads>, Int<kBlockK / 64>>{}),
+                   make_coord(0, 0));
 
-        Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem_a.data()), SmemLayoutA{});
-        Tensor sB = make_tensor(make_smem_ptr(shared_storage.smem_b.data()), SmemLayoutB{});
-        Tensor sE = make_tensor(make_smem_ptr(shared_storage.smem_e.data()), SmemLayoutE{});
+    auto [tAgA, tAsA] = tma_partition(mainloop_params.tma_load_A,
+                                      _0{},
+                                      Layout<ClusterShape>{},
+                                      group_modes<0, 2>(sA),
+                                      group_modes<0, 2>(gA));
 
-        Tensor mA = mainloop_params.tma_load_A.get_tma_tensor(mainloop_params.layout_A.shape());
-        Tensor mB = mainloop_params.tma_load_B.get_tma_tensor(mainloop_params.layout_B.shape());
-        Tensor mE = mainloop_params.tma_load_E.get_tma_tensor(mainloop_params.layout_E.shape());
+    auto [tEgE, tEsE] = tma_partition(mainloop_params.tma_load_E,
+                                      _0{},
+                                      Layout<ClusterShape>{},
+                                      group_modes<0, 2>(sE),
+                                      group_modes<0, 2>(gE));
 
-        Tensor gA = local_tile(mA(_, _, _, bidm, bidb), select<0, 1>(Shape<Int<kBlockM / 2>, Int<kBlockK>>{}), make_coord(0,0,_));
+    int lane_predicate = cute::elect_one_sync();
+    int warp_idx_in_warpgroup =
+        __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);
 
-        Tensor gE = local_tile(mE(_, _, _, bidm, bidb), select<0, 1>(Shape<Int<NumMmaThreads>, Int<kBlockK / 64>>{}), make_coord(0, 0));
+    if constexpr (TokenPackSize == 0) {
+      Tensor gB = get_local_no_packed_tensor(mB, pre_fix_tokens, tokens, bidn);
+      auto [tBgB, tBsB] = tma_partition(mainloop_params.tma_load_B,
+                                        _0{},
+                                        Layout<ClusterShape>{},
+                                        group_modes<0, 2>(sB),
+                                        group_modes<0, 2>(gB));
 
-        auto [tAgA, tAsA] = tma_partition(mainloop_params.tma_load_A, _0{}, Layout<ClusterShape>{}, group_modes<0, 2>(sA), group_modes<0, 2>(gA));
-
-        auto [tEgE, tEsE] = tma_partition(mainloop_params.tma_load_E, _0{}, Layout<ClusterShape>{},  group_modes<0, 2>(sE), group_modes<0, 2>(gE));
-
-        int lane_predicate = cute::elect_one_sync();
-        int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);
-
-        if constexpr (TokenPackSize == 0) {
-            Tensor gB = get_local_no_packed_tensor(
-                mB,
-                pre_fix_tokens,
-                tokens,
-                bidn);
-            auto [tBgB, tBsB] = tma_partition(mainloop_params.tma_load_B, _0{}, Layout<ClusterShape>{}, group_modes<0, 2>(sB), group_modes<0, 2>(gB));
-
-            const int kIters = kTiles / kStages;
-            if (tidx == 0) {
-                #pragma unroll
-                for (int kiter = 0; kiter < kIters; ++kiter) {
-                    #pragma unroll
-                    for (int s = 0; s < kStages; s++) {
-                        const int i = kiter * kStages + s;
-                        pipeline.producer_acquire(smem_pipe_write);
-                        copy(mainloop_params.tma_load_A.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                        tAgA(_, i), tAsA(_, s));
-                        copy(mainloop_params.tma_load_E.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                        tEgE(_, i), tEsE(_, s));
-                        copy(mainloop_params.tma_load_B.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                        tBgB(_, i), tBsB(_, s));
-                        ++smem_pipe_write;
-                    }
-                }
-
-                #pragma unroll
-                for (int i = kIters * kStages; i < kTiles; ++i) {
-                    pipeline.producer_acquire(smem_pipe_write);
-                    copy(mainloop_params.tma_load_A.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                    tAgA(_, i), tAsA(_, smem_pipe_write.index()));
-                    copy(mainloop_params.tma_load_E.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                    tEgE(_, i), tEsE(_, smem_pipe_write.index()));
-                    copy(mainloop_params.tma_load_B.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                    tBgB(_, i), tBsB(_, smem_pipe_write.index()));
-                    ++smem_pipe_write;
-                }
-            }
-        } else {
-            auto mB_this_batch = make_tensor(
-                mB(_, _, bidb).data(),
-                make_layout(
-                    cute::make_shape(tokens, size<1>(mB)),
-                    mB.stride()
-                ));
-            Tensor gB = local_tile(mB_this_batch, select<1, 2>(TileShape_MNK{}), make_coord(bidn, _));
-            auto [tBgB, tBsB] = tma_partition(mainloop_params.tma_load_B, _0{}, Layout<ClusterShape>{}, group_modes<0, 2>(sB), group_modes<0, 2>(gB));
-
-            const int kIters = kTiles / kStages;
-            if (tidx == 0) {
-                #pragma unroll
-                for (int kiter = 0; kiter < kIters; ++kiter) {
-                    #pragma unroll
-                    for (int s = 0; s < kStages; s++) {
-                        const int i = kiter * kStages + s;
-                        pipeline.producer_acquire(smem_pipe_write);
-                        copy(mainloop_params.tma_load_A.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                        tAgA(_, i), tAsA(_, s));
-                        copy(mainloop_params.tma_load_E.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                        tEgE(_, i), tEsE(_, s));
-                        copy(mainloop_params.tma_load_B.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                        tBgB(_, i), tBsB(_, s));
-                        ++smem_pipe_write;
-                    }
-                }
-
-                #pragma unroll
-                for (int i = kIters * kStages; i < kTiles; ++i) {
-                    pipeline.producer_acquire(smem_pipe_write);
-                    copy(mainloop_params.tma_load_A.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                    tAgA(_, i), tAsA(_, smem_pipe_write.index()));
-                    copy(mainloop_params.tma_load_E.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                    tEgE(_, i), tEsE(_, smem_pipe_write.index()));
-                    copy(mainloop_params.tma_load_B.with(*pipeline.producer_get_barrier(smem_pipe_write), 0),
-                    tBgB(_, i), tBsB(_, smem_pipe_write.index()));
-                    ++smem_pipe_write;
-                }
-            }
-        }
-    }
-
-    template <int CUR_N, typename SharedStorage>
-    CUTLASS_DEVICE void
-    mma(Params const& mainloop_params,
-            MainloopPipeline pipeline,
-            PipelineState& smem_pipe_read,
-            SharedStorage& shared_storage,
-            float *acc_s,
-            const int tidx) {
-
-        using sMemBLayout = std::conditional_t<
-            CUR_N == kBlockN,
-            SmemLayoutB,
-            SmemLayoutB_TAIL
-        >;
-
-        using Mma = std::conditional_t<
-            CUR_N == kBlockN,
-            typename Ktraits::Mma,
-            typename Ktraits::Mma_TAIL
-        >;
-
-        Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem_a.data()), SmemLayoutA{});
-        Tensor sB = make_tensor(make_smem_ptr(shared_storage.smem_b.data()), sMemBLayout{});
-        Tensor sE = make_tensor(make_smem_ptr(shared_storage.smem_e.data()), SmemLayoutE{});
-
-        const int wg_idx = tidx / 128;
-        const int wg_offset = wg_idx * 64;
-
-        auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
-            auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
-            pipeline.consumer_wait(smem_pipe_read, barrier_token);
-        };
-
-        constexpr int E_STEP = kBlockK / 64 * NumMmaThreads;
-        constexpr int B_STEPS = CUR_N == 0 ? 1 : (kBlockN / CUR_N);
-
-        const int kIters = kTiles / kStages;
-        #pragma unroll
+      const int kIters = kTiles / kStages;
+      if (tidx == 0) {
+#pragma unroll
         for (int kiter = 0; kiter < kIters; ++kiter) {
-            #pragma unroll
-            for (int s = 0; s < kStages; s++) {
-                consumer_wait(pipeline, smem_pipe_read);
-
-                gemm<Mma, kBlockK, NumMmaThreads>(
-                    sA(_, _, s).data().get().get() + wg_offset,
-                    sB(_, _, s * B_STEPS).data().get().get(),
-                    acc_s,
-                    shared_storage.smem_e.data() + s * E_STEP + tidx);
-
-                pipeline.consumer_release(smem_pipe_read);
-                ++smem_pipe_read;
-            }
+#pragma unroll
+          for (int s = 0; s < kStages; s++) {
+            const int i = kiter * kStages + s;
+            pipeline.producer_acquire(smem_pipe_write);
+            copy(mainloop_params.tma_load_A.with(
+                     *pipeline.producer_get_barrier(smem_pipe_write), 0),
+                 tAgA(_, i),
+                 tAsA(_, s));
+            copy(mainloop_params.tma_load_E.with(
+                     *pipeline.producer_get_barrier(smem_pipe_write), 0),
+                 tEgE(_, i),
+                 tEsE(_, s));
+            copy(mainloop_params.tma_load_B.with(
+                     *pipeline.producer_get_barrier(smem_pipe_write), 0),
+                 tBgB(_, i),
+                 tBsB(_, s));
+            ++smem_pipe_write;
+          }
         }
 
-        #pragma unroll
-        for (int i = 0; i < kTiles % kStages; ++i) {
-            consumer_wait(pipeline, smem_pipe_read);
-
-            gemm<Mma, kBlockK, NumMmaThreads>(
-                sA(_, _, i).data().get().get() + wg_offset,
-                sB(_, _, i * B_STEPS).data().get().get(),
-                acc_s,
-                shared_storage.smem_e.data() + i * E_STEP + tidx);
-
-            pipeline.consumer_release(smem_pipe_read);
-            ++smem_pipe_read;
+#pragma unroll
+        for (int i = kIters * kStages; i < kTiles; ++i) {
+          pipeline.producer_acquire(smem_pipe_write);
+          copy(mainloop_params.tma_load_A.with(
+                   *pipeline.producer_get_barrier(smem_pipe_write), 0),
+               tAgA(_, i),
+               tAsA(_, smem_pipe_write.index()));
+          copy(mainloop_params.tma_load_E.with(
+                   *pipeline.producer_get_barrier(smem_pipe_write), 0),
+               tEgE(_, i),
+               tEsE(_, smem_pipe_write.index()));
+          copy(mainloop_params.tma_load_B.with(
+                   *pipeline.producer_get_barrier(smem_pipe_write), 0),
+               tBgB(_, i),
+               tBsB(_, smem_pipe_write.index()));
+          ++smem_pipe_write;
         }
+      }
+    } else {
+      auto mB_this_batch = make_tensor(
+          mB(_, _, bidb).data(),
+          make_layout(cute::make_shape(tokens, size<1>(mB)), mB.stride()));
+      Tensor gB = local_tile(
+          mB_this_batch, select<1, 2>(TileShape_MNK{}), make_coord(bidn, _));
+      auto [tBgB, tBsB] = tma_partition(mainloop_params.tma_load_B,
+                                        _0{},
+                                        Layout<ClusterShape>{},
+                                        group_modes<0, 2>(sB),
+                                        group_modes<0, 2>(gB));
+
+      const int kIters = kTiles / kStages;
+      if (tidx == 0) {
+#pragma unroll
+        for (int kiter = 0; kiter < kIters; ++kiter) {
+#pragma unroll
+          for (int s = 0; s < kStages; s++) {
+            const int i = kiter * kStages + s;
+            pipeline.producer_acquire(smem_pipe_write);
+            copy(mainloop_params.tma_load_A.with(
+                     *pipeline.producer_get_barrier(smem_pipe_write), 0),
+                 tAgA(_, i),
+                 tAsA(_, s));
+            copy(mainloop_params.tma_load_E.with(
+                     *pipeline.producer_get_barrier(smem_pipe_write), 0),
+                 tEgE(_, i),
+                 tEsE(_, s));
+            copy(mainloop_params.tma_load_B.with(
+                     *pipeline.producer_get_barrier(smem_pipe_write), 0),
+                 tBgB(_, i),
+                 tBsB(_, s));
+            ++smem_pipe_write;
+          }
+        }
+
+#pragma unroll
+        for (int i = kIters * kStages; i < kTiles; ++i) {
+          pipeline.producer_acquire(smem_pipe_write);
+          copy(mainloop_params.tma_load_A.with(
+                   *pipeline.producer_get_barrier(smem_pipe_write), 0),
+               tAgA(_, i),
+               tAsA(_, smem_pipe_write.index()));
+          copy(mainloop_params.tma_load_E.with(
+                   *pipeline.producer_get_barrier(smem_pipe_write), 0),
+               tEgE(_, i),
+               tEsE(_, smem_pipe_write.index()));
+          copy(mainloop_params.tma_load_B.with(
+                   *pipeline.producer_get_barrier(smem_pipe_write), 0),
+               tBgB(_, i),
+               tBsB(_, smem_pipe_write.index()));
+          ++smem_pipe_write;
+        }
+      }
+    }
+  }
+
+  template <int CUR_N, typename SharedStorage>
+  CUTLASS_DEVICE void mma(Params const& mainloop_params,
+                          MainloopPipeline pipeline,
+                          PipelineState& smem_pipe_read,
+                          SharedStorage& shared_storage,
+                          float* acc_s,
+                          const int tidx) {
+    using sMemBLayout =
+        std::conditional_t<CUR_N == kBlockN, SmemLayoutB, SmemLayoutB_TAIL>;
+
+    using Mma = std::conditional_t<CUR_N == kBlockN,
+                                   typename Ktraits::Mma,
+                                   typename Ktraits::Mma_TAIL>;
+
+    Tensor sA =
+        make_tensor(make_smem_ptr(shared_storage.smem_a.data()), SmemLayoutA{});
+    Tensor sB =
+        make_tensor(make_smem_ptr(shared_storage.smem_b.data()), sMemBLayout{});
+    Tensor sE =
+        make_tensor(make_smem_ptr(shared_storage.smem_e.data()), SmemLayoutE{});
+
+    const int wg_idx = tidx / 128;
+    const int wg_offset = wg_idx * 64;
+
+    auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
+      auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+      pipeline.consumer_wait(smem_pipe_read, barrier_token);
+    };
+
+    constexpr int E_STEP = kBlockK / 64 * NumMmaThreads;
+    constexpr int B_STEPS = CUR_N == 0 ? 1 : (kBlockN / CUR_N);
+
+    const int kIters = kTiles / kStages;
+#pragma unroll
+    for (int kiter = 0; kiter < kIters; ++kiter) {
+#pragma unroll
+      for (int s = 0; s < kStages; s++) {
+        consumer_wait(pipeline, smem_pipe_read);
+
+        gemm<Mma, kBlockK, NumMmaThreads>(
+            sA(_, _, s).data().get().get() + wg_offset,
+            sB(_, _, s * B_STEPS).data().get().get(),
+            acc_s,
+            shared_storage.smem_e.data() + s * E_STEP + tidx);
+
+        pipeline.consumer_release(smem_pipe_read);
+        ++smem_pipe_read;
+      }
     }
 
+#pragma unroll
+    for (int i = 0; i < kTiles % kStages; ++i) {
+      consumer_wait(pipeline, smem_pipe_read);
+
+      gemm<Mma, kBlockK, NumMmaThreads>(
+          sA(_, _, i).data().get().get() + wg_offset,
+          sB(_, _, i * B_STEPS).data().get().get(),
+          acc_s,
+          shared_storage.smem_e.data() + i * E_STEP + tidx);
+
+      pipeline.consumer_release(smem_pipe_read);
+      ++smem_pipe_read;
+    }
+  }
 };
diff --git a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/utils.hpp b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/utils.hpp
index 33551e9196..47449d66e5 100644
--- a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/utils.hpp
+++ b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/utils.hpp
@@ -32,69 +32,71 @@
 #include <cutlass/numeric_conversion.h>
 #include <cutlass/numeric_types.h>
 
-
 using namespace cute;
 
-
-template<typename T>
+template <typename T>
 struct PackedHalf;
 
-template<>
+template <>
 struct PackedHalf<cutlass::half_t> {
-    using Type = __half2;
+  using Type = __half2;
 };
 
-template<>
+template <>
 struct PackedHalf<cutlass::bfloat16_t> {
-    using Type = nv_bfloat162;
+  using Type = nv_bfloat162;
 };
 
 template <class PointerType>
-__device__ GmmaDescriptor make_smem_desc(
-        PointerType smem_ptr,
-        int layout_type,
-        int leading_byte_offset = 0,
-        int stride_byte_offset = 1024) {
-
-    GmmaDescriptor desc;
-    auto uint_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
-    desc.bitfield.start_address_ = uint_ptr >> 4;
-    desc.bitfield.layout_type_ = layout_type;
-    desc.bitfield.leading_byte_offset_ = leading_byte_offset >> 4;
-    desc.bitfield.stride_byte_offset_ = stride_byte_offset >> 4;
-    desc.bitfield.base_offset_ = 0;
-    return desc;
+__device__ GmmaDescriptor make_smem_desc(PointerType smem_ptr,
+                                         int layout_type,
+                                         int leading_byte_offset = 0,
+                                         int stride_byte_offset = 1024) {
+  GmmaDescriptor desc;
+  auto uint_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  desc.bitfield.start_address_ = uint_ptr >> 4;
+  desc.bitfield.layout_type_ = layout_type;
+  desc.bitfield.leading_byte_offset_ = leading_byte_offset >> 4;
+  desc.bitfield.stride_byte_offset_ = stride_byte_offset >> 4;
+  desc.bitfield.base_offset_ = 0;
+  return desc;
 }
 
-template <typename Mma, size_t ...Idx>
-__forceinline__ __device__ static void gemm(uint64_t const& desc_a, uint64_t const& desc_b, float* d, const uint32_t e, std::index_sequence<Idx...>) {
-    Mma::fma(desc_a, desc_b, d[Idx]..., e, GMMA::ScaleOut::One);
+template <typename Mma, size_t... Idx>
+__forceinline__ __device__ static void gemm(uint64_t const& desc_a,
+                                            uint64_t const& desc_b,
+                                            float* d,
+                                            const uint32_t e,
+                                            std::index_sequence<Idx...>) {
+  Mma::fma(desc_a, desc_b, d[Idx]..., e, GMMA::ScaleOut::One);
 }
 
 template <typename Mma, int kBlockK, int NumMmaThreads, typename T>
-__forceinline__ __device__ void gemm(
-        const T * sA,
-        const T * sB,
-        float * acc_c,
-        const uint32_t *E) {
+__forceinline__ __device__ void gemm(const T* sA,
+                                     const T* sB,
+                                     float* acc_c,
+                                     const uint32_t* E) {
+  constexpr int acc_num = sizeof(Mma::CRegisters) / sizeof(float);
 
-    constexpr int acc_num = sizeof(Mma::CRegisters) / sizeof(float);
+  warpgroup_arrive();
+// 选择的下标   对应的16进制
+//    01          4
+//    02          8
+//    03          12
+//    12          9
+//    13          13
+//    23          14
+#pragma unroll
+  for (int i = 0; i < kBlockK / 64; i++) {
+    GmmaDescriptor a_desc = make_smem_desc(sA + i * 32, 1, 0, 1024);
+    GmmaDescriptor b_desc = make_smem_desc(sB + i * 64, 1, 0, 1024);
+    gemm<Mma>(a_desc,
+              b_desc,
+              acc_c,
+              E[i * NumMmaThreads],
+              std::make_index_sequence<acc_num>{});
+  }
 
-    warpgroup_arrive();
-    // 选择的下标   对应的16进制
-    //    01          4
-    //    02          8
-    //    03          12
-    //    12          9
-    //    13          13
-    //    23          14
-    #pragma unroll
-    for (int i = 0; i < kBlockK / 64; i++) {
-        GmmaDescriptor a_desc = make_smem_desc(sA + i * 32, 1, 0, 1024);
-        GmmaDescriptor b_desc = make_smem_desc(sB + i * 64, 1, 0, 1024);
-        gemm<Mma>(a_desc, b_desc, acc_c, E[i * NumMmaThreads], std::make_index_sequence<acc_num>{});
-    }
-
-    warpgroup_commit_batch();
-    warpgroup_wait<0>();
+  warpgroup_commit_batch();
+  warpgroup_wait<0>();
 }
diff --git a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/w8a8_sparse_gemm_kernel.hpp b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/w8a8_sparse_gemm_kernel.hpp
index c86cba01d6..3be9c1c619 100644
--- a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/w8a8_sparse_gemm_kernel.hpp
+++ b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/w8a8_sparse_gemm_kernel.hpp
@@ -27,283 +27,307 @@
 #include "mainloop_fwd.h"
 
 template <typename Ktraits>
-void  __global__ __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp, 1) w8a8_sparse_gemm_kernel(
-        CUTE_GRID_CONSTANT typename CollectiveMainloopFwd<Ktraits>::Params const mainloop_params) {
+void __global__ __launch_bounds__(Ktraits::kNWarps *cutlass::NumThreadsPerWarp,
+                                  1)
+    w8a8_sparse_gemm_kernel(
+        CUTE_GRID_CONSTANT
+        typename CollectiveMainloopFwd<Ktraits>::Params const mainloop_params) {
+  using Element = typename Ktraits::Element;
+  static_assert(cutlass::sizeof_bits_v<Element> == 8);
 
-    using Element = typename Ktraits::Element;
-    static_assert(cutlass::sizeof_bits_v<Element> == 8);
+  using TileShape_MNK = typename Ktraits::TileShape_MNK;
+  using ClusterShape = typename Ktraits::ClusterShape_MNK;
 
-    using TileShape_MNK = typename Ktraits::TileShape_MNK;
-    using ClusterShape = typename Ktraits::ClusterShape_MNK;
+  static constexpr int NumMmaThreads = size(typename Ktraits::TiledMma{});
+  static constexpr int NumCopyThreads = cutlass::NumThreadsPerWarpGroup;
+  static constexpr int TokenPackSize = Ktraits::TokenPackSize;
+  static constexpr int kBlockM = Ktraits::kBlockM;
+  static constexpr int kBlockN = Ktraits::kBlockN;
+  static constexpr int TAIL_N = Ktraits::TAIL_N;
+  static constexpr int M = Ktraits::M;
 
-    static constexpr int NumMmaThreads = size(typename Ktraits::TiledMma{});
-    static constexpr int NumCopyThreads = cutlass::NumThreadsPerWarpGroup;
-    static constexpr int TokenPackSize = Ktraits::TokenPackSize;
-    static constexpr int kBlockM = Ktraits::kBlockM;
-    static constexpr int kBlockN = Ktraits::kBlockN;
-    static constexpr int TAIL_N = Ktraits::TAIL_N;
-    static constexpr int M = Ktraits::M;
+  using CollectiveMainloop = CollectiveMainloopFwd<Ktraits>;
 
-    using CollectiveMainloop = CollectiveMainloopFwd<Ktraits>;
+  using MainloopPipeline = typename Ktraits::MainloopPipeline;
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState = typename MainloopPipeline::PipelineState;
 
-    using MainloopPipeline = typename Ktraits::MainloopPipeline;
-    using PipelineParams = typename MainloopPipeline::Params;
-    using PipelineState = typename MainloopPipeline::PipelineState;
+  extern __shared__ char shared_memory[];
+  auto &shared_storage =
+      *reinterpret_cast<typename Ktraits::SharedStorage *>(shared_memory);
 
-    extern __shared__ char shared_memory[];
-    auto &shared_storage = *reinterpret_cast<typename Ktraits::SharedStorage*>(shared_memory);
+  int const lane_predicate = cute::elect_one_sync();
+  int const warp_idx = cutlass::canonical_warp_idx_sync();
 
-    int const lane_predicate = cute::elect_one_sync();
-    int const warp_idx = cutlass::canonical_warp_idx_sync();
+  if (warp_idx == 0 && lane_predicate) {
+    CollectiveMainloop::prefetch_tma_descriptors(mainloop_params);
+  }
 
-    if (warp_idx == 0 && lane_predicate) {
-        CollectiveMainloop::prefetch_tma_descriptors(mainloop_params);
+  // Obtain warp index
+  int const warp_group_thread_idx =
+      threadIdx.x % cutlass::NumThreadsPerWarpGroup;
+
+  PipelineParams pipeline_params;
+  pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesA +
+                                      CollectiveMainloop::TmaTransactionBytesE +
+                                      CollectiveMainloop::TmaTransactionBytesB;
+  int warp_group_idx = cutlass::canonical_warp_group_idx();
+  pipeline_params.role = warp_group_idx == 0
+                             ? MainloopPipeline::ThreadCategory::Producer
+                             : MainloopPipeline::ThreadCategory::Consumer;
+  pipeline_params.is_leader = warp_group_thread_idx == 0;
+  pipeline_params.num_consumers = NumMmaThreads;
+
+  MainloopPipeline pipeline(
+      shared_storage.pipeline, pipeline_params, ClusterShape{});
+
+  pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesB;
+
+  CollectiveMainloop collective_mainloop;
+
+  if constexpr (size(ClusterShape{}) > 1) {
+    cute::cluster_arrive_relaxed();
+    cute::cluster_wait();
+  } else {
+    __syncthreads();
+  }
+
+  const int bidm = blockIdx.x;
+  const int bidn = blockIdx.y;
+  const int bidb = blockIdx.z;
+  const int tidx = threadIdx.x;
+
+  const int pre_fix_tokens =
+      TokenPackSize == 0 ? mainloop_params.tokens[bidb] : 0;
+
+  const int tokens = TokenPackSize == 0
+                         ? mainloop_params.tokens[bidb + 1] - pre_fix_tokens
+                         : mainloop_params.tokens[bidb];
+
+  if (bidn * kBlockN >= tokens) {
+    return;
+  }
+
+  if (warp_group_idx == 0) {
+    cutlass::arch::warpgroup_reg_dealloc<40>();
+    PipelineState smem_pipe_write =
+        cutlass::make_producer_start_state<MainloopPipeline>();
+    collective_mainloop.load(mainloop_params,
+                             pipeline,
+                             smem_pipe_write,
+                             shared_storage,
+                             pre_fix_tokens,
+                             tokens,
+                             bidm,
+                             bidn,
+                             bidb,
+                             tidx);
+  } else {
+    cutlass::arch::warpgroup_reg_alloc<232>();
+    PipelineState smem_pipe_read;
+
+    constexpr int acc_num =
+        sizeof(typename Ktraits::Mma::CRegisters) / sizeof(float);
+    float acc_s[acc_num];
+
+#pragma unroll
+    for (int i = 0; i < acc_num; ++i) {
+      acc_s[i] = 0.0f;
     }
 
-    // Obtain warp index
-    int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup;
+    const int reamin_tokens = tokens - bidn * kBlockN;
 
-    PipelineParams pipeline_params;
-    pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesA + CollectiveMainloop::TmaTransactionBytesE + CollectiveMainloop::TmaTransactionBytesB;
-    int warp_group_idx = cutlass::canonical_warp_group_idx();
-    pipeline_params.role = warp_group_idx == 0
-        ? MainloopPipeline::ThreadCategory::Producer
-        : MainloopPipeline::ThreadCategory::Consumer;
-    pipeline_params.is_leader = warp_group_thread_idx == 0;
-    pipeline_params.num_consumers = NumMmaThreads;
+    const int mma_tidx = tidx - NumCopyThreads;
 
-    MainloopPipeline pipeline(shared_storage.pipeline, pipeline_params, ClusterShape{});
+    const float2 weight_scale = reinterpret_cast<const float2 *>(
+        mainloop_params.weight_scale + bidb * M + bidm * kBlockM)[mma_tidx / 4];
 
-    pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesB;
+    if (TAIL_N > 0 && reamin_tokens < kBlockN) {
+      collective_mainloop.mma<TAIL_N>(mainloop_params,
+                                      pipeline,
+                                      smem_pipe_read,
+                                      shared_storage,
+                                      acc_s,
+                                      mma_tidx);
 
-    CollectiveMainloop collective_mainloop;
-
-    if constexpr (size(ClusterShape{}) > 1) {
-        cute::cluster_arrive_relaxed();
-        cute::cluster_wait();
+      collective_mainloop.store<TAIL_N>(
+          mainloop_params,
+          acc_s,
+          shared_storage,
+          pre_fix_tokens,
+          tokens,
+          reinterpret_cast<const float *>(&weight_scale),
+          bidm,
+          bidn,
+          bidb,
+          mma_tidx);
     } else {
-        __syncthreads();
+      collective_mainloop.mma<kBlockN>(mainloop_params,
+                                       pipeline,
+                                       smem_pipe_read,
+                                       shared_storage,
+                                       acc_s,
+                                       mma_tidx);
+
+      collective_mainloop.store<kBlockN>(
+          mainloop_params,
+          acc_s,
+          shared_storage,
+          pre_fix_tokens,
+          tokens,
+          reinterpret_cast<const float *>(&weight_scale),
+          bidm,
+          bidn,
+          bidb,
+          mma_tidx);
     }
-
-
-    const int bidm = blockIdx.x;
-    const int bidn = blockIdx.y;
-    const int bidb = blockIdx.z;
-    const int tidx = threadIdx.x;
-
-    const int pre_fix_tokens = TokenPackSize == 0 ? mainloop_params.tokens[bidb] : 0;
-
-    const int tokens = TokenPackSize == 0 ? mainloop_params.tokens[bidb + 1] - pre_fix_tokens : mainloop_params.tokens[bidb];
-
-
-    if (bidn * kBlockN >= tokens) {
-        return;
-    }
-
-    if (warp_group_idx == 0) {
-        cutlass::arch::warpgroup_reg_dealloc<40>();
-        PipelineState smem_pipe_write = cutlass::make_producer_start_state<MainloopPipeline>();
-        collective_mainloop.load(
-                mainloop_params,
-                pipeline,
-                smem_pipe_write,
-                shared_storage,
-                pre_fix_tokens,
-                tokens,
-                bidm,
-                bidn,
-                bidb,
-                tidx);
-    } else {
-        cutlass::arch::warpgroup_reg_alloc<232>();
-        PipelineState smem_pipe_read;
-
-        constexpr int acc_num = sizeof(typename Ktraits::Mma::CRegisters) / sizeof(float);
-        float acc_s[acc_num];
-
-        #pragma unroll
-        for (int i = 0; i < acc_num; ++i) {
-            acc_s[i] = 0.0f;
-        }
-
-        const int reamin_tokens = tokens - bidn * kBlockN;
-
-        const int mma_tidx = tidx - NumCopyThreads;
-
-        const float2 weight_scale = reinterpret_cast<const float2*>(mainloop_params.weight_scale + bidb * M + bidm * kBlockM)[mma_tidx / 4];
-
-
-        if (TAIL_N > 0 && reamin_tokens < kBlockN) {
-            collective_mainloop.mma<TAIL_N>(
-                mainloop_params,
-                pipeline,
-                smem_pipe_read,
-                shared_storage,
-                acc_s,
-                mma_tidx);
-
-            collective_mainloop.store<TAIL_N>(
-                mainloop_params,
-                acc_s,
-                shared_storage,
-                pre_fix_tokens,
-                tokens,
-                reinterpret_cast<const float*>(&weight_scale),
-                bidm,
-                bidn,
-                bidb,
-                mma_tidx);
-        } else {
-            collective_mainloop.mma<kBlockN>(
-                mainloop_params,
-                pipeline,
-                smem_pipe_read,
-                shared_storage,
-                acc_s,
-                mma_tidx);
-
-            collective_mainloop.store<kBlockN>(
-                mainloop_params,
-                acc_s,
-                shared_storage,
-                pre_fix_tokens,
-                tokens,
-                reinterpret_cast<const float*>(&weight_scale),
-                bidm,
-                bidn,
-                bidb,
-                mma_tidx);
-        }
-    }
-
+  }
 }
 
 template <int Batch>
 auto get_gmem_layout(int Rows, int Cols) {
-    return  make_layout(
-                make_shape(
-                    static_cast<int64_t>(Rows),
-                    static_cast<int64_t>(Cols),
-                    static_cast<int64_t>(Batch)),
-                make_stride(
-                    static_cast<int64_t>(Cols),
-                    cute::_1{},
-                    static_cast<int64_t>(Rows * Cols)));
+  return make_layout(make_shape(static_cast<int64_t>(Rows),
+                                static_cast<int64_t>(Cols),
+                                static_cast<int64_t>(Batch)),
+                     make_stride(static_cast<int64_t>(Cols),
+                                 cute::_1{},
+                                 static_cast<int64_t>(Rows * Cols)));
 }
 
 template <int Batch>
 auto get_weight_gmem_layout(int m_nums, int k_nums, int Rows, int Cols) {
-    return  make_layout(
-                make_shape(
-                    static_cast<int64_t>(Rows),
-                    static_cast<int64_t>(Cols),
-                    static_cast<int64_t>(k_nums),
-                    static_cast<int64_t>(m_nums),
-                    static_cast<int64_t>(Batch)),
-                make_stride(
-                    static_cast<int64_t>(Cols),
-                    cute::_1{},
-                    static_cast<int64_t>(Rows * Cols),
-                    static_cast<int64_t>(Rows * Cols * k_nums),
-                    static_cast<int64_t>(Rows * Cols * k_nums * m_nums)));
+  return make_layout(
+      make_shape(static_cast<int64_t>(Rows),
+                 static_cast<int64_t>(Cols),
+                 static_cast<int64_t>(k_nums),
+                 static_cast<int64_t>(m_nums),
+                 static_cast<int64_t>(Batch)),
+      make_stride(static_cast<int64_t>(Cols),
+                  cute::_1{},
+                  static_cast<int64_t>(Rows * Cols),
+                  static_cast<int64_t>(Rows * Cols * k_nums),
+                  static_cast<int64_t>(Rows * Cols * k_nums * m_nums)));
 }
 
 template <int Batch>
 auto get_gmem_e_layout(int ms, int ks, int ks_in, int Cols) {
-    return  make_layout(
-                make_shape(
-                    static_cast<int64_t>(Cols),
-                    static_cast<int64_t>(ks_in),
-                    static_cast<int64_t>(ks),
-                    static_cast<int64_t>(ms),
-                    static_cast<int64_t>(Batch)),
-                make_stride(
-                    cute::_1{},
-                    static_cast<int64_t>(Cols),
-                    static_cast<int64_t>(ks_in * Cols),
-                    static_cast<int64_t>(ks * ks_in * Cols),
-                    static_cast<int64_t>(ms * ks * Cols * 2)));
+  return make_layout(make_shape(static_cast<int64_t>(Cols),
+                                static_cast<int64_t>(ks_in),
+                                static_cast<int64_t>(ks),
+                                static_cast<int64_t>(ms),
+                                static_cast<int64_t>(Batch)),
+                     make_stride(cute::_1{},
+                                 static_cast<int64_t>(Cols),
+                                 static_cast<int64_t>(ks_in * Cols),
+                                 static_cast<int64_t>(ks * ks_in * Cols),
+                                 static_cast<int64_t>(ms * ks * Cols * 2)));
 }
 
-template <typename InputType, typename OutputType, typename Kernel_traits, int M, int K, int Batch, int kPackTokenSize>
-void run_gemm(
-        const InputType * A,
-        const uint32_t *E,
-        const InputType * B,
-        OutputType * C,
-        const float *weight_scale,
-        const int *tokens_idx,
-        const int max_tokens,
-        cudaStream_t stream) {
+template <typename InputType,
+          typename OutputType,
+          typename Kernel_traits,
+          int M,
+          int K,
+          int Batch,
+          int kPackTokenSize>
+void run_gemm(const InputType *A,
+              const uint32_t *E,
+              const InputType *B,
+              OutputType *C,
+              const float *weight_scale,
+              const int *tokens_idx,
+              const int max_tokens,
+              cudaStream_t stream) {
+  using ElementOutput = typename Kernel_traits::ElementOutput;
+  using Element = typename Kernel_traits::Element;
+  using CollectiveMainloop = CollectiveMainloopFwd<Kernel_traits>;
+  using ClusterShape = typename Kernel_traits::ClusterShape_MNK;
+  constexpr int NumMmaThreads = Kernel_traits::NumMmaThreads;
+  constexpr int kBlockK = Kernel_traits::kBlockK;
+  constexpr int kBlockM = Kernel_traits::kBlockM;
 
-    using ElementOutput = typename Kernel_traits::ElementOutput;
-    using Element = typename Kernel_traits::Element;
-    using CollectiveMainloop = CollectiveMainloopFwd<Kernel_traits>;
-    using ClusterShape = typename Kernel_traits::ClusterShape_MNK;
-    constexpr int NumMmaThreads = Kernel_traits::NumMmaThreads;
-    constexpr int kBlockK = Kernel_traits::kBlockK;
-    constexpr int kBlockM = Kernel_traits::kBlockM;
+  static_assert(M % Kernel_traits::kBlockM == 0);
+  constexpr int M_nums = M / Kernel_traits::kBlockM;
+  const int N_nums =
+      (max_tokens + Kernel_traits::kBlockN - 1) / Kernel_traits::kBlockN;
 
-    static_assert(M % Kernel_traits::kBlockM == 0);
-    constexpr int M_nums = M / Kernel_traits::kBlockM;
-    const int N_nums = (max_tokens + Kernel_traits::kBlockN - 1) / Kernel_traits::kBlockN;
+  constexpr int kTiles = Kernel_traits::kTiles;
 
-    constexpr int kTiles = Kernel_traits::kTiles;
+  typename CollectiveMainloop::Params mainloop_params =
+      CollectiveMainloop::to_underlying_arguments({
+          static_cast<Element const *>(A),
+          get_weight_gmem_layout<Batch>(M_nums, kTiles, kBlockM / 2, kBlockK),
+          static_cast<uint32_t const *>(E),
+          get_gmem_e_layout<Batch>(M_nums, kTiles, kBlockK / 64, NumMmaThreads),
+          static_cast<Element const *>(B),
+          get_gmem_layout<Batch>(
+              kPackTokenSize == 0 ? max_tokens * Batch : kPackTokenSize, K),
+          static_cast<ElementOutput *>(C),
+          get_gmem_layout<Batch>(
+              M, kPackTokenSize == 0 ? max_tokens : kPackTokenSize),
+          tokens_idx,
+          weight_scale,
+      });
 
-    typename CollectiveMainloop::Params mainloop_params =
-        CollectiveMainloop::to_underlying_arguments({
-            static_cast<Element const*>(A),
-            get_weight_gmem_layout<Batch>(M_nums, kTiles, kBlockM / 2, kBlockK),
-            static_cast<uint32_t const*>(E),
-            get_gmem_e_layout<Batch>(M_nums, kTiles, kBlockK / 64, NumMmaThreads),
-            static_cast<Element const*>(B),
-            get_gmem_layout<Batch>(kPackTokenSize == 0 ? max_tokens * Batch : kPackTokenSize, K),
-            static_cast<ElementOutput*>(C),
-            get_gmem_layout<Batch>(M, kPackTokenSize == 0 ? max_tokens : kPackTokenSize),
-            tokens_idx,
-            weight_scale,
-        });
+  void *kernel;
+  kernel = (void *)w8a8_sparse_gemm_kernel<Kernel_traits>;
 
-    void *kernel;
-    kernel = (void *)w8a8_sparse_gemm_kernel<Kernel_traits>;
+  int smem_size = sizeof(typename Kernel_traits::SharedStorage);
 
-    int smem_size = sizeof(typename Kernel_traits::SharedStorage);
+  if (smem_size >= 48 * 1024) {
+    cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  }
 
-    if (smem_size >= 48 * 1024) {
-       cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-    }
-
-    dim3 grid_dims;
-    grid_dims.x = M_nums;
-    grid_dims.y = N_nums;
-    grid_dims.z = Batch;
-    static constexpr int ctaSize = Kernel_traits::kNWarps * 32;
-    dim3 block_dims(ctaSize);
-    dim3 cluster_dims(size<0>(ClusterShape{}), size<1>(ClusterShape{}), size<2>(ClusterShape{}));
-    cutlass::ClusterLaunchParams launch_params{grid_dims, block_dims, cluster_dims, smem_size, stream};
-    cutlass::launch_kernel_on_cluster(
-        launch_params, kernel, mainloop_params);
+  dim3 grid_dims;
+  grid_dims.x = M_nums;
+  grid_dims.y = N_nums;
+  grid_dims.z = Batch;
+  static constexpr int ctaSize = Kernel_traits::kNWarps * 32;
+  dim3 block_dims(ctaSize);
+  dim3 cluster_dims(size<0>(ClusterShape{}),
+                    size<1>(ClusterShape{}),
+                    size<2>(ClusterShape{}));
+  cutlass::ClusterLaunchParams launch_params{
+      grid_dims, block_dims, cluster_dims, smem_size, stream};
+  cutlass::launch_kernel_on_cluster(launch_params, kernel, mainloop_params);
 }
 
-template <typename InputType, typename OutputType, int M, int K, int Batch, int kPackTokenSize>
-void w8a8_sparse_gemm(
-        const InputType * A,
-        const uint32_t * E,
-        const InputType * B,
-        OutputType * C,
-        const float *weight_scale,
-        const int *tokens_idx,
-        const int max_tokens,
-        cudaStream_t stream) {
-    constexpr static int kBlockM = 128;
-    constexpr static int kBlockK = 128;
-    constexpr static int kNWarps = 4 + kBlockM / 16;
-    constexpr static int kStages = 5;
-    constexpr int kCluster = 1;
-    static_assert(K % kBlockK == 0);
-    constexpr int kTiles = K / kBlockK;
-    const int max_tokens_pack16 = (max_tokens + 31) / 32 * 32;
+template <typename InputType,
+          typename OutputType,
+          int M,
+          int K,
+          int Batch,
+          int kPackTokenSize>
+void w8a8_sparse_gemm(const InputType *A,
+                      const uint32_t *E,
+                      const InputType *B,
+                      OutputType *C,
+                      const float *weight_scale,
+                      const int *tokens_idx,
+                      const int max_tokens,
+                      cudaStream_t stream) {
+  constexpr static int kBlockM = 128;
+  constexpr static int kBlockK = 128;
+  constexpr static int kNWarps = 4 + kBlockM / 16;
+  constexpr static int kStages = 5;
+  constexpr int kCluster = 1;
+  static_assert(K % kBlockK == 0);
+  constexpr int kTiles = K / kBlockK;
+  const int max_tokens_pack16 = (max_tokens + 31) / 32 * 32;
 
-    using Kernel_traits = Kernel_traits<kBlockM, 256, kBlockK, kNWarps, kStages, kTiles, M, kPackTokenSize, 0, kCluster, InputType, OutputType>;
-    run_gemm<InputType, OutputType, Kernel_traits, M, K, Batch, kPackTokenSize>(A, E, B, C, weight_scale, tokens_idx, max_tokens_pack16, stream);
+  using Kernel_traits = Kernel_traits<kBlockM,
+                                      256,
+                                      kBlockK,
+                                      kNWarps,
+                                      kStages,
+                                      kTiles,
+                                      M,
+                                      kPackTokenSize,
+                                      0,
+                                      kCluster,
+                                      InputType,
+                                      OutputType>;
+  run_gemm<InputType, OutputType, Kernel_traits, M, K, Batch, kPackTokenSize>(
+      A, E, B, C, weight_scale, tokens_idx, max_tokens_pack16, stream);
 }
diff --git a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8afp8_sparse_gemm.cu b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8afp8_sparse_gemm.cu
index 03d3c16a1b..e09eb755e2 100644
--- a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8afp8_sparse_gemm.cu
+++ b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8afp8_sparse_gemm.cu
@@ -21,92 +21,93 @@
 #include "wfp8Afp8_sparse_gemm_template.h"
 
 template <typename OutputType>
-void DisPatchWFp8AFp8Gemm(
-        const cutlass::float_e4m3_t* input,
-        const uint32_t* sparse_idx,
-        const cutlass::float_e4m3_t* weight,
-        const int * tokens,
-        const float * weight_scale,
-        OutputType * out,
-        const int token_padding_size,
-        const int max_tokens,
-        const int batch_size,
-        const int M,
-        const int K,
-        cudaStream_t stream) {
+void DisPatchWFp8AFp8Gemm(const cutlass::float_e4m3_t* input,
+                          const uint32_t* sparse_idx,
+                          const cutlass::float_e4m3_t* weight,
+                          const int* tokens,
+                          const float* weight_scale,
+                          OutputType* out,
+                          const int token_padding_size,
+                          const int max_tokens,
+                          const int batch_size,
+                          const int M,
+                          const int K,
+                          cudaStream_t stream) {
+  const int max_tokens_pack32 = (max_tokens + 31) / 32 * 32;
 
-    const int max_tokens_pack32 = (max_tokens + 31) / 32 * 32;
-
-    int kBlockN = 256;
-    int TailN = max_tokens_pack32 % kBlockN;
-    if (max_tokens < 256) {
-        kBlockN = max_tokens_pack32;
-        TailN = 0;
-    }
-    if constexpr (std::is_same_v<OutputType, cutlass::bfloat16_t>) {
-        SPARSE_GEMM_SWITCH_BF16(M, K, batch_size, token_padding_size, kBlockN, TailN,
-            weight,
-            sparse_idx,
-            input,
-            out,
-            weight_scale,
-            tokens,
-            max_tokens,
-            stream)
-    } else {
-        PD_THROW("Only supported dtype in ['BFLOAT16'].");
-    }
+  int kBlockN = 256;
+  int TailN = max_tokens_pack32 % kBlockN;
+  if (max_tokens < 256) {
+    kBlockN = max_tokens_pack32;
+    TailN = 0;
+  }
+  if constexpr (std::is_same_v<OutputType, cutlass::bfloat16_t>) {
+    SPARSE_GEMM_SWITCH_BF16(M,
+                            K,
+                            batch_size,
+                            token_padding_size,
+                            kBlockN,
+                            TailN,
+                            weight,
+                            sparse_idx,
+                            input,
+                            out,
+                            weight_scale,
+                            tokens,
+                            max_tokens,
+                            stream)
+  } else {
+    PD_THROW("Only supported dtype in ['BFLOAT16'].");
+  }
 }
 
-void WFp8AFp8Gemm(
-        const paddle::Tensor& input,
-        const paddle::Tensor& sparse_idx,
-        const paddle::Tensor& weight,
-        const paddle::Tensor& tokens, // If tokenpadding=0, this tensor represents the prefix sum of tensors, otherwise it represents the number of tokens in each group
-        const paddle::Tensor& weight_scale,
-        const paddle::Tensor& out,
-        const int token_padding_size,
-        const int max_tokens,
-        const bool is_bfloat16) {
+void WFp8AFp8Gemm(const paddle::Tensor& input,
+                  const paddle::Tensor& sparse_idx,
+                  const paddle::Tensor& weight,
+                  const paddle::Tensor&
+                      tokens,  // If tokenpadding=0, this tensor represents the
+                               // prefix sum of tensors, otherwise it represents
+                               // the number of tokens in each group
+                  const paddle::Tensor& weight_scale,
+                  const paddle::Tensor& out,
+                  const int token_padding_size,
+                  const int max_tokens,
+                  const bool is_bfloat16) {
+  const int batch_size = weight.dims()[0];
+  const int M = weight.dims()[1];
+  const int K = weight.dims()[2] * 2;
 
-    const int batch_size = weight.dims()[0];
-    const int M = weight.dims()[1];
-    const int K = weight.dims()[2] * 2;
+  if (input.dtype() != paddle::DataType::FLOAT8_E4M3FN) {
+    PD_THROW("Only supported dtype in ['FLOAT8_E4M3FN'].");
+  }
 
-    if (input.dtype() != paddle::DataType::FLOAT8_E4M3FN) {
-        PD_THROW("Only supported dtype in ['FLOAT8_E4M3FN'].");
-    }
-
-    if (is_bfloat16) {
-        DisPatchWFp8AFp8Gemm(
-            reinterpret_cast<const cutlass::float_e4m3_t*>(input.data<phi::dtype::float8_e4m3fn>()),
-            reinterpret_cast<const uint32_t*>(sparse_idx.data<int32_t>()),
-            reinterpret_cast<const cutlass::float_e4m3_t*>(weight.data<phi::dtype::float8_e4m3fn>()),
-            tokens.data<int>(),
-            weight_scale.data<float>(),
-            reinterpret_cast<cutlass::bfloat16_t*>(const_cast<phi::dtype::bfloat16*>(out.data<phi::dtype::bfloat16>())),
-            token_padding_size,
-            max_tokens,
-            batch_size,
-            M,
-            K,
-            input.stream()
-        );
-    } else {
-        PD_THROW("Only supported dtype in ['BFLOAT16'].");
-    }
+  if (is_bfloat16) {
+    DisPatchWFp8AFp8Gemm(
+        reinterpret_cast<const cutlass::float_e4m3_t*>(
+            input.data<phi::dtype::float8_e4m3fn>()),
+        reinterpret_cast<const uint32_t*>(sparse_idx.data<int32_t>()),
+        reinterpret_cast<const cutlass::float_e4m3_t*>(
+            weight.data<phi::dtype::float8_e4m3fn>()),
+        tokens.data<int>(),
+        weight_scale.data<float>(),
+        reinterpret_cast<cutlass::bfloat16_t*>(
+            const_cast<phi::dtype::bfloat16*>(
+                out.data<phi::dtype::bfloat16>())),
+        token_padding_size,
+        max_tokens,
+        batch_size,
+        M,
+        K,
+        input.stream());
+  } else {
+    PD_THROW("Only supported dtype in ['BFLOAT16'].");
+  }
 }
 
 PD_BUILD_STATIC_OP(wfp8afp8_sparse_gemm)
-    .Inputs({"input",
-             "sparse_idx",
-             "weight",
-             "tokens",
-             "weight_scale",
-             "ffn_out"})
+    .Inputs(
+        {"input", "sparse_idx", "weight", "tokens", "weight_scale", "ffn_out"})
     .Outputs({"out"})
     .SetInplaceMap({{"ffn_out", "out"}})
-    .Attrs({"token_padding_size: int",
-            "max_tokens: int",
-            "is_bfloat16: bool"})
+    .Attrs({"token_padding_size: int", "max_tokens: int", "is_bfloat16: bool"})
     .SetKernelFn(PD_KERNEL(WFp8AFp8Gemm));
diff --git a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8afp8_sparse_gemm_weight.cu b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8afp8_sparse_gemm_weight.cu
index 9871dea4b6..35851429f6 100644
--- a/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8afp8_sparse_gemm_weight.cu
+++ b/custom_ops/gpu_ops/wfp8afp8_sparse_gemm/wfp8afp8_sparse_gemm_weight.cu
@@ -19,78 +19,95 @@
 #include "helper.h"
 #include "paddle/extension.h"
 
-
-void pack_E(const uint8_t *E_src, int32_t *E_dst, const int M, const int K, const int Batch) {
-    // 选择的下标   对应的16进制
-    //    01          4
-    //    02          8
-    //    03          12
-    //    12          9
-    //    13          13
-    //    23          14
-    const int ld1 = K / 4;
-    const int ld2 = K / 4 / 8;
-    const uint8_t select_idx[6] = {14, 13, 9, 12, 8, 4};
-    for (int b = 0; b < Batch; ++b) {
-        for (int m = 0; m < M; ++m) {
-            for (int k = 0; k < ld1; k+=8) {
-                uint32_t dst = 0;
-                for (int k2 = 7; k2 > 0; --k2) {
-                    dst |= select_idx[E_src[b * M * ld1 + m * ld1 + k + k2]];
-                    dst <<= 4;
-                }
-                dst |= select_idx[E_src[b * M * ld1 + m * ld1 + k]];
-                E_dst[b * M * ld2 + m * ld2 + k / 8] = dst;
-            }
+void pack_E(const uint8_t *E_src,
+            int32_t *E_dst,
+            const int M,
+            const int K,
+            const int Batch) {
+  // 选择的下标   对应的16进制
+  //    01          4
+  //    02          8
+  //    03          12
+  //    12          9
+  //    13          13
+  //    23          14
+  const int ld1 = K / 4;
+  const int ld2 = K / 4 / 8;
+  const uint8_t select_idx[6] = {14, 13, 9, 12, 8, 4};
+  for (int b = 0; b < Batch; ++b) {
+    for (int m = 0; m < M; ++m) {
+      for (int k = 0; k < ld1; k += 8) {
+        uint32_t dst = 0;
+        for (int k2 = 7; k2 > 0; --k2) {
+          dst |= select_idx[E_src[b * M * ld1 + m * ld1 + k + k2]];
+          dst <<= 4;
         }
+        dst |= select_idx[E_src[b * M * ld1 + m * ld1 + k]];
+        E_dst[b * M * ld2 + m * ld2 + k / 8] = dst;
+      }
     }
+  }
 }
 
-void peruate_E(const int32_t *E_src, int32_t *E_dst, const int M, const int K, const int Batch) {
-    const int m_nums = M / 128;
-    const int k_nums = K / 128;
-    for (int b = 0; b < Batch; ++b) {
-        for (int m = 0; m < m_nums; ++m) {
-            for (int k = 0; k < k_nums; ++k) {
-                const int dst_idx = b * m_nums * k_nums * 512 + m * k_nums * 512 + k * 512;
-                for (int i = 0; i < 8; ++i) {
-                    for (int j = 0; j < 8; ++j) {
-                        E_dst[dst_idx + 0 + j * 32 + 4 * i] = E_src[dst_idx + 0 + j * 64 + 4 * i];
-                        E_dst[dst_idx + 2 + j * 32 + 4 * i] = E_src[dst_idx + 1 + j * 64 + 4 * i];
-                        E_dst[dst_idx + 1 + j * 32 + 4 * i] = E_src[dst_idx + 32 + j * 64 + 4 * i];
-                        E_dst[dst_idx + 3 + j * 32 + 4 * i] = E_src[dst_idx + 33 + j * 64 + 4 * i];
-                    }
-                    for (int j = 0; j < 8; ++j) {
-                        E_dst[dst_idx + 256 + j * 32 + 4 * i] = E_src[dst_idx + 2 + j * 64 + 4 * i];
-                        E_dst[dst_idx + 258 + j * 32 + 4 * i] = E_src[dst_idx + 3 + j * 64 + 4 * i];
-                        E_dst[dst_idx + 257 + j * 32 + 4 * i] = E_src[dst_idx + 34 + j * 64 + 4 * i];
-                        E_dst[dst_idx + 259 + j * 32 + 4 * i] = E_src[dst_idx + 35 + j * 64 + 4 * i];
-                    }
-                }
-            }
+void peruate_E(const int32_t *E_src,
+               int32_t *E_dst,
+               const int M,
+               const int K,
+               const int Batch) {
+  const int m_nums = M / 128;
+  const int k_nums = K / 128;
+  for (int b = 0; b < Batch; ++b) {
+    for (int m = 0; m < m_nums; ++m) {
+      for (int k = 0; k < k_nums; ++k) {
+        const int dst_idx =
+            b * m_nums * k_nums * 512 + m * k_nums * 512 + k * 512;
+        for (int i = 0; i < 8; ++i) {
+          for (int j = 0; j < 8; ++j) {
+            E_dst[dst_idx + 0 + j * 32 + 4 * i] =
+                E_src[dst_idx + 0 + j * 64 + 4 * i];
+            E_dst[dst_idx + 2 + j * 32 + 4 * i] =
+                E_src[dst_idx + 1 + j * 64 + 4 * i];
+            E_dst[dst_idx + 1 + j * 32 + 4 * i] =
+                E_src[dst_idx + 32 + j * 64 + 4 * i];
+            E_dst[dst_idx + 3 + j * 32 + 4 * i] =
+                E_src[dst_idx + 33 + j * 64 + 4 * i];
+          }
+          for (int j = 0; j < 8; ++j) {
+            E_dst[dst_idx + 256 + j * 32 + 4 * i] =
+                E_src[dst_idx + 2 + j * 64 + 4 * i];
+            E_dst[dst_idx + 258 + j * 32 + 4 * i] =
+                E_src[dst_idx + 3 + j * 64 + 4 * i];
+            E_dst[dst_idx + 257 + j * 32 + 4 * i] =
+                E_src[dst_idx + 34 + j * 64 + 4 * i];
+            E_dst[dst_idx + 259 + j * 32 + 4 * i] =
+                E_src[dst_idx + 35 + j * 64 + 4 * i];
+          }
         }
+      }
     }
+  }
 }
 
 std::vector<paddle::Tensor> WFp8AFp8GemmSparseIdxConvert(
-        const paddle::Tensor& weight,
-        const int batch_size,
-        const int M,
-        const int K) {
-
-    paddle::Tensor weight_temp = paddle::empty({batch_size, M, K / 32}, paddle::DataType::INT32, weight.place());
-    paddle::Tensor weight_new = paddle::empty({batch_size, M, K / 32}, paddle::DataType::INT32, weight.place());
-    pack_E(weight.data<uint8_t>(), weight_temp.data<int32_t>(), M, K, batch_size);
-    peruate_E(weight_temp.data<int32_t>(), weight_new.data<int32_t>(), M, K, batch_size);
-    return {weight_new};
+    const paddle::Tensor &weight,
+    const int batch_size,
+    const int M,
+    const int K) {
+  paddle::Tensor weight_temp = paddle::empty(
+      {batch_size, M, K / 32}, paddle::DataType::INT32, weight.place());
+  paddle::Tensor weight_new = paddle::empty(
+      {batch_size, M, K / 32}, paddle::DataType::INT32, weight.place());
+  pack_E(weight.data<uint8_t>(), weight_temp.data<int32_t>(), M, K, batch_size);
+  peruate_E(weight_temp.data<int32_t>(),
+            weight_new.data<int32_t>(),
+            M,
+            K,
+            batch_size);
+  return {weight_new};
 }
 
-
-
 PD_BUILD_STATIC_OP(wfp8afp8_gemm_sparse_idx_convert)
     .Inputs({"weight"})
     .Outputs({"converted_weight"})
-    .Attrs({"batch: int",
-            "M: int",
-            "K: int"})
+    .Attrs({"batch: int", "M: int", "K: int"})
     .SetKernelFn(PD_KERNEL(WFp8AFp8GemmSparseIdxConvert));
diff --git a/custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/mtp_kernel/draft_model_postprocess.xpu b/custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/mtp_kernel/draft_model_postprocess.xpu
index 61c21e79d5..d66146db40 100644
--- a/custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/mtp_kernel/draft_model_postprocess.xpu
+++ b/custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/mtp_kernel/draft_model_postprocess.xpu
@@ -109,7 +109,7 @@ __global__ void draft_model_postprocess(const int64_t* base_model_draft_tokens,
     }
     cur_row_to_all_clusetr = min(bsz - offset_loop, clusetr_limit_row);
     offset_cluster = 0;
-    //计算offset_cluster
+    // 计算offset_cluster
     for (int start_cluster = 0; start_cluster < cluster_id();
          start_cluster += 1) {
       offset_cluster += (rounddown_div(cur_row_to_all_clusetr, nclusters) +
diff --git a/custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/quant2d_per_channel.xpu b/custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/quant2d_per_channel.xpu
index 43daf966c2..2bf62a537a 100644
--- a/custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/quant2d_per_channel.xpu
+++ b/custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/quant2d_per_channel.xpu
@@ -3,8 +3,8 @@
 #include "xpu/kernel/cluster_primitive.h"
 #include "xpu/kernel/cluster_primitive_template.h"
 // TODO()
-//#include "xpu/quant_xpu.h"
-//#include "xpu_plugin.h"
+// #include "xpu/quant_xpu.h"
+// #include "xpu_plugin.h"
 namespace xpu3 {
 namespace plugin {
 #define MAX_SM_SIZE 32768