chore: fotmat code

src/core/compute_workspace.cpp

@@ -187,8 +187,8 @@ namespace sd {
                 }
                 copy->buffer = buffer;
                 // Capability checks inspect alignment even before weights are resident.
-                copy->data   = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external));
-                copy->extra  = nullptr;
+                copy->data  = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external));
+                copy->extra = nullptr;
             }
         }
         auto copy_graph = ggml_new_graph_custom(context, graph_size, false);

src/core/cuda/affine.cu

@@ -129,15 +129,15 @@ namespace SD_CUDA_NAMESPACE {
             ggml_ext_gpu_context context;
             if (!cuda_context(execution->backend, &context))
                 return GGML_STATUS_FAILED;
-            auto stream          = static_cast<CUDAStream>(context.stream);
-            const auto* node     = nodes[0];
-            const auto* a        = node->src[0];
-            const auto* b        = node->op == GGML_OP_CONT ? nullptr : node->src[1];
-            auto* dst            = nodes[count - 1];
-            const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
-            const bool gate      = residual != nullptr;
-            const bool alias     = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
-                               (residual && overlaps(dst, residual) && dst->data != residual->data);
+            auto stream           = static_cast<CUDAStream>(context.stream);
+            const auto* node      = nodes[0];
+            const auto* a         = node->src[0];
+            const auto* b         = node->op == GGML_OP_CONT ? nullptr : node->src[1];
+            auto* dst             = nodes[count - 1];
+            const auto* residual  = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
+            const bool gate       = residual != nullptr;
+            const bool alias      = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
+                                    (residual && overlaps(dst, residual) && dst->data != residual->data);
             void* output          = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
             const int64_t count4  = ggml_nelements(dst) / 4;
             const unsigned blocks = unsigned((count4 + 255) / 256);

src/core/cuda/cute_attention.cu

@@ -209,8 +209,8 @@ namespace sd_fattn_cute {
         using Element                    = cutlass::half_t;
         static constexpr int Threads     = Warps * 32;
         using Mma                        = TiledMMA<MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
-                             Layout<Shape<Int<Warps>, _1, _1>>,
-                             Tile<Int<16 * Warps>, _16, _16>>;
+                                                    Layout<Shape<Int<Warps>, _1, _1>>,
+                                                    Tile<Int<16 * Warps>, _16, _16>>;
         static constexpr int SmemWidth   = D % 64 == 0 ? 64 : D % 32 == 0 ? 32
                                                                           : 16;
         static constexpr int SwizzleBits = SmemWidth == 64 ? 3 : SmemWidth == 32 ? 2
@@ -449,7 +449,7 @@ static __global__ void copy_to_half(const char* src, half* dst, int64_t count, i
     if (i < count) {
         const int64_t row = i / dim;
         const auto* p     = reinterpret_cast<const T*>(src + (row % rows) * nb1 +
-                                                   ((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3);
+                                                       ((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3);
         dst[i]            = __float2half_rn(float(p[i % dim]) * scale);
     }
 }
@@ -784,9 +784,9 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
     desc.preferred = nullptr;
     desc.optimize  = nullptr;
     GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
-    desc.name     = "sd.i8_rope_cute_attention.cuda.v1";
+    desc.name      = "sd.i8_rope_cute_attention.cuda.v1";
     desc.op        = GGML_OP_CUSTOM;
     desc.custom_op = sd_ggml_i8_op();
-    desc.priority = 400;
+    desc.priority  = 400;
     GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
 }

src/core/cuda/i8_gemm.cu

@@ -197,8 +197,8 @@ namespace sd_i8_cutlass {
         using GemmBase        = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
         using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
         using Kernel          = std::conditional_t<StreamK,
-                                          cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>,
-                                          cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>;
+                                                   cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>,
+                                                   cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>;
         using Gemm            = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
     };
     struct Plan {

src/core/cuda/int8.cu

@@ -661,10 +661,10 @@ namespace SD_CUDA_NAMESPACE {
 #else
                     constexpr auto compute_type = HIPBLAS_R_32I;
 #endif
-                    const auto status           = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
-                                                                int(p.n), int(p.padded), int(p.k), &alpha,
-                                                                op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k),
-                                                                &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
+                    const auto status = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
+                                                      int(p.n), int(p.padded), int(p.k), &alpha,
+                                                      op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k),
+                                                      &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
                     if (status != HIPBLAS_STATUS_SUCCESS)
 #else
                     const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N,

src/core/cuda/sage_attention.cu

@@ -287,9 +287,9 @@ namespace sd_cuda {
 
         constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D;
         auto kernel                   = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D,
-                                                               sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp,
-                                                               sageattention::QuantGranularity::kPerWarp, float, false, half,
-                                                               sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>;
+                                                                                 sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp,
+                                                                                 sageattention::QuantGranularity::kPerWarp, float, false, half,
+                                                                                 sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>;
         kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), shared_bytes, stream>>>(
             qi, ki, static_cast<half*>(v->data), out, nullptr, qs, ks, nullptr,
             lq, lk, hq / hk,

src/core/cuda/sageattention/cp_async.cuh

@@ -24,17 +24,17 @@
 
 namespace sageattention {
 
-namespace cp_async {
+    namespace cp_async {
 
-enum class SharedMemFillMode {
-  kFillZero,  // Fill zero to shared memory when predicate is false
-  kNoFill     // Do not fill zero to shared memory when predicate is false
-};
+        enum class SharedMemFillMode {
+            kFillZero,  // Fill zero to shared memory when predicate is false
+            kNoFill     // Do not fill zero to shared memory when predicate is false
+        };
 
-enum class PrefetchMode {
-  kNoPrefetch,  // Do not fetch additional data from global memory to L2
-  kPrefetch     // Fetch additional data from global memory to L2
-};
+        enum class PrefetchMode {
+            kNoPrefetch,  // Do not fetch additional data from global memory to L2
+            kPrefetch     // Fetch additional data from global memory to L2
+        };
 
 #if (__CUDACC_VER_MAJOR__ >= 11)
 #if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
@@ -42,107 +42,107 @@ enum class PrefetchMode {
 #endif
 #endif
 
-/*!
- * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
- *   cp.async instructions to a group
- */
-__device__ __forceinline__ void commit_group() {
+        /*!
+         * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
+         *   cp.async instructions to a group
+         */
+        __device__ __forceinline__ void commit_group() {
 #ifdef CP_ASYNC_ENABLED
-  asm volatile("cp.async.commit_group;\n" ::);
+            asm volatile("cp.async.commit_group;\n" ::);
 #endif
-}
+        }
 
-/*!
- * \brief Wrapper of PTX cp.async.wait_group instruction
- * \tparam n Wait till most recent n groups are committed
- */
-template <size_t n>
-__device__ __forceinline__ void wait_group() {
+        /*!
+         * \brief Wrapper of PTX cp.async.wait_group instruction
+         * \tparam n Wait till most recent n groups are committed
+         */
+        template <size_t n>
+        __device__ __forceinline__ void wait_group() {
 #ifdef CP_ASYNC_ENABLED
-  asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
+            asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
 #endif
-}
+        }
 
-/*!
- * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
- *   global memory to shared memory
- * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
- * \tparam T Data type
- * \param smem_ptr Pointer to shared memory
- * \param gmem_ptr Pointer to global memory
- */
-template <PrefetchMode prefetch_mode, typename T>
-__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
+        /*!
+         * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
+         *   global memory to shared memory
+         * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
+         * \tparam T Data type
+         * \param smem_ptr Pointer to shared memory
+         * \param gmem_ptr Pointer to global memory
+         */
+        template <PrefetchMode prefetch_mode, typename T>
+        __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
 #ifdef CP_ASYNC_ENABLED
-  uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
-  if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
-    asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
-                 "l"(gmem_ptr), "n"(16), "r"(16));
-  } else {
-    asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
-                 "l"(gmem_ptr), "n"(16), "r"(16));
-  }
+            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+            if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+                asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+                             "l"(gmem_ptr), "n"(16), "r"(16));
+            } else {
+                asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+                             "l"(gmem_ptr), "n"(16), "r"(16));
+            }
 #else
-  *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
+            *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
 #endif
-}
+        }
 
-/*!
- * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
- *   global memory to shared memory with predicate.
- * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
- * \tparam fill_mode Whether to fill zero to shared memory when predicate is false
- * \tparam T Data type
- * \param smem_ptr Pointer to shared memory
- * \param gmem_ptr Pointer to global memory
- * \param predicate Predicate value
- * \note fill zero is slower than not fill zero
- */
-template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
-__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
+        /*!
+         * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
+         *   global memory to shared memory with predicate.
+         * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
+         * \tparam fill_mode Whether to fill zero to shared memory when predicate is false
+         * \tparam T Data type
+         * \param smem_ptr Pointer to shared memory
+         * \param gmem_ptr Pointer to global memory
+         * \param predicate Predicate value
+         * \note fill zero is slower than not fill zero
+         */
+        template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
+        __device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
 #ifdef CP_ASYNC_ENABLED
-  uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
-  if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
-    int src_in_bytes = predicate ? 16 : 0;
-    if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
-      asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
-                   "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
-    } else {
-      asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
-                   "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
-    }
-  } else {
-    if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
-      asm volatile(
-          "{\n"
-          " .reg .pred p;\n"
-          " setp.ne.b32 p, %0, 0;\n"
-          " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
-          "}\n" ::"r"((int)predicate),
-          "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
-    } else {
-      asm volatile(
-          "{\n"
-          " .reg .pred p;\n"
-          " setp.ne.b32 p, %0, 0;\n"
-          " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
-          "}\n" ::"r"((int)predicate),
-          "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
-    }
-  }
+            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+            if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+                int src_in_bytes = predicate ? 16 : 0;
+                if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+                    asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+                                 "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
+                } else {
+                    asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+                                 "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
+                }
+            } else {
+                if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+                    asm volatile(
+                        "{\n"
+                        " .reg .pred p;\n"
+                        " setp.ne.b32 p, %0, 0;\n"
+                        " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
+                        "}\n" ::"r"((int)predicate),
+                        "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
+                } else {
+                    asm volatile(
+                        "{\n"
+                        " .reg .pred p;\n"
+                        " setp.ne.b32 p, %0, 0;\n"
+                        " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
+                        "}\n" ::"r"((int)predicate),
+                        "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
+                }
+            }
 #else
-  if (predicate) {
-    *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
-  } else {
-    if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
-      *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0);
-    }
-  }
+            if (predicate) {
+                *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
+            } else {
+                if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+                    *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0);
+                }
+            }
 #endif
-}
+        }
 
-} // namespace cp_async
+    }  // namespace cp_async
 
-} // namespace sageattention
+}  // namespace sageattention
 
 #endif

src/core/cuda/sageattention/math.cuh

@@ -26,137 +26,141 @@
 namespace sageattention {
 #ifndef USHORT_TYPE
 #define USHORT_TYPE
-typedef unsigned short ushort;
+    typedef unsigned short ushort;
 #endif
 
-namespace math {
+    namespace math {
 
-// log2(e)
-constexpr float log2e = 1.44269504088896340736f;
-constexpr float log2e_recp = 1.0f / log2e;
+        // log2(e)
+        constexpr float log2e      = 1.44269504088896340736f;
+        constexpr float log2e_recp = 1.0f / log2e;
 
-__forceinline__ __device__ half2 uint32_as_half2(uint32_t x) { return *(half2*)&x; }
+        __forceinline__ __device__ half2 uint32_as_half2(uint32_t x) {
+            return *(half2*)&x;
+        }
 
-__forceinline__ __device__ uint32_t half2_as_uint32(half2 x) { return *(uint32_t*)&x; }
+        __forceinline__ __device__ uint32_t half2_as_uint32(half2 x) {
+            return *(uint32_t*)&x;
+        }
 
-/*!
- * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x
- * \param x input
- */
-__forceinline__ __device__ float ptx_exp2(float x) {
-  float y;
-  asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
-  return y;
-}
+        /*!
+         * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x
+         * \param x input
+         */
+        __forceinline__ __device__ float ptx_exp2(float x) {
+            float y;
+            asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
 
-/*!
- * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x)
- * \param x input
- */
-__forceinline__ __device__ float ptx_log2(float x) {
-  float y;
-  asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
-  return y;
-}
+        /*!
+         * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x)
+         * \param x input
+         */
+        __forceinline__ __device__ float ptx_log2(float x) {
+            float y;
+            asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
 
-/*!
- * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x
- * \param x input
- */
-__forceinline__ __device__ half2 ptx_exp2(half2 x) {
-  uint32_t y_u32;
-  uint32_t x_u32 = half2_as_uint32(x);
-  asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
-  return uint32_as_half2(y_u32);
-}
+        /*!
+         * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x
+         * \param x input
+         */
+        __forceinline__ __device__ half2 ptx_exp2(half2 x) {
+            uint32_t y_u32;
+            uint32_t x_u32 = half2_as_uint32(x);
+            asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
+            return uint32_as_half2(y_u32);
+        }
 
-/*!
- * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x
- * \param x input
- */
-__forceinline__ __device__ half ptx_exp2(half x) {
-  ushort y_u16;
-  asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
-  return __ushort_as_half(y_u16);
-}
+        /*!
+         * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x
+         * \param x input
+         */
+        __forceinline__ __device__ half ptx_exp2(half x) {
+            ushort y_u16;
+            asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
+            return __ushort_as_half(y_u16);
+        }
 
-/*!
- * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x
- * \param x input
- */
-__forceinline__ __device__ float ptx_rcp(float x) {
-  float y;
-  asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
-  return y;
-}
+        /*!
+         * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x
+         * \param x input
+         */
+        __forceinline__ __device__ float ptx_rcp(float x) {
+            float y;
+            asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
 
-/*!
- * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle
- *   between threads in a warp.
- * \param x The value in the source lane
- * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta]
- */
-__forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) {
-  float y;
-  asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;"
-               : "=f"(y)
-               : "f"(x), "r"(lane_mask));
-  return y;
-}
+        /*!
+         * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle
+         *   between threads in a warp.
+         * \param x The value in the source lane
+         * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta]
+         */
+        __forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) {
+            float y;
+            asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;"
+                         : "=f"(y)
+                         : "f"(x), "r"(lane_mask));
+            return y;
+        }
 
-/*!
- * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly
- *   shuffle between threads in a warp.
- * \param x The value in the source lane
- * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask]
- */
-__forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) {
-  return __shfl_xor_sync(0xffffffff, x, lane_mask);
-}
+        /*!
+         * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly
+         *   shuffle between threads in a warp.
+         * \param x The value in the source lane
+         * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask]
+         */
+        __forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) {
+            return __shfl_xor_sync(0xffffffff, x, lane_mask);
+        }
 
-/*!
- * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x)
- * \param x input
- */
-__forceinline__ __device__ float rsqrt(float x) {
-  float y;
-  asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
-  return y;
-}
+        /*!
+         * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x)
+         * \param x input
+         */
+        __forceinline__ __device__ float rsqrt(float x) {
+            float y;
+            asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
 
-/*!
- * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x)
- * \param x input
- */
-__forceinline__ __device__ float tanh(float x) {
-  float y;
-  asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x));
-  return y;
-}
+        /*!
+         * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x)
+         * \param x input
+         */
+        __forceinline__ __device__ float tanh(float x) {
+            float y;
+            asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
 
-/*!
- * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x)
- * \param x input
- */
-__forceinline__ __device__ half2 tanh(half2 x) {
-  uint32_t y_u32;
-  uint32_t x_u32 = half2_as_uint32(x);
-  asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
-  return uint32_as_half2(y_u32);
-}
+        /*!
+         * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x)
+         * \param x input
+         */
+        __forceinline__ __device__ half2 tanh(half2 x) {
+            uint32_t y_u32;
+            uint32_t x_u32 = half2_as_uint32(x);
+            asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
+            return uint32_as_half2(y_u32);
+        }
 
-/*!
- * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x)
- * \param x input
- */
-__forceinline__ __device__ half tanh(half x) {
-  ushort y_u16;
-  asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
-  return __ushort_as_half(y_u16);
-}
+        /*!
+         * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x)
+         * \param x input
+         */
+        __forceinline__ __device__ half tanh(half x) {
+            ushort y_u16;
+            asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
+            return __ushort_as_half(y_u16);
+        }
 
-}  // namespace math
+    }  // namespace math
 
-} // namespace sageattention
+}  // namespace sageattention
 
 #endif

src/core/cuda/sageattention/numeric_conversion.cuh

@@ -38,118 +38,111 @@ namespace sageattention {
 #define RUNTIME_ASSERT(x) assert(0 && x)
 #endif
 
-__device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) {
-  uint16_t h0 = source & 0xFFFF;
-  uint16_t h1 = (source >> 16) & 0xFFFF;
-  asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0));
-  asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1));
-}
+    __device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) {
+        uint16_t h0 = source & 0xFFFF;
+        uint16_t h1 = (source >> 16) & 0xFFFF;
+        asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0));
+        asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1));
+    }
 
-__device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t *dest, float *source0, float *source1)
-{
+    __device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t* dest, float* source0, float* source1) {
 #ifdef FP8_CAST_ENABLED
-  asm volatile( \
-      "{\n" \
-      ".reg .b16 lo;\n" \
-      ".reg .b16 hi;\n" \
-      "cvt.rn.satfinite.e4m3x2.f32   lo, %2, %1;\n" \
-      "cvt.rn.satfinite.e4m3x2.f32   hi, %4, %3;\n" \
-      "mov.b32 %0, {lo, hi};\n" \
-      "}" \
-      : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo;\n"
+            ".reg .b16 hi;\n"
+            "cvt.rn.satfinite.e4m3x2.f32   lo, %2, %1;\n"
+            "cvt.rn.satfinite.e4m3x2.f32   hi, %4, %3;\n"
+            "mov.b32 %0, {lo, hi};\n"
+            "}"
+            : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
 #else
-  RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
 #endif
-}
+    }
 
-__device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t *dest, float *source0, float *source1)
-{
+    __device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t* dest, float* source0, float* source1) {
 #ifdef FP8_CAST_ENABLED
-  asm volatile( \
-      "{\n" \
-      ".reg .b16 lo;\n" \
-      ".reg .b16 hi;\n" \
-      "cvt.rn.satfinite.e5m2x2.f32   lo, %2, %1;\n" \
-      "cvt.rn.satfinite.e5m2x2.f32   hi, %4, %3;\n" \
-      "mov.b32 %0, {lo, hi};\n" \
-      "}" \
-      : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo;\n"
+            ".reg .b16 hi;\n"
+            "cvt.rn.satfinite.e5m2x2.f32   lo, %2, %1;\n"
+            "cvt.rn.satfinite.e5m2x2.f32   hi, %4, %3;\n"
+            "mov.b32 %0, {lo, hi};\n"
+            "}"
+            : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
 #else
-  RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
 #endif
-}
+    }
 
-__device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t *dest, uint32_t *source0, uint32_t *source1)
-{
+    __device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
 #ifdef FP8_CAST_ENABLED
-  asm volatile( \
-      "{\n" \
-      ".reg .b16 lo;\n" \
-      ".reg .b16 hi;\n" \
-      "cvt.rn.satfinite.e4m3x2.f16x2   lo, %1;\n" \
-      "cvt.rn.satfinite.e4m3x2.f16x2   hi, %2;\n" \
-      "mov.b32 %0, {lo, hi};\n" \
-      "}" \
-      : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo;\n"
+            ".reg .b16 hi;\n"
+            "cvt.rn.satfinite.e4m3x2.f16x2   lo, %1;\n"
+            "cvt.rn.satfinite.e4m3x2.f16x2   hi, %2;\n"
+            "mov.b32 %0, {lo, hi};\n"
+            "}"
+            : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
 #else
-  RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
 #endif
-}
+    }
 
-__device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t *dest, uint32_t *source0, uint32_t *source1)
-{
+    __device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
 #ifdef FP8_CAST_ENABLED
-  asm volatile( \
-      "{\n" \
-      ".reg .b16 lo;\n" \
-      ".reg .b16 hi;\n" \
-      "cvt.rn.satfinite.e5m2x2.f16x2   lo, %1;\n" \
-      "cvt.rn.satfinite.e5m2x2.f16x2   hi, %2;\n" \
-      "mov.b32 %0, {lo, hi};\n" \
-      "}" \
-      : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo;\n"
+            ".reg .b16 hi;\n"
+            "cvt.rn.satfinite.e5m2x2.f16x2   lo, %1;\n"
+            "cvt.rn.satfinite.e5m2x2.f16x2   hi, %2;\n"
+            "mov.b32 %0, {lo, hi};\n"
+            "}"
+            : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
 #else
-  RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
 #endif
-}
+    }
 
-__device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source)
-{
+    __device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
 #ifdef FP8_CAST_ENABLED
-  asm volatile( \
-      "{\n" \
-      ".reg .b16 lo, hi;\n" \
-      "mov.b32 {lo, hi}, %2;\n" \
-      "cvt.rn.f16x2.e4m3x2 %0, lo;\n" \
-      "cvt.rn.f16x2.e4m3x2 %1, hi;\n" \
-      "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo, hi;\n"
+            "mov.b32 {lo, hi}, %2;\n"
+            "cvt.rn.f16x2.e4m3x2 %0, lo;\n"
+            "cvt.rn.f16x2.e4m3x2 %1, hi;\n"
+            "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
 #else
-  RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
 #endif
-}
+    }
 
-__device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source)
-{
+    __device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
 #ifdef FP8_CAST_ENABLED
-  asm volatile( \
-      "{\n" \
-      ".reg .b16 lo, hi;\n" \
-      "mov.b32 {lo, hi}, %2;\n" \
-      "cvt.rn.f16x2.e5m2x2 %0, lo;\n" \
-      "cvt.rn.f16x2.e5m2x2 %1, hi;\n" \
-      "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo, hi;\n"
+            "mov.b32 {lo, hi}, %2;\n"
+            "cvt.rn.f16x2.e5m2x2 %0, lo;\n"
+            "cvt.rn.f16x2.e5m2x2 %1, hi;\n"
+            "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
 #else
-  RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
 #endif
-}
+    }
 
-__device__ __forceinline__ int8_t float_to_int8_rn(float x)
-{
-    uint32_t dst;
-    asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
-    return reinterpret_cast<const int8_t&>(dst);
-}
+    __device__ __forceinline__ int8_t float_to_int8_rn(float x) {
+        uint32_t dst;
+        asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
+        return reinterpret_cast<const int8_t&>(dst);
+    }
 
-} // namespace sageattention
+}  // namespace sageattention
 
 #endif

src/core/cuda/sageattention/permuted_smem.cuh

@@ -23,180 +23,179 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
-
 #include "cp_async.cuh"
 #include "mma.cuh"
 
 namespace sageattention {
 
-enum class SwizzleMode {
-  k32B, // for k32B mode, a line of shared memory must have 32B (16 half value)
-  k64B, // for k64B mode, a line of shared memory must have 64B (32 half value)
-  k128B, // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B.
-};
+    enum class SwizzleMode {
+        k32B,   // for k32B mode, a line of shared memory must have 32B (16 half value)
+        k64B,   // for k64B mode, a line of shared memory must have 64B (32 half value)
+        k128B,  // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B.
+    };
 
-// Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth
-using b128_t = uint4;
+    // Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth
+    using b128_t = uint4;
 
-/*!
- * \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure
- * that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through
- * permutation in the granularity of 128bit.
- *
- * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless,
- * meaning it does not maintain any information about the current pointer position. The offset returnd by
- * the struct can be used to access the shared memory through the provided interface.
- *
- * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j).
- * We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp.
- * Under this assumption, we do not permute for k32B mode.
- */
-template <SwizzleMode swizzle_mode, uint32_t stride>
-struct smem_t {
-  // The base pointer.
-  b128_t* base;
-  // How many b128_t value a row contains
-  // uint32_t stride;
+    /*!
+     * \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure
+     * that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through
+     * permutation in the granularity of 128bit.
+     *
+     * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless,
+     * meaning it does not maintain any information about the current pointer position. The offset returnd by
+     * the struct can be used to access the shared memory through the provided interface.
+     *
+     * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j).
+     * We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp.
+     * Under this assumption, we do not permute for k32B mode.
+     */
+    template <SwizzleMode swizzle_mode, uint32_t stride>
+    struct smem_t {
+        // The base pointer.
+        b128_t* base;
+        // How many b128_t value a row contains
+        // uint32_t stride;
 
-  __device__ __forceinline__ smem_t() : base(nullptr) {}
-  template <typename T>
-  __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {
-    if constexpr (swizzle_mode == SwizzleMode::k128B) {
-      static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode");
-    } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
-      static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode");
-    } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
-      static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode");
-    } else {
-      static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode");
-    }
-  }
+        __device__ __forceinline__ smem_t() : base(nullptr) {}
+        template <typename T>
+        __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode");
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode");
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode");
+            } else {
+                static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode");
+            }
+        }
 
-  /*!
-   * \brief Set the base pointer.
-   */
-  template <typename T>
-  __device__ __forceinline__ void set_base(T* new_base) {
-    base = (b128_t*)new_base;
-  }
+        /*!
+         * \brief Set the base pointer.
+         */
+        template <typename T>
+        __device__ __forceinline__ void set_base(T* new_base) {
+            base = (b128_t*)new_base;
+        }
 
-  /*!
-   * \brief Compute the element offset given coordinates in a permuted shared memory.
-   * \param i The row index.
-   * \param j The column index.
-   */
-  static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t &i, const uint32_t &j) {
-    if constexpr (swizzle_mode == SwizzleMode::k128B) {
-      return i * stride + (j ^ (i % 8));
-    } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
-      return i * stride + (j ^ ((i / 2) % 4));
-    } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
-      return i * stride + j;
-    }
-  }
+        /*!
+         * \brief Compute the element offset given coordinates in a permuted shared memory.
+         * \param i The row index.
+         * \param j The column index.
+         */
+        static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t& i, const uint32_t& j) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                return i * stride + (j ^ (i % 8));
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                return i * stride + (j ^ ((i / 2) % 4));
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                return i * stride + j;
+            }
+        }
 
-  /*!
-  * \tparam step_size The step size to advance the offset in the permuted shared memory.
-  * \param offset The current offset.
-  */
-  template <uint32_t step_size>
-  static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset) {
-    if constexpr (swizzle_mode == SwizzleMode::k128B) {
-      static_assert(step_size % 8 == 0,
-                    "Unsupported step size");
-      return offset + step_size;
-    } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
-      static_assert(step_size == 4, "Unsupported step size");
-      return offset + step_size;
-    } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
-      static_assert(step_size == 2, "Unsupported step size");
-      return offset + step_size;
-    }
-  }
+        /*!
+         * \tparam step_size The step size to advance the offset in the permuted shared memory.
+         * \param offset The current offset.
+         */
+        template <uint32_t step_size>
+        static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                static_assert(step_size % 8 == 0,
+                              "Unsupported step size");
+                return offset + step_size;
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                static_assert(step_size == 4, "Unsupported step size");
+                return offset + step_size;
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                static_assert(step_size == 2, "Unsupported step size");
+                return offset + step_size;
+            }
+        }
 
-  // ! use with care
-  template <uint32_t step_size>
-  static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset, const uint32_t &step_idx) {
-    if constexpr (swizzle_mode == SwizzleMode::k128B) {
-      static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0,
-                    "Unsupported step size");
-      if constexpr (step_size == 2) {
-        return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8;
-      } else if constexpr (step_size == 4) {
-        return (offset ^ 0x4) + (step_idx % 2 == 1) * 8;
-      } else {
-        // step_size % 8 == 0
-        return offset + step_size;
-      }
-    } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
-      static_assert(step_size == 2 || step_size == 4, "Unsupported step size");
-      if constexpr (step_size == 2) {
-        return (offset ^ 0x2) + (step_idx % 2 == 1) * 4;
-      } else {
-        return offset + step_size;
-      }
-    } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
-      return offset + step_size;
-    }
-  }
+        // ! use with care
+        template <uint32_t step_size>
+        static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset, const uint32_t& step_idx) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0,
+                              "Unsupported step size");
+                if constexpr (step_size == 2) {
+                    return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8;
+                } else if constexpr (step_size == 4) {
+                    return (offset ^ 0x4) + (step_idx % 2 == 1) * 8;
+                } else {
+                    // step_size % 8 == 0
+                    return offset + step_size;
+                }
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                static_assert(step_size == 2 || step_size == 4, "Unsupported step size");
+                if constexpr (step_size == 2) {
+                    return (offset ^ 0x2) + (step_idx % 2 == 1) * 4;
+                } else {
+                    return offset + step_size;
+                }
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                return offset + step_size;
+            }
+        }
 
-  template <uint32_t step_size>
-  static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t &offset) {
-    if constexpr (swizzle_mode == SwizzleMode::k128B) {
-      static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
-      if constexpr (step_size == 4) {
-        return (offset ^ 0x4) + step_size * stride;
-      } else {
-        // step_size % 8 == 0
-        return offset + step_size * stride;
-      }
-    } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
-      static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
-      if constexpr (step_size == 4) {
-        return (offset ^ 0x2) + step_size * stride;
-      } else {
-        // step_size % 8 == 0
-        return offset + step_size * stride;
-      }
-    } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
-      return offset + step_size * stride;
-    }
-  }
+        template <uint32_t step_size>
+        static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t& offset) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
+                if constexpr (step_size == 4) {
+                    return (offset ^ 0x4) + step_size * stride;
+                } else {
+                    // step_size % 8 == 0
+                    return offset + step_size * stride;
+                }
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
+                if constexpr (step_size == 4) {
+                    return (offset ^ 0x2) + step_size * stride;
+                } else {
+                    // step_size % 8 == 0
+                    return offset + step_size * stride;
+                }
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                return offset + step_size * stride;
+            }
+        }
 
-  __device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t &offset, uint32_t* R) const {
-    b128_t* smem_ptr = base + offset;
-    mma::ldmatrix_m8n8x2(R, smem_ptr);
-  }
+        __device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t& offset, uint32_t* R) const {
+            b128_t* smem_ptr = base + offset;
+            mma::ldmatrix_m8n8x2(R, smem_ptr);
+        }
 
-  __device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t &offset, uint32_t* R) const {
-    b128_t* smem_ptr = base + offset;
-    mma::ldmatrix_m8n8x4(R, smem_ptr);
-  }
+        __device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t& offset, uint32_t* R) const {
+            b128_t* smem_ptr = base + offset;
+            mma::ldmatrix_m8n8x4(R, smem_ptr);
+        }
 
-  __device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t &offset, uint32_t* R) const {
-    b128_t* smem_ptr = base + offset;
-    mma::ldmatrix_m8n8x4_trans(R, smem_ptr);
-  }
+        __device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t& offset, uint32_t* R) const {
+            b128_t* smem_ptr = base + offset;
+            mma::ldmatrix_m8n8x4_trans(R, smem_ptr);
+        }
 
-  template <cp_async::SharedMemFillMode fill_mode, typename T>
-  __device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr, bool predicate) const {
-    b128_t* smem_ptr = base + offset;
-    cp_async::pred_load_128b<cp_async::PrefetchMode::kPrefetch, fill_mode>(
-        smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
-  }
+        template <cp_async::SharedMemFillMode fill_mode, typename T>
+        __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr, bool predicate) const {
+            b128_t* smem_ptr = base + offset;
+            cp_async::pred_load_128b<cp_async::PrefetchMode::kPrefetch, fill_mode>(
+                smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
+        }
 
-  template <typename T>
-  __device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr) const {
-    b128_t* smem_ptr = base + offset;
-    cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr));
-  }
+        template <typename T>
+        __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr) const {
+            b128_t* smem_ptr = base + offset;
+            cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr));
+        }
 
-  template <typename T>
-  __device__ __forceinline__ void store_128b(const uint32_t &offset, T* gptr) const {
-    *reinterpret_cast<b128_t*>(gptr) = *(base + offset);
-  }
-};
+        template <typename T>
+        __device__ __forceinline__ void store_128b(const uint32_t& offset, T* gptr) const {
+            *reinterpret_cast<b128_t*>(gptr) = *(base + offset);
+        }
+    };
 
-} // namespace sageattention
+}  // namespace sageattention
 
 #endif

src/core/ggml_extend.h

@@ -4,9 +4,9 @@
 #include <cstdint>
 #include <vector>
 
+#include "core/ggml_extensions.h"
 #include "ggml-backend.h"
 #include "ggml.h"
-#include "core/ggml_extensions.h"
 
 #define EPS 1e-05f
 

src/core/vulkan/int8.cpp

@@ -613,7 +613,7 @@ namespace sd_vulkan_i8 {
         auto destination                   = concat ? range(dst, ggml_nbytes(concat->src[0]), ggml_nbytes(tail)) : dst;
         ggml_ext_vulkan_binding bindings[] = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}, {op.input, nullptr, scale_offset, size_t(rows) * 4}, {op.bias ? op.bias : op.weight_scale}, destination, {up_mat ? up.weight : gate}, {up_mat ? up.weight_scale : residual}};
         auto dispatch                      = [&](ggml_ext_vulkan_pipeline pipeline, uint32_t tile_cols, uint32_t tile_rows,
-                            std::initializer_list<uint32_t> constants, uint32_t row_offset, uint32_t dispatch_rows) {
+                                                 std::initializer_list<uint32_t> constants, uint32_t row_offset, uint32_t dispatch_rows) {
             if (!dispatch_rows)
                 return GGML_STATUS_SUCCESS;
             const uint32_t cols     = glu ? uint32_t(glu->ne[0]) : n;

src/core/vulkan/shaders/modulate_convrot.comp

@@ -2,16 +2,28 @@
 
 #extension GL_EXT_control_flow_attributes : require
 
-layout(constant_id = 0) const uint K = 4096;
+layout(constant_id = 0) const uint K      = 4096;
 layout(constant_id = 1) const bool AFFINE = false;
 layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
-layout(binding = 0) readonly buffer X { float data_x[]; };
-layout(binding = 1) readonly buffer M { float data_m[]; };
-layout(binding = 2) readonly buffer H { float data_h[]; };
-layout(binding = 3) writeonly buffer D { uint data_d[]; };
-layout(binding = 4) writeonly buffer S { float data_s[]; };
+layout(binding = 0) readonly buffer X {
+    float data_x[];
+};
+layout(binding = 1) readonly buffer M {
+    float data_m[];
+};
+layout(binding = 2) readonly buffer H {
+    float data_h[];
+};
+layout(binding = 3) writeonly buffer D {
+    uint data_d[];
+};
+layout(binding = 4) writeonly buffer S {
+    float data_s[];
+};
 #ifdef RMS_AFFINE
-layout(binding = 5) readonly buffer W { float data_w[]; };
+layout(binding = 5) readonly buffer W {
+    float data_w[];
+};
 #endif
 layout(push_constant) uniform parameter {
     uint rows;
@@ -21,7 +33,8 @@ layout(push_constant) uniform parameter {
     float eps;
     uint k;
     uint w_offset;
-} p;
+}
+p;
 shared float values[K];
 shared vec2 sums[512];
 
@@ -36,7 +49,7 @@ float rms_row(uint row, uint tid) {
     [[unroll]] for (uint col = tid, iter = 0; iter < 12; col += 512, ++iter) {
         float x = 0.0;
         if (col < p.k) {
-            x = data_x[p.x_offset + row * K + col];
+            x           = data_x[p.x_offset + row * K + col];
             values[col] = x;
         }
         sum += x * x;
@@ -63,30 +76,31 @@ void main() {
     for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) {
 #ifdef RMS_AFFINE
         float inv_std = rms_row(row, tid);
-        float mean = 0.0;
+        float mean    = 0.0;
 #else
         vec2 sum = vec2(0.0);
         [[unroll]] for (uint col = tid; col < K; col += 512) {
             const float x = data_x[p.x_offset + row * K + col];
-            values[col] = x;
+            values[col]   = x;
             if (!AFFINE) {
                 sum.x += x;
                 sum.y = fma(x, x, sum.y);
             }
         }
-        float mean = 0.0;
+        float mean    = 0.0;
         float inv_std = 1.0;
         if (!AFFINE) {
             sums[tid] = sum;
             barrier();
             [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) {
-                if (tid < stride) sums[tid] += sums[tid + stride];
+                if (tid < stride)
+                    sums[tid] += sums[tid + stride];
                 barrier();
             }
             mean = sums[0].x / p.k;
             // Match NORM's contraction even when precise propagates from the quantized output.
             const float variance = fma(-mean, mean, sums[0].y / p.k);
-            inv_std = inversesqrt(variance + p.eps);
+            inv_std              = inversesqrt(variance + p.eps);
         }
 #endif
         [[unroll]] for (uint col = tid; col < K; col += 512) {
@@ -96,7 +110,7 @@ void main() {
 #endif
             if (AFFINE) {
                 precise float modulated = values[col] * data_m[p.m_offset + col];
-                value = values[col] + modulated;
+                value                   = values[col] + modulated;
                 value += data_h[p.h_offset + col];
             } else {
                 value = (values[col] - mean) * inv_std;
@@ -107,15 +121,15 @@ void main() {
         barrier();
         [[unroll]] for (uint stride = 1; stride < 256; stride *= 4) {
             for (uint i = tid; i < K / 4; i += 512) {
-                const uint group = i / 64;
-                const uint lane = i % 64;
-                const uint base = group * 256 + (lane / stride) * 4 * stride + lane % stride;
-                const float a = values[base];
-                const float b = values[base + stride];
-                const float c = values[base + 2 * stride];
-                const float d = values[base + 3 * stride];
-                values[base] = a + b + c - d;
-                values[base + stride] = a + b - c + d;
+                const uint group          = i / 64;
+                const uint lane           = i % 64;
+                const uint base           = group * 256 + (lane / stride) * 4 * stride + lane % stride;
+                const float a             = values[base];
+                const float b             = values[base + stride];
+                const float c             = values[base + 2 * stride];
+                const float d             = values[base + 3 * stride];
+                values[base]              = a + b + c - d;
+                values[base + stride]     = a + b - c + d;
                 values[base + 2 * stride] = a - b + c + d;
                 values[base + 3 * stride] = -a + b + c + d;
             }
@@ -133,16 +147,13 @@ void main() {
             }
             barrier();
         }
-        const float scale = sums[0].x / 127.0;
+        const float scale     = sums[0].x / 127.0;
         const float inv_scale = scale == 0.0 ? 0.0 : 1.0 / scale;
         if (tid == 0) {
             data_s[row] = scale;
         }
         for (uint i = tid; i < K / 4; i += 512) {
-            data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale)
-                | (quantize(values[i * 4 + 1], inv_scale) << 8)
-                | (quantize(values[i * 4 + 2], inv_scale) << 16)
-                | (quantize(values[i * 4 + 3], inv_scale) << 24);
+            data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) | (quantize(values[i * 4 + 1], inv_scale) << 8) | (quantize(values[i * 4 + 2], inv_scale) << 16) | (quantize(values[i * 4 + 3], inv_scale) << 24);
         }
         barrier();
     }

src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp

@@ -5,43 +5,55 @@
 
 layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
 
-layout(push_constant) uniform parameter
-{
+layout(push_constant) uniform parameter {
     uint n;
     uint rows;
     uint k;
     uint has_bias;
     uint row_offset;
-} p;
+}
+p;
 
-layout(binding = 0, std430) readonly buffer W { int data_w[]; };
-layout(binding = 1, std430) readonly buffer A { int data_a[]; };
-layout(binding = 2, std430) readonly buffer WS { float data_ws[]; };
-layout(binding = 3, std430) readonly buffer AS { float data_as[]; };
-layout(binding = 4, std430) readonly buffer B { float data_b[]; };
-layout(binding = 5, std430) writeonly buffer D { float data_d[]; };
+layout(binding = 0, std430) readonly buffer W {
+    int data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 5, std430) writeonly buffer D {
+    float data_d[];
+};
 
 shared int weights[16 * 8];
 shared int activations[16 * 8];
 
 void main() {
     const uint output_index = gl_WorkGroupID.x * 16 + gl_LocalInvocationID.x;
-    const uint row = p.row_offset + gl_WorkGroupID.y * 16 + gl_LocalInvocationID.y;
-    const uint lx = gl_LocalInvocationID.x;
-    const uint ly = gl_LocalInvocationID.y;
-    const uint k_packed = p.k / 4;
-    int sum = 0;
+    const uint row          = p.row_offset + gl_WorkGroupID.y * 16 + gl_LocalInvocationID.y;
+    const uint lx           = gl_LocalInvocationID.x;
+    const uint ly           = gl_LocalInvocationID.y;
+    const uint k_packed     = p.k / 4;
+    int sum                 = 0;
 
     for (uint k_base = 0; k_base < k_packed; k_base += 8) {
         if (ly < 8) {
             weights[lx * 8 + ly] = output_index < p.n
-                ? data_w[output_index * k_packed + k_base + ly]
-                : 0;
+                                       ? data_w[output_index * k_packed + k_base + ly]
+                                       : 0;
         }
         if (lx < 8) {
             activations[ly * 8 + lx] = row < p.rows
-                ? data_a[row * k_packed + k_base + lx]
-                : 0;
+                                           ? data_a[row * k_packed + k_base + lx]
+                                           : 0;
         }
         barrier();
 

src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp

@@ -9,35 +9,47 @@
 #extension GL_KHR_shader_subgroup_basic : require
 
 layout(constant_id = 0) const uint BLOCK_SIZE = 32;
-layout(constant_id = 1) const uint CM_M = 16;
-layout(constant_id = 2) const uint CM_N = 16;
-layout(constant_id = 3) const uint CM_K = 16;
-const uint TILE_M = 4;
-const uint TILE_N = 4;
+layout(constant_id = 1) const uint CM_M       = 16;
+layout(constant_id = 2) const uint CM_N       = 16;
+layout(constant_id = 3) const uint CM_K       = 16;
+const uint TILE_M                             = 4;
+const uint TILE_N                             = 4;
 
 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
-layout(push_constant) uniform parameter
-{
+layout(push_constant) uniform parameter {
     uint n;
     uint rows;
     uint k;
     uint has_bias;
     uint row_offset;
-} p;
+}
+p;
 
-layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; };
-layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; };
-layout(binding = 2, std430) readonly buffer WS { float data_ws[]; };
-layout(binding = 3, std430) readonly buffer AS { float data_as[]; };
-layout(binding = 4, std430) readonly buffer B { float data_b[]; };
-layout(binding = 5, std430) writeonly buffer D { float data_d[]; };
+layout(binding = 0, std430) readonly buffer W {
+    int8_t data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int8_t data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 5, std430) writeonly buffer D {
+    float data_d[];
+};
 
 shared int32_t accumulators[1024];
 
 void main() {
     const uint output_base = gl_WorkGroupID.x * CM_N * TILE_N;
-    const uint row_base = gl_WorkGroupID.y * CM_M * TILE_M;
+    const uint row_base    = gl_WorkGroupID.y * CM_M * TILE_M;
     coopmat<int32_t, gl_ScopeSubgroup, CM_M, CM_N, gl_MatrixUseAccumulator> sums[TILE_M * TILE_N];
     [[unroll]]
     for (uint i = 0; i < TILE_M * TILE_N; ++i) {
@@ -64,7 +76,7 @@ void main() {
             [[unroll]]
             for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) {
                 const uint index = tile_m * TILE_N + tile_n;
-                sums[index] = coopMatMulAdd(activations[tile_m], weights[tile_n], sums[index]);
+                sums[index]      = coopMatMulAdd(activations[tile_m], weights[tile_n], sums[index]);
             }
         }
     }
@@ -79,9 +91,9 @@ void main() {
                            gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
 
             for (uint index = gl_SubgroupInvocationID; index < CM_M * CM_N; index += gl_SubgroupSize) {
-                const uint row = row_base + tile_m * CM_M + index / CM_N;
+                const uint row          = row_base + tile_m * CM_M + index / CM_N;
                 const uint output_index = output_base + tile_n * CM_N + index % CM_N;
-                float value = float(accumulators[index]) * data_as[row] * data_ws[output_index];
+                float value             = float(accumulators[index]) * data_as[row] * data_ws[output_index];
                 if (p.has_bias != 0) {
                     value += data_b[output_index];
                 }

src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp

@@ -7,15 +7,14 @@
 #extension GL_NV_cooperative_matrix2 : require
 #extension GL_KHR_memory_scope_semantics : require
 
-layout(constant_id = 0) const uint BN = 128;
+layout(constant_id = 0) const uint BN        = 128;
 layout(constant_id = 1) const uint ALIGNMENT = 0;
-layout(constant_id = 2) const bool SWIGLU = false;
+layout(constant_id = 2) const bool SWIGLU    = false;
 layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;
 layout(constant_id = 3) const uint BM = 128;
 layout(constant_id = 4) const uint BK = 64;
 
-layout(push_constant) uniform parameter
-{
+layout(push_constant) uniform parameter {
     uint n;
     uint rows;
     uint k;
@@ -23,25 +22,45 @@ layout(push_constant) uniform parameter
     uint row_offset;
     float output_scale;
     uint group_cols;
-} p;
+}
+p;
 
-layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; };
-layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; };
-layout(binding = 2, std430) readonly buffer WS { float data_ws[]; };
-layout(binding = 3, std430) readonly buffer AS { float data_as[]; };
-layout(binding = 4, std430) readonly buffer B { float data_b[]; };
+layout(binding = 0, std430) readonly buffer W {
+    int8_t data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int8_t data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
 // Store FP32 bit patterns without changing the integer accumulator fragment layout.
-layout(binding = 5, std430) writeonly buffer D { int32_t data_d[]; };
+layout(binding = 5, std430) writeonly buffer D {
+    int32_t data_d[];
+};
 #ifdef SPLIT_SWIGLU
-layout(binding = 6, std430) readonly buffer UW { int8_t data_uw[]; };
-layout(binding = 7, std430) readonly buffer US { float data_us[]; };
+layout(binding = 6, std430) readonly buffer UW {
+    int8_t data_uw[];
+};
+layout(binding = 7, std430) readonly buffer US {
+    float data_us[];
+};
 #endif
 #ifdef RESIDUAL
-layout(binding = 6, std430) readonly buffer G { float data_g[]; };
-layout(binding = 7, std430) readonly buffer R { float data_r[]; };
+layout(binding = 6, std430) readonly buffer G {
+    float data_g[];
+};
+layout(binding = 7, std430) readonly buffer R {
+    float data_r[];
+};
 #endif
-int32_t dequantize(const in uint row, const in uint col, const in int32_t acc,
-                const in uint row_base, const in uint col_base) {
+int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) {
     const uint r = row_base + col;
     const uint c = col_base + row;
     if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) {
@@ -75,32 +94,32 @@ int32_t dequantize_swiglu(const in uint row, const in uint col, const in int32_t
 }
 
 void main() {
-    const uint k_size = p.k & ~255u;
+    const uint k_size                                                   = p.k & ~255u;
     tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> a_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
     tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> w_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
     tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
-    a_layout = setTensorLayoutDimensionNV(a_layout, p.rows, k_size);
-    w_layout = setTensorLayoutDimensionNV(w_layout, p.n, k_size);
-    d_layout = setTensorLayoutDimensionNV(d_layout, p.rows, SWIGLU ? p.n / 2 : p.n);
-    tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0);
+    a_layout                                                            = setTensorLayoutDimensionNV(a_layout, p.rows, k_size);
+    w_layout                                                            = setTensorLayoutDimensionNV(w_layout, p.n, k_size);
+    d_layout                                                            = setTensorLayoutDimensionNV(d_layout, p.rows, SWIGLU ? p.n / 2 : p.n);
+    tensorViewNV<2, false, 1, 0> transpose                              = createTensorViewNV(2, false, 1, 0);
     // Keep a weight panel in L2 while visiting the activation tiles.
-    const uint col_tiles = gl_NumWorkGroups.x;
-    const uint row_tiles = gl_NumWorkGroups.y;
-    const uint group_cols = clamp(p.group_cols, 1u, col_tiles);
-    const uint linear = gl_WorkGroupID.y * col_tiles + gl_WorkGroupID.x;
-    const uint group = linear / (group_cols * row_tiles);
-    const uint first_col = group * group_cols;
-    const uint cols = min(col_tiles - first_col, group_cols);
-    const uint within_group = linear - group * group_cols * row_tiles;
-    const uint row = p.row_offset + (within_group / cols) * BN;
-    const uint col = (first_col + within_group % cols) * BM;
-    coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+    const uint col_tiles                                                        = gl_NumWorkGroups.x;
+    const uint row_tiles                                                        = gl_NumWorkGroups.y;
+    const uint group_cols                                                       = clamp(p.group_cols, 1u, col_tiles);
+    const uint linear                                                           = gl_WorkGroupID.y * col_tiles + gl_WorkGroupID.x;
+    const uint group                                                            = linear / (group_cols * row_tiles);
+    const uint first_col                                                        = group * group_cols;
+    const uint cols                                                             = min(col_tiles - first_col, group_cols);
+    const uint within_group                                                     = linear - group * group_cols * row_tiles;
+    const uint row                                                              = p.row_offset + (within_group / cols) * BN;
+    const uint col                                                              = (first_col + within_group % cols) * BM;
+    coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum    = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
     coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum_up = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
     if (ALIGNMENT != 0) {
         tensorLayoutNV<2> a_full = createTensorLayoutNV(2);
         tensorLayoutNV<2> w_full = createTensorLayoutNV(2);
-        a_full = setTensorLayoutDimensionNV(a_full, p.rows, k_size);
-        w_full = setTensorLayoutDimensionNV(w_full, p.n, k_size);
+        a_full                   = setTensorLayoutDimensionNV(a_full, p.rows, k_size);
+        w_full                   = setTensorLayoutDimensionNV(w_full, p.n, k_size);
         for (uint k = 0; k < p.k; k += BK) {
             coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
             coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
@@ -145,7 +164,7 @@ void main() {
     // The aligned dispatch contains only complete tiles and a 128-aligned output width.
     if (ALIGNMENT == 1) {
         tensorLayoutNV<2> d_full = createTensorLayoutNV(2);
-        d_full = setTensorLayoutDimensionNV(d_full, p.rows, (SWIGLU ? p.n / 2 : p.n) & ~127u);
+        d_full                   = setTensorLayoutDimensionNV(d_full, p.rows, (SWIGLU ? p.n / 2 : p.n) & ~127u);
         coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_full, row, BN, col, BM), transpose);
     } else {
         coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row, BN, col, BM), transpose);

src/core/vulkan/shaders/quantize_i8_convrot.comp

@@ -4,45 +4,51 @@
 
 layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
 
-layout(push_constant) uniform parameter
-{
+layout(push_constant) uniform parameter {
     uint k;
     uint rows;
-} p;
+}
+p;
 
-layout(binding = 0, std430) readonly buffer A { float data_a[]; };
-layout(binding = 1, std430) writeonly buffer D { uint data_d[]; };
-layout(binding = 2, std430) writeonly buffer S { float data_s[]; };
+layout(binding = 0, std430) readonly buffer A {
+    float data_a[];
+};
+layout(binding = 1, std430) writeonly buffer D {
+    uint data_d[];
+};
+layout(binding = 2, std430) writeonly buffer S {
+    float data_s[];
+};
 
 shared float values[256];
 shared float maxima[64];
 
 void transform_group(const uint row, const uint group) {
-    const uint tid = gl_LocalInvocationID.x;
+    const uint tid         = gl_LocalInvocationID.x;
     const uint base_offset = row * p.k + group * 256;
 
     [[unroll]]
     for (uint i = 0; i < 4; ++i) {
         const uint index = tid + i * 64;
-        values[index] = data_a[base_offset + index] * (1.0 / 16.0);
+        values[index]    = data_a[base_offset + index] * (1.0 / 16.0);
     }
     barrier();
 
     [[unroll]]
     for (uint stride = 1; stride < 256; stride *= 4) {
         const uint base = (tid / stride) * 4 * stride + tid % stride;
-        const uint i0 = base;
-        const uint i1 = i0 + stride;
-        const uint i2 = i1 + stride;
-        const uint i3 = i2 + stride;
-        const float a = values[i0];
-        const float b = values[i1];
-        const float c = values[i2];
-        const float d = values[i3];
-        values[i0] =  a + b + c - d;
-        values[i1] =  a + b - c + d;
-        values[i2] =  a - b + c + d;
-        values[i3] = -a + b + c + d;
+        const uint i0   = base;
+        const uint i1   = i0 + stride;
+        const uint i2   = i1 + stride;
+        const uint i3   = i2 + stride;
+        const float a   = values[i0];
+        const float b   = values[i1];
+        const float c   = values[i2];
+        const float d   = values[i3];
+        values[i0]      = a + b + c - d;
+        values[i1]      = a + b - c + d;
+        values[i2]      = a - b + c + d;
+        values[i3]      = -a + b + c + d;
         barrier();
     }
 }
@@ -53,9 +59,9 @@ uint quantize_value(const float value, const float inv_scale) {
 }
 
 void quantize_row(const uint row) {
-    const uint tid = gl_LocalInvocationID.x;
+    const uint tid    = gl_LocalInvocationID.x;
     const uint groups = p.k / 256;
-    float local_max = 0.0;
+    float local_max   = 0.0;
 
     for (uint group = 0; group < groups; ++group) {
         transform_group(row, group);
@@ -85,11 +91,8 @@ void quantize_row(const uint row) {
     const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
     for (uint group = 0; group < groups; ++group) {
         transform_group(row, group);
-        const uint index = tid * 4;
-        const uint packed = quantize_value(values[index], inv_scale)
-                          | (quantize_value(values[index + 1], inv_scale) << 8)
-                          | (quantize_value(values[index + 2], inv_scale) << 16)
-                          | (quantize_value(values[index + 3], inv_scale) << 24);
+        const uint index                            = tid * 4;
+        const uint packed                           = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24);
         data_d[(row * p.k + group * 256) / 4 + tid] = packed;
         barrier();
     }

src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp

@@ -8,12 +8,24 @@
 #extension GL_EXT_control_flow_attributes : require
 
 layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
-layout(push_constant) uniform parameter { uint k; uint rows; } p;
-layout(binding = 0, std430) readonly buffer A { float data_a[]; };
-layout(binding = 1, std430) writeonly buffer D { int8_t data_d[]; };
-layout(binding = 2, std430) writeonly buffer S { float data_s[]; };
+layout(push_constant) uniform parameter {
+    uint k;
+    uint rows;
+}
+p;
+layout(binding = 0, std430) readonly buffer A {
+    float data_a[];
+};
+layout(binding = 1, std430) writeonly buffer D {
+    int8_t data_d[];
+};
+layout(binding = 2, std430) writeonly buffer S {
+    float data_s[];
+};
 #if defined(SWIGLU) || defined(SIGMOID_MUL)
-layout(binding = 3, std430) readonly buffer B { float data_b[]; };
+layout(binding = 3, std430) readonly buffer B {
+    float data_b[];
+};
 #endif
 shared float maxima[8];
 #ifdef SINGLE_PASS
@@ -26,7 +38,7 @@ layout(constant_id = 0) const uint K = 4096;
 
 vec4 rotate(float a, float b, float c, float d) {
     precise vec4 r = vec4(a + b + c - d, a + b - c + d,
-                         a - b + c + d, -a + b + c + d);
+                          a - b + c + d, -a + b + c + d);
     return r;
 }
 
@@ -37,12 +49,12 @@ void transform_group(uint row, uint group, out float v[8]) {
         float value = data_a[base + lane + j * 32];
 #ifdef SWIGLU
         precise float activated = value / (1.0 + exp(-value)) * data_b[base + lane + j * 32];
-        value = activated;
+        value                   = activated;
 #endif
 #ifdef SIGMOID_MUL
-        precise float gate = 1.0 / (1.0 + exp(-value));
+        precise float gate      = 1.0 / (1.0 + exp(-value));
         precise float activated = data_b[base + lane + j * 32] * gate;
-        value = activated;
+        value                   = activated;
 #endif
         v[j] = value * (1.0 / 16.0);
     }
@@ -51,18 +63,21 @@ void transform_group(uint row, uint group, out float v[8]) {
         [[unroll]] for (uint j = 0; j < 8; ++j) {
             vec4 r = rotate(subgroupShuffle(v[j], first), subgroupShuffle(v[j], first + stride),
                             subgroupShuffle(v[j], first + 2 * stride), subgroupShuffle(v[j], first + 3 * stride));
-            v[j] = r[(lane / stride) & 3u];
+            v[j]   = r[(lane / stride) & 3u];
         }
     }
     [[unroll]] for (uint j = 0; j < 8; j += 2) {
-        vec4 r = rotate(subgroupShuffle(v[j], lane & 15u), subgroupShuffle(v[j], (lane & 15u) + 16),
-                        subgroupShuffle(v[j + 1], lane & 15u), subgroupShuffle(v[j + 1], (lane & 15u) + 16));
-        v[j] = r[lane / 16];
+        vec4 r   = rotate(subgroupShuffle(v[j], lane & 15u), subgroupShuffle(v[j], (lane & 15u) + 16),
+                          subgroupShuffle(v[j + 1], lane & 15u), subgroupShuffle(v[j + 1], (lane & 15u) + 16));
+        v[j]     = r[lane / 16];
         v[j + 1] = r[2 + lane / 16];
     }
     [[unroll]] for (uint j = 0; j < 2; ++j) {
-        vec4 r = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]);
-        v[j] = r.x; v[j + 2] = r.y; v[j + 4] = r.z; v[j + 6] = r.w;
+        vec4 r   = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]);
+        v[j]     = r.x;
+        v[j + 2] = r.y;
+        v[j + 4] = r.z;
+        v[j + 6] = r.w;
     }
 }
 
@@ -72,26 +87,32 @@ void main() {
 #ifdef SINGLE_PASS
         float values[((K + 2047) / 2048) * 8];
         [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
-            if (gl_SubgroupID + g * 8 >= K / 256) continue;
-            float v[8]; transform_group(row, gl_SubgroupID + g * 8, v);
+            if (gl_SubgroupID + g * 8 >= K / 256)
+                continue;
+            float v[8];
+            transform_group(row, gl_SubgroupID + g * 8, v);
             [[unroll]] for (uint j = 0; j < 8; ++j) {
                 values[g * 8 + j] = v[j];
-                local_max = max(local_max, abs(v[j]));
+                local_max         = max(local_max, abs(v[j]));
             }
         }
 #else
         for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
-            float v[8]; transform_group(row, group, v);
-            [[unroll]] for (uint j = 0; j < 8; ++j) local_max = max(local_max, abs(v[j]));
+            float v[8];
+            transform_group(row, group, v);
+            [[unroll]] for (uint j = 0; j < 8; ++j)
+                local_max = max(local_max, abs(v[j]));
         }
 #endif
         local_max = subgroupMax(local_max);
-        if (gl_SubgroupInvocationID == 0) maxima[gl_SubgroupID] = local_max;
+        if (gl_SubgroupInvocationID == 0)
+            maxima[gl_SubgroupID] = local_max;
         barrier();
         if (gl_LocalInvocationID.x == 0) {
             float row_max = 0.0;
-            for (uint j = 0; j < 8; ++j) row_max = max(row_max, maxima[j]);
-            maxima[0] = row_max;
+            for (uint j = 0; j < 8; ++j)
+                row_max = max(row_max, maxima[j]);
+            maxima[0]   = row_max;
             data_s[row] = row_max / 127.0;
         }
         barrier();
@@ -99,14 +120,16 @@ void main() {
         const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
 #ifdef SINGLE_PASS
         [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
-            if (gl_SubgroupID + g * 8 >= K / 256) continue;
+            if (gl_SubgroupID + g * 8 >= K / 256)
+                continue;
             [[unroll]] for (uint j = 0; j < 8; ++j) {
                 data_d[row * p.k + (gl_SubgroupID + g * 8) * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(values[g * 8 + j] * inv_scale)), -127, 127));
             }
         }
 #else
         for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
-            float v[8]; transform_group(row, group, v);
+            float v[8];
+            transform_group(row, group, v);
             [[unroll]] for (uint j = 0; j < 8; ++j) {
                 data_d[row * p.k + group * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(v[j] * inv_scale)), -127, 127));
             }

src/core/vulkan/shaders/rms_rope_apply.comp

@@ -5,16 +5,26 @@
 #extension GL_EXT_control_flow_attributes : require
 
 layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;
-layout(binding = 0) readonly buffer X { float data_x[]; };
-layout(binding = 1) readonly buffer PE { float data_pe[]; };
+layout(binding = 0) readonly buffer X {
+    float data_x[];
+};
+layout(binding = 1) readonly buffer PE {
+    float data_pe[];
+};
 #ifdef CAST_F16
 #extension GL_EXT_shader_16bit_storage : require
 #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
-layout(binding = 2) writeonly buffer D { float16_t data_d[]; };
+layout(binding = 2) writeonly buffer D {
+    float16_t data_d[];
+};
 #else
-layout(binding = 2) writeonly buffer D { float data_d[]; };
+layout(binding = 2) writeonly buffer D {
+    float data_d[];
+};
 #endif
-layout(binding = 3) readonly buffer W { float data_w[]; };
+layout(binding = 3) readonly buffer W {
+    float data_w[];
+};
 
 layout(push_constant) uniform parameter {
     uint half_dim;
@@ -31,10 +41,11 @@ layout(push_constant) uniform parameter {
     float eps;
     uint w_offset;
     float output_scale;
-} pc;
+}
+pc;
 
 float shuffle_channel(vec4 values, uint channel) {
-    const uint lane = channel & 31u;
+    const uint lane     = channel & 31u;
     const vec4 shuffled = vec4(subgroupShuffle(values.x, lane), subgroupShuffle(values.y, lane),
                                subgroupShuffle(values.z, lane), subgroupShuffle(values.w, lane));
     return shuffled[channel / 32];
@@ -45,43 +56,43 @@ void main() {
     if (head >= pc.heads) {
         return;
     }
-    const uint lane = gl_SubgroupInvocationID;
+    const uint lane    = gl_SubgroupInvocationID;
     const uint batches = pc.pairs / (64 * pc.heads * pc.tokens);
     for (uint batch = gl_WorkGroupID.z; batch < batches; batch += gl_NumWorkGroups.z) {
         for (uint token = gl_WorkGroupID.y; token < pc.tokens; token += gl_NumWorkGroups.y) {
-            const uint base = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3];
-            vec4 x = vec4(data_x[base + lane], data_x[base + lane + 32],
-                          data_x[base + lane + 64], data_x[base + lane + 96]);
+            const uint base      = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3];
+            vec4 x               = vec4(data_x[base + lane], data_x[base + lane + 32],
+                                        data_x[base + lane + 64], data_x[base + lane + 96]);
             precise vec4 squares = x * x;
-            precise float sum = (squares.x + squares.z) + (squares.y + squares.w);
+            precise float sum    = (squares.x + squares.z) + (squares.y + squares.w);
             [[unroll]] for (uint stride = 16; stride > 0; stride >>= 1) {
                 sum += subgroupShuffle(sum, lane ^ stride);
             }
-            sum = subgroupShuffle(sum, 0);
-            const float scale = inversesqrt(sum / 128.0 + pc.eps);
+            sum                   = subgroupShuffle(sum, 0);
+            const float scale     = inversesqrt(sum / 128.0 + pc.eps);
             precise vec4 weighted = scale * x;
             [[unroll]] for (uint i = 0; i < 4; ++i) {
                 weighted[i] *= data_w[pc.w_offset + lane + i * 32];
             }
             [[unroll]] for (uint i = 0; i < 4; ++i) {
-                const uint col = lane + i * 32;
+                const uint col      = lane + i * 32;
                 precise float value = weighted[i];
                 if (pc.interleaved != 0) {
                     const float a = subgroupShuffle(value, lane & ~1u);
                     const float b = subgroupShuffle(value, lane | 1u);
                     if (col / 2 < pc.rot_half) {
                         const uint matrix = pc.pe_offset + (col / 2) * pc.pe_stride[2] + token * pc.pe_stride[3] + (col & 1u) * pc.pe_stride[1];
-                        value = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]];
+                        value             = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]];
                     }
                 } else {
                     const uint pair = col % pc.rot_half;
-                    const float a = shuffle_channel(weighted, pair);
-                    const float b = shuffle_channel(weighted, pair + pc.rot_half);
+                    const float a   = shuffle_channel(weighted, pair);
+                    const float b   = shuffle_channel(weighted, pair + pc.rot_half);
                     if (col < 2 * pc.rot_half) {
                         const uint matrix = pc.pe_offset + pair * pc.pe_stride[2] + token * pc.pe_stride[3];
-                        const float c = data_pe[matrix];
-                        const float s = data_pe[matrix + pc.pe_stride[1]];
-                        value = col < pc.rot_half ? a * c - b * s : a * s + b * c;
+                        const float c     = data_pe[matrix];
+                        const float s     = data_pe[matrix + pc.pe_stride[1]];
+                        value             = col < pc.rot_half ? a * c - b * s : a * s + b * c;
                     }
                 }
 #ifdef CAST_F16

src/core/vulkan/shaders/rope_apply.comp

@@ -2,9 +2,15 @@
 
 layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
 
-layout(binding = 0) readonly buffer X { float data_x[]; };
-layout(binding = 1) readonly buffer PE { float data_pe[]; };
-layout(binding = 2) writeonly buffer D { float data_d[]; };
+layout(binding = 0) readonly buffer X {
+    float data_x[];
+};
+layout(binding = 1) readonly buffer PE {
+    float data_pe[];
+};
+layout(binding = 2) writeonly buffer D {
+    float data_d[];
+};
 
 layout(push_constant) uniform parameter {
     uint half_dim;
@@ -18,41 +24,42 @@ layout(push_constant) uniform parameter {
     uint pe_offset;
     uint dst_offset;
     uint interleaved;
-} pc;
+}
+pc;
 
 void main() {
-    const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x*gl_WorkGroupSize.x*gl_GlobalInvocationID.y;
+    const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x * gl_WorkGroupSize.x * gl_GlobalInvocationID.y;
     if (index >= pc.pairs) {
         return;
     }
-    const uint j     = index % pc.half_dim;
-    const uint row   = index / pc.half_dim;
-    const uint token = row % pc.tokens;
-    const uint head  = (row / pc.tokens) % pc.heads;
-    const uint batch = row / (pc.tokens*pc.heads);
-    const uint i0    = pc.interleaved != 0 || j >= pc.rot_half ? 2*j : j;
-    const uint i1    = pc.interleaved != 0 || j >= pc.rot_half ? 2*j + 1 : j + pc.rot_half;
-    const uint x_row = pc.x_offset + head*pc.x_stride[1] + token*pc.x_stride[2] + batch*pc.x_stride[3];
-    const uint out_row = pc.dst_offset + row*pc.half_dim*2;
-    const float a = data_x[x_row + i0*pc.x_stride[0]];
-    const float b = data_x[x_row + i1*pc.x_stride[0]];
+    const uint j       = index % pc.half_dim;
+    const uint row     = index / pc.half_dim;
+    const uint token   = row % pc.tokens;
+    const uint head    = (row / pc.tokens) % pc.heads;
+    const uint batch   = row / (pc.tokens * pc.heads);
+    const uint i0      = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j : j;
+    const uint i1      = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j + 1 : j + pc.rot_half;
+    const uint x_row   = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3];
+    const uint out_row = pc.dst_offset + row * pc.half_dim * 2;
+    const float a      = data_x[x_row + i0 * pc.x_stride[0]];
+    const float b      = data_x[x_row + i1 * pc.x_stride[0]];
     if (j >= pc.rot_half) {
         data_d[out_row + i0] = a;
         data_d[out_row + i1] = b;
         return;
     }
-    const uint matrix = pc.pe_offset + j*pc.pe_stride[2] + token*pc.pe_stride[3];
-    const float c = data_pe[matrix];
-    const float s = data_pe[matrix + pc.pe_stride[1]];
+    const uint matrix = pc.pe_offset + j * pc.pe_stride[2] + token * pc.pe_stride[3];
+    const float c     = data_pe[matrix];
+    const float s     = data_pe[matrix + pc.pe_stride[1]];
     // Preserve the separate multiply/add rounding of the unfused graph.
     precise float y0;
     precise float y1;
     if (pc.interleaved != 0) {
-        y0 = a*c + b*data_pe[matrix + pc.pe_stride[0]];
-        y1 = a*s + b*data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]];
+        y0 = a * c + b * data_pe[matrix + pc.pe_stride[0]];
+        y1 = a * s + b * data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]];
     } else {
-        y0 = a*c - b*s;
-        y1 = a*s + b*c;
+        y0 = a * c - b * s;
+        y1 = a * s + b * c;
     }
     data_d[out_row + i0] = y0;
     data_d[out_row + i1] = y1;