Skip to content

[ARK] Optimize SYCL-TLA woqgemm_s8 #2133

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Zhenzhong1 wants to merge 2 commits into
base: main
Choose a base branch
from
+275 −12
Draft

[ARK] Optimize SYCL-TLA woqgemm_s8 #2133

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ebf4a78
inference pass v0.1
Zhenzhong1 Aug 7, 2026
9f7bbca
inference pass v1.0
Zhenzhong1 Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
287 changes: 275 additions & 12 deletions auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,201 @@ class S8FinalizeKernelName;
template <int TileM, int TileN, class SGLayout>
class S8AccumKernelName;

template <bool HasBias, bool FullTile, class ElementOut, int TileM, int TileN, class SGLayout>
class S8KBlockDequantKernelName;

template <bool HasBias, bool FullTile, class TiledMMA, class ElementOut>
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<float>(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;

int k_tile_size = int(get<2>(wg_tile));

if (blocksize % k_tile_size == 0) {
auto A = make_tensor(make_gmem_ptr(const_cast<int8_t*>(a)), make_shape(m, k), make_stride(k, _1{}));
auto B = make_tensor(make_gmem_ptr(const_cast<int8_t*>(b)), make_shape(n, k), 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_tiles_per_block = blocksize / k_tile_size;
int k_tile_count = blks * k_tiles_per_block;
int k_tile_prefetch = 0;

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 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);

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);
}

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<float>(scale_b[col * blks + ib]);
tFrC(i) += static_cast<float>(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);

auto A = make_tensor(make_gmem_ptr(const_cast<int8_t*>(a_blk)), make_shape(m, blocksize), make_stride(k, _1{}));
auto B = make_tensor(make_gmem_ptr(const_cast<int8_t*>(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);

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);
}

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<float>(scale_b[col * blks + ib]);
tFrC(i) += static_cast<float>(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<float>(scale_a[row]);
if constexpr (HasBias) {
value += static_cast<float>(bias[col]);
}
c[row * n + col] = static_cast<ElementOut>(value);
}
}

template <bool AccumBlock, bool HasBias, bool FullTile, class ATensor, class BTensor, class TiledMMA, class ElementOut>
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,
Expand Down Expand Up @@ -143,6 +338,39 @@ void igemm_device_impl(ATensor const& A, BTensor const& B, TiledMMA const& mma,
}
}

template <bool HasBias, class ElementOut, int TileM, int TileN, class SGLayout>
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<TileM>, Int<TileN>, _64>;
using MMA = typename TiledMMAHelper<MMA_Atom<Op>, Layout<WGTile>, 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<S8KBlockDequantKernelName<HasBias, true, ElementOut, TileM, TileN, SGLayout>>(
sycl::nd_range<2>(global, local), props, [=](auto) {
igemm_kblock_device_impl<HasBias, true>(mma, a, b, c, scale_a, scale_b, bias, m, n, k, blocksize, blks);
});
} else {
q->parallel_for<S8KBlockDequantKernelName<HasBias, false, ElementOut, TileM, TileN, SGLayout>>(
sycl::nd_range<2>(global, local), props, [=](auto) {
igemm_kblock_device_impl<HasBias, false>(mma, a, b, c, scale_a, scale_b, bias, m, n, k, blocksize, blks);
});
}
}

template <bool AccumBlock, class ElementOut, int TileM, int TileN, class SGLayout>
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,
Expand Down Expand Up @@ -200,6 +428,52 @@ void launch_igemm_tile(sycl::queue* q, int m, int n, int gemm_k, int lda, int ld
}
}

template <class ElementOut>
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<Shape<_1, _4, _1>, Stride<_0, _1, _0>>;
using SmallMidTileSG = Layout<Shape<_2, _4, _1>, Stride<_4, _1, _0>>;
using MediumTileSG = Layout<Shape<_4, _4, _1>, Stride<_4, _1, _0>>;
using LargeTileSG = Layout<Shape<_8, _4, _1>, Stride<_4, _1, _0>>;

bool has_bias = bias != nullptr;

if (m < 16) {
if (has_bias) {
launch_igemm_kblock_tile<true, ElementOut, 8, 128, SmallTileSG>(
q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks);
} else {
launch_igemm_kblock_tile<false, ElementOut, 8, 128, SmallTileSG>(
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<true, ElementOut, 64, 128, SmallMidTileSG>(
q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks);
} else {
launch_igemm_kblock_tile<false, ElementOut, 64, 128, SmallMidTileSG>(
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<true, ElementOut, 128, 128, MediumTileSG>(
q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks);
} else {
launch_igemm_kblock_tile<false, ElementOut, 128, 128, MediumTileSG>(
q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks);
}
} else {
if (has_bias) {
launch_igemm_kblock_tile<true, ElementOut, 256, 128, LargeTileSG>(
q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks);
} else {
launch_igemm_kblock_tile<false, ElementOut, 256, 128, LargeTileSG>(
q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks);
}
}
}


template <bool AccumBlock, class ElementOut>
void launch_igemm(sycl::queue* q, int m, int n, int gemm_k, int lda, int ldb,
Expand Down Expand Up @@ -254,18 +528,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<float*>(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<true, ElementOut>(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<ElementOut>(q, m, n, k, a, b, c, scale_a, scale_b, bias, blocksize, blks);
}

} // namespace sycl_tla_s8_detail
Expand Down