From ebf4a7896402ea9df4a28d021ceb846342963a76 Mon Sep 17 00:00:00 2001 From: Zhenzhong1 Date: Fri, 7 Aug 2026 14:50:17 +0800 Subject: [PATCH 1/2] inference pass v0.1 Signed-off-by: Zhenzhong1 --- .../wrapper/include/sycl_tla_s8_gemm.hpp | 222 +++++++++++++++++- 1 file changed, 210 insertions(+), 12 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp index 51fe8fac5..5a0e84883 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp @@ -40,6 +40,136 @@ class S8FinalizeKernelName; template class S8AccumKernelName; +template +class S8KBlockDequantKernelName; +template +void igemm_kblock_device_impl(TiledMMA const& mma, const int8_t* a, const int8_t* b, ElementOut* c, + const ElementOut* scale_a, const ElementOut* scale_b, const ElementOut* bias, int m, + int n, int k, int blocksize, int blks) { + auto item = sycl::ext::oneapi::this_work_item::get_nd_item<2>(); + int wg_m = int(item.get_group(1)); + int wg_n = int(item.get_group(0)); + int local_id = int(item.get_local_id(0)); + + auto wg_tile = mma.tile_mnk(); + auto wg_coord = make_coord(wg_m, wg_n, 0); + + Tensor cC = make_identity_tensor(make_shape(m, n)); + Tensor gC = local_tile(cC, wg_tile, wg_coord, Step<_1, _1, X>{}); + + auto thr_mma = mma.get_slice(local_id); + + Tensor tCrC = partition_fragment_C(mma, select<0, 1>(wg_tile)); + Tensor tFrC = make_tensor_like(tCrC); + Tensor tCgC = thr_mma.partition_C(gC); + + CUTE_UNROLL + for (int i = 0; i < size(tFrC); ++i) { + tFrC(i) = 0.0f; + } + + constexpr SPIRVScope barrier_scope = ScopeWorkgroup; + constexpr int prefetch_dist = 3; + + for (int ib = 0; ib < blks; ++ib) { + const int8_t* a_blk = a + size_t(ib) * size_t(blocksize); + const int8_t* b_blk = b + size_t(ib) * size_t(blocksize); + + auto A = make_tensor(make_gmem_ptr(const_cast(a_blk)), make_shape(m, blocksize), make_stride(k, _1{})); + auto B = make_tensor(make_gmem_ptr(const_cast(b_blk)), make_shape(n, blocksize), make_stride(k, _1{})); + + Tensor cA = make_identity_tensor(A.shape()); + Tensor cB = make_identity_tensor(B.shape()); + + Tensor gA = local_tile(cA, select<0, 2>(wg_tile), make_coord(wg_m, _)); + Tensor gB = local_tile(cB, select<1, 2>(wg_tile), make_coord(wg_n, _)); + + auto copy_a = make_block_2d_copy_A(mma, A); + auto copy_b = make_block_2d_copy_B(mma, B); + + auto thr_copy_a = copy_a.get_slice(local_id); + auto thr_copy_b = copy_b.get_slice(local_id); + + auto tCrA = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); + auto tCrB = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); + + auto tArA = thr_copy_a.partition_sg_fragment_D(gA(_, _, 0)); + auto tBrB = thr_copy_b.partition_sg_fragment_D(gB(_, _, 0)); + + Tensor tAgA = thr_copy_a.partition_S(gA); + Tensor tBgB = thr_copy_b.partition_S(gB); + + auto prefetch_a = make_block_2d_prefetch(copy_a); + auto prefetch_b = make_block_2d_prefetch(copy_b); + + auto thr_prefetch_A = prefetch_a.get_slice(local_id); + auto thr_prefetch_B = prefetch_b.get_slice(local_id); + + auto pAgA = thr_prefetch_A.partition_S(gA); + auto pBgB = thr_prefetch_B.partition_S(gB); + + int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); + int k_tile_prefetch = 0; + + clear(tCrC); + + CUTE_UNROLL + for (; k_tile_prefetch < prefetch_dist && k_tile_prefetch < k_tile_count; ++k_tile_prefetch) { + prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); + prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + } + + for (int k_tile = 0; k_tile < k_tile_count; ++k_tile, ++k_tile_prefetch) { + barrier_arrive(barrier_scope); + + copy(copy_a, tAgA(_, _, _, k_tile), tArA); + copy(copy_b, tBgB(_, _, _, k_tile), tBrB); + + if (k_tile_prefetch < k_tile_count) { + prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); + prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + } + + reorder(tArA, tCrA); + reorder(tBrB, tCrB); + gemm(mma, tCrA, tCrB, tCrC); + + barrier_wait(barrier_scope); + } + + CUTE_UNROLL + for (int i = 0; i < size(tCrC); ++i) { + auto coord = tCgC(i); + int row = int(get<0>(coord)); + int col = int(get<1>(coord)); + + if constexpr (!FullTile) { + if (row >= m || col >= n) continue; + } + + float sb = static_cast(scale_b[col * blks + ib]); + tFrC(i) += static_cast(tCrC(i)) * sb; + } + } + + CUTE_UNROLL + for (int i = 0; i < size(tFrC); ++i) { + auto coord = tCgC(i); + int row = int(get<0>(coord)); + int col = int(get<1>(coord)); + + if constexpr (!FullTile) { + if (row >= m || col >= n) continue; + } + + float value = tFrC(i) * static_cast(scale_a[row]); + if constexpr (HasBias) { + value += static_cast(bias[col]); + } + c[row * n + col] = static_cast(value); + } +} + template void igemm_device_impl(ATensor const& A, BTensor const& B, TiledMMA const& mma, ElementOut* c, float* accum, const ElementOut* scale_a, const ElementOut* scale_b, const ElementOut* bias, int m, int n, @@ -143,6 +273,39 @@ void igemm_device_impl(ATensor const& A, BTensor const& B, TiledMMA const& mma, } } +template +void launch_igemm_kblock_tile(sycl::queue* q, int m, int n, int k, const int8_t* a, const int8_t* b, ElementOut* c, + const ElementOut* scale_a, const ElementOut* scale_b, const ElementOut* bias, + int blocksize, int blks) { + compat::set_default_queue(*q); + + using Op = XE_DPAS_TT<8, int32_t, int8_t, int8_t>; + using WGTile = Shape, Int, _64>; + using MMA = typename TiledMMAHelper, Layout, SGLayout>::TiledMMA; + MMA mma{}; + + sycl::range<2> local = {size(mma), 1}; + sycl::range<2> global = {local[0] * ceil_div(n, get<1>(mma.tile_mnk())), + local[1] * ceil_div(m, get<0>(mma.tile_mnk()))}; + + namespace syclex = sycl::ext::oneapi::experimental; + namespace intelex = sycl::ext::intel::experimental; + syclex::properties props{syclex::sub_group_size<16>, intelex::grf_size<256>}; + + bool full_tile = (m % TileM == 0) && (n % TileN == 0); + + if (full_tile) { + q->parallel_for>( + sycl::nd_range<2>(global, local), props, [=](auto) { + igemm_kblock_device_impl(mma, a, b, c, scale_a, scale_b, bias, m, n, k, blocksize, blks); + }); + } else { + q->parallel_for>( + sycl::nd_range<2>(global, local), props, [=](auto) { + igemm_kblock_device_impl(mma, a, b, c, scale_a, scale_b, bias, m, n, k, blocksize, blks); + }); + } +} template void launch_igemm_tile(sycl::queue* q, int m, int n, int gemm_k, int lda, int ldb, const int8_t* a, const int8_t* b, @@ -200,6 +363,52 @@ void launch_igemm_tile(sycl::queue* q, int m, int n, int gemm_k, int lda, int ld } } +template +void launch_igemm_kblock(sycl::queue* q, int m, int n, int k, const int8_t* a, const int8_t* b, ElementOut* c, + const ElementOut* scale_a, const ElementOut* scale_b, const ElementOut* bias, int blocksize, + int blks) { + using SmallTileSG = Layout, Stride<_0, _1, _0>>; + using SmallMidTileSG = Layout, Stride<_4, _1, _0>>; + using MediumTileSG = Layout, Stride<_4, _1, _0>>; + using LargeTileSG = Layout, Stride<_4, _1, _0>>; + + bool has_bias = bias != nullptr; + + if (m < 16) { + if (has_bias) { + launch_igemm_kblock_tile( + q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); + } else { + launch_igemm_kblock_tile( + q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); + } + } else if (m < 128) { + if (has_bias) { + launch_igemm_kblock_tile( + q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); + } else { + launch_igemm_kblock_tile( + q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); + } + } else if (m <= 1024) { + if (has_bias) { + launch_igemm_kblock_tile( + q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); + } else { + launch_igemm_kblock_tile( + q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); + } + } else { + if (has_bias) { + launch_igemm_kblock_tile( + q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); + } else { + launch_igemm_kblock_tile( + q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); + } + } +} + template void launch_igemm(sycl::queue* q, int m, int n, int gemm_k, int lda, int ldb, @@ -254,18 +463,7 @@ void run_typed(sycl::queue* q, int m, int n, int k, const int8_t* a, const int8_ } int blks = k / blocksize; - size_t bytes = size_t(m) * size_t(n) * sizeof(float); - auto* accum = static_cast(DeviceMemoryPool::Instance()->get_scratch_mem(bytes, 3, q)); - q->memset(accum, 0, bytes); - - for (int ib = 0; ib < blks; ++ib) { - const int8_t* a_blk = a + size_t(ib) * size_t(blocksize); - const int8_t* b_blk = b + size_t(ib) * size_t(blocksize); - launch_igemm(q, m, n, blocksize, k, k, a_blk, b_blk, nullptr, accum, nullptr, scale_b, nullptr, - ib, blks); - } - - finalize_block_output(q, m, n, accum, c, scale_a, bias); + launch_igemm_kblock(q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks); } } // namespace sycl_tla_s8_detail From 9f7bbca2994f3595fc230cca9c4c48ad2c16ceb2 Mon Sep 17 00:00:00 2001 From: Zhenzhong1 Date: Fri, 7 Aug 2026 15:02:54 +0800 Subject: [PATCH 2/2] inference pass v1.0 Signed-off-by: Zhenzhong1 --- .../wrapper/include/sycl_tla_s8_gemm.hpp | 121 ++++++++++++++---- 1 file changed, 93 insertions(+), 28 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp index 5a0e84883..fbfb4faaf 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp @@ -42,6 +42,7 @@ class S8AccumKernelName; template class S8KBlockDequantKernelName; + template void igemm_kblock_device_impl(TiledMMA const& mma, const int8_t* a, const int8_t* b, ElementOut* c, const ElementOut* scale_a, const ElementOut* scale_b, const ElementOut* bias, int m, @@ -71,12 +72,11 @@ void igemm_kblock_device_impl(TiledMMA const& mma, const int8_t* a, const int8_t constexpr SPIRVScope barrier_scope = ScopeWorkgroup; constexpr int prefetch_dist = 3; - for (int ib = 0; ib < blks; ++ib) { - const int8_t* a_blk = a + size_t(ib) * size_t(blocksize); - const int8_t* b_blk = b + size_t(ib) * size_t(blocksize); + int k_tile_size = int(get<2>(wg_tile)); - auto A = make_tensor(make_gmem_ptr(const_cast(a_blk)), make_shape(m, blocksize), make_stride(k, _1{})); - auto B = make_tensor(make_gmem_ptr(const_cast(b_blk)), make_shape(n, blocksize), make_stride(k, _1{})); + if (blocksize % k_tile_size == 0) { + auto A = make_tensor(make_gmem_ptr(const_cast(a)), make_shape(m, k), make_stride(k, _1{})); + auto B = make_tensor(make_gmem_ptr(const_cast(b)), make_shape(n, k), make_stride(k, _1{})); Tensor cA = make_identity_tensor(A.shape()); Tensor cB = make_identity_tensor(B.shape()); @@ -108,47 +108,112 @@ void igemm_kblock_device_impl(TiledMMA const& mma, const int8_t* a, const int8_t auto pAgA = thr_prefetch_A.partition_S(gA); auto pBgB = thr_prefetch_B.partition_S(gB); - int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); + int k_tiles_per_block = blocksize / k_tile_size; + int k_tile_count = blks * k_tiles_per_block; int k_tile_prefetch = 0; - clear(tCrC); - CUTE_UNROLL for (; k_tile_prefetch < prefetch_dist && k_tile_prefetch < k_tile_count; ++k_tile_prefetch) { prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); } - for (int k_tile = 0; k_tile < k_tile_count; ++k_tile, ++k_tile_prefetch) { - barrier_arrive(barrier_scope); + for (int ib = 0; ib < blks; ++ib) { + clear(tCrC); + + for (int bk = 0; bk < k_tiles_per_block; ++bk) { + int k_tile = ib * k_tiles_per_block + bk; + + barrier_arrive(barrier_scope); - copy(copy_a, tAgA(_, _, _, k_tile), tArA); - copy(copy_b, tBgB(_, _, _, k_tile), tBrB); + copy(copy_a, tAgA(_, _, _, k_tile), tArA); + copy(copy_b, tBgB(_, _, _, k_tile), tBrB); - if (k_tile_prefetch < k_tile_count) { - prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); - prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + if (k_tile_prefetch < k_tile_count) { + prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); + prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + } + ++k_tile_prefetch; + + reorder(tArA, tCrA); + reorder(tBrB, tCrB); + gemm(mma, tCrA, tCrB, tCrC); + + barrier_wait(barrier_scope); } - reorder(tArA, tCrA); - reorder(tBrB, tCrB); - gemm(mma, tCrA, tCrB, tCrC); + CUTE_UNROLL + for (int i = 0; i < size(tCrC); ++i) { + auto coord = tCgC(i); + int row = int(get<0>(coord)); + int col = int(get<1>(coord)); - barrier_wait(barrier_scope); + if constexpr (!FullTile) { + if (row >= m || col >= n) continue; + } + + float sb = static_cast(scale_b[col * blks + ib]); + tFrC(i) += static_cast(tCrC(i)) * sb; + } } + } else { + for (int ib = 0; ib < blks; ++ib) { + const int8_t* a_blk = a + size_t(ib) * size_t(blocksize); + const int8_t* b_blk = b + size_t(ib) * size_t(blocksize); - CUTE_UNROLL - for (int i = 0; i < size(tCrC); ++i) { - auto coord = tCgC(i); - int row = int(get<0>(coord)); - int col = int(get<1>(coord)); + auto A = make_tensor(make_gmem_ptr(const_cast(a_blk)), make_shape(m, blocksize), make_stride(k, _1{})); + auto B = make_tensor(make_gmem_ptr(const_cast(b_blk)), make_shape(n, blocksize), make_stride(k, _1{})); - if constexpr (!FullTile) { - if (row >= m || col >= n) continue; + Tensor cA = make_identity_tensor(A.shape()); + Tensor cB = make_identity_tensor(B.shape()); + + Tensor gA = local_tile(cA, select<0, 2>(wg_tile), make_coord(wg_m, _)); + Tensor gB = local_tile(cB, select<1, 2>(wg_tile), make_coord(wg_n, _)); + + auto copy_a = make_block_2d_copy_A(mma, A); + auto copy_b = make_block_2d_copy_B(mma, B); + + auto thr_copy_a = copy_a.get_slice(local_id); + auto thr_copy_b = copy_b.get_slice(local_id); + + auto tCrA = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); + auto tCrB = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); + + auto tArA = thr_copy_a.partition_sg_fragment_D(gA(_, _, 0)); + auto tBrB = thr_copy_b.partition_sg_fragment_D(gB(_, _, 0)); + + Tensor tAgA = thr_copy_a.partition_S(gA); + Tensor tBgB = thr_copy_b.partition_S(gB); + + clear(tCrC); + + int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); + for (int k_tile = 0; k_tile < k_tile_count; ++k_tile) { + barrier_arrive(barrier_scope); + + copy(copy_a, tAgA(_, _, _, k_tile), tArA); + copy(copy_b, tBgB(_, _, _, k_tile), tBrB); + + reorder(tArA, tCrA); + reorder(tBrB, tCrB); + gemm(mma, tCrA, tCrB, tCrC); + + barrier_wait(barrier_scope); } - float sb = static_cast(scale_b[col * blks + ib]); - tFrC(i) += static_cast(tCrC(i)) * sb; + CUTE_UNROLL + for (int i = 0; i < size(tCrC); ++i) { + auto coord = tCgC(i); + int row = int(get<0>(coord)); + int col = int(get<1>(coord)); + + if constexpr (!FullTile) { + if (row >= m || col >= n) continue; + } + + float sb = static_cast(scale_b[col * blks + ib]); + tFrC(i) += static_cast(tCrC(i)) * sb; + } } }