diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index cc73da5c0..a98543c7f 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -600,6 +600,7 @@ def woqgemm( weight_type, scale_type, asym, + out: torch.Tensor | None = None, ): _validate_packed_blob(B, n, k, groupsize, compute_type, weight_type, scale_type, asym) m = A.shape[0] @@ -607,7 +608,18 @@ def woqgemm( ct = cvtstr_dtype(compute_type) wt = cvtstr_dtype(weight_type) st = cvtstr_dtype(scale_type) - C = torch.zeros(m, n, dtype=A.dtype, device=A.device) + if out is None: + C = torch.zeros(m, n, dtype=A.dtype, device=A.device) + else: + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if out.dtype != A.dtype: + raise ValueError(f"out dtype must be {A.dtype}, got {out.dtype}") + if out.device != A.device: + raise ValueError(f"out device must be {A.device}, got {out.device}") + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + C = out stream = get_stream(A) lib.woqgemm( stream, @@ -617,7 +629,7 @@ def woqgemm( A.contiguous().data_ptr(), cvt_dtype(A.dtype), B.contiguous().data_ptr(), - C.contiguous().data_ptr(), + C.data_ptr(), bias.contiguous().data_ptr(), groupsize, ct, diff --git a/auto_round_extension/ark/auto_round_kernel/sycl_tla_dense_woq_s4_dpas.cpp.in b/auto_round_extension/ark/auto_round_kernel/sycl_tla_dense_woq_s4_dpas.cpp.in new file mode 100644 index 000000000..4d24aecd4 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/sycl_tla_dense_woq_s4_dpas.cpp.in @@ -0,0 +1,33 @@ +// Generated by CMake. Do not edit directly. + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + +#include "sycl_tla_dense_woq_s4_dpas.hpp" + +namespace ark { +namespace dense_woq_s4_dpas { +namespace detail { + +void @DENSE_WOQ_S4_FUNCTION_NAME@(sycl::queue* q, const void* activations, const void* weights, const void* scales, + const void* bias, void* outputs, int gemm_m, int gemm_n, int gemm_k, + int group_size) { + using ElementA = cute_scalar_t; + + if constexpr (@DENSE_WOQ_S4_SCALE_GROUP_MAJOR@) { + DenseWoqS4GEMMLauncherUpTo( + *q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(bias), static_cast(outputs), + gemm_m, gemm_n, gemm_k, group_size); + } else { + DenseWoqS4GEMMLauncher<'R', 'C', @DENSE_WOQ_S4_POLICY_NAME@, false>( + *q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(bias), static_cast(outputs), + gemm_m, gemm_n, gemm_k, group_size); + } +} + +} // namespace detail +} // namespace dense_woq_s4_dpas +} // namespace ark + +#endif \ No newline at end of file diff --git a/auto_round_extension/ark/auto_round_kernel/sycl_tla_generation.cmake b/auto_round_extension/ark/auto_round_kernel/sycl_tla_generation.cmake index 9e0f6212c..ce2aa3967 100644 --- a/auto_round_extension/ark/auto_round_kernel/sycl_tla_generation.cmake +++ b/auto_round_extension/ark/auto_round_kernel/sycl_tla_generation.cmake @@ -127,6 +127,29 @@ foreach(_s8_dtype IN LISTS _s8_dtypes) endforeach() endforeach() +function(generate_dense_woq_s4_policy policy_suffix policy_name) + foreach(_dense_woq_s4_scale_layout IN ITEMS group_n n_group) + if(_dense_woq_s4_scale_layout STREQUAL "group_n") + set(DENSE_WOQ_S4_SCALE_GROUP_MAJOR true) + else() + set(DENSE_WOQ_S4_SCALE_GROUP_MAJOR false) + endif() + set(DENSE_WOQ_S4_POLICY_NAME ${policy_name}) + set(DENSE_WOQ_S4_FUNCTION_NAME run_${policy_suffix}_${_dense_woq_s4_scale_layout}) + generate_sycl_tla_source(sycl_tla_dense_woq_s4_dpas.cpp.in + sycl_tla_dense_woq_s4_dpas_${policy_suffix}_${_dense_woq_s4_scale_layout}.cpp) + endforeach() + set(SYCL_TLA_GENERATED_SRCS ${SYCL_TLA_GENERATED_SRCS} PARENT_SCOPE) +endfunction() + +generate_dense_woq_s4_policy(m4_n128 dpas_w4a16_dense_policy_m_4_n128) +generate_dense_woq_s4_policy(m8_n128 dpas_w4a16_dense_policy_m_8_n128) +generate_dense_woq_s4_policy(m16 dpas_w4a16_dense_policy_m_16) +generate_dense_woq_s4_policy(m32 dpas_w4a16_dense_policy_m_32) +generate_dense_woq_s4_policy(m32_n256 dpas_w4a16_dense_policy_m_32_n256) +generate_dense_woq_s4_policy(m64_n256 dpas_w4a16_dense_policy_m_64_n256) +generate_dense_woq_s4_policy(m128 dpas_w4a16_dense_policy_m_128) + set(MOE_SOURCE_MODE 8) set(MOE_DEFINE_FP8_HELPERS 0) set(MOE_DEFINE_INT8_HELPERS 0) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_woq_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_woq_s4_dpas.hpp new file mode 100644 index 000000000..42f506893 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_woq_s4_dpas.hpp @@ -0,0 +1,760 @@ +// Copyright (C) 2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 + +// SYCL-TLA Dense WOQ S4 DPAS Wrapper + +#pragma once + +#include +#include +#include + +#ifdef ARK_XPU +#include +#endif + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) +#include + +#include "cute/tensor.hpp" +#include "cute/util/compat.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/integer_subbyte.h" +#include "cutlass/platform/platform.h" +#include "cutlass/util/sycl_event_manager.hpp" +#include "sycl_tla_common.hpp" +#include "sycl_tla_dense_woq_s4_dpas_helpers.hpp" + +#if defined(__clang__) +#pragma clang diagnostic ignored "-Wpass-failed" +#pragma clang diagnostic ignored "-Wdeprecated-declarations" +#elif defined(__GNUC__) +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#endif +#endif + +namespace ark { + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + +namespace dense_woq_s4_dpas { + +using namespace cute; + +template +struct cute_scalar { + using type = ScalarT; +}; + +template <> +struct cute_scalar { + using type = cutlass::half_t; +}; + +template <> +struct cute_scalar { + using type = cutlass::bfloat16_t; +}; + +template +using cute_scalar_t = typename cute_scalar::type; + +class dpas_policy_base { + public: + using WGTile = Shape<_256, _256, _32>; + using SGLayout = Layout, Stride<_4, _1, _0>>; + + using GmemTiledCopyA = void; + using GmemTiledCopyB = void; + using GmemTiledCopyD = void; +}; + +class dpas_w4a16_policy : public dpas_policy_base { + public: + using WGTile = Shape<_128, _256, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; + + using GmemTiledCopyD = XE_STORE_2D<16, 8, 32>; +}; + +class dpas_w4a16_policy_m_8 : public dpas_policy_base { + public: + using WGTile = Shape<_8, _64, _32>; + using SGLayout = Layout, Stride<_4, _1, _0>>; +}; + +class dpas_w4a16_policy_m_16 : public dpas_policy_base { + public: + using WGTile = Shape<_16, _64, _32>; + using SGLayout = Layout, Stride<_4, _1, _0>>; +}; + +class dpas_w4a16_policy_m_32 : public dpas_policy_base { + public: + using WGTile = Shape<_32, _64, _32>; + using SGLayout = Layout, Stride<_4, _1, _0>>; +}; + +class dpas_w4a16_dense_policy_m_4_n128 : public dpas_policy_base { + public: + using WGTile = Shape<_4, _128, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; +}; + +class dpas_w4a16_dense_policy_m_8_n128 : public dpas_policy_base { + public: + using WGTile = Shape<_8, _128, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; +}; + +class dpas_w4a16_dense_policy_m_16 : public dpas_policy_base { + public: + using WGTile = Shape<_16, _128, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; +}; + +class dpas_w4a16_dense_policy_m_32 : public dpas_policy_base { + public: + using WGTile = Shape<_32, _128, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; +}; + +class dpas_w4a16_dense_policy_m_32_n256 : public dpas_policy_base { + public: + using WGTile = Shape<_32, _256, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; +}; + +class dpas_w4a16_dense_policy_m_64_n256 : public dpas_policy_base { + public: + using WGTile = Shape<_64, _256, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; +}; + +class dpas_w4a16_dense_policy_m_128 : public dpas_policy_base { + public: + using WGTile = Shape<_128, _256, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; +}; + +template +class DenseWoqS4DpasName; + +template +CUTE_DEVICE auto make_dense_tensor(T* ptr, int r, int c) { + auto shape = make_shape(r, c); + auto gmem_ptr = make_gmem_ptr(ptr); + if constexpr (LayoutKind == 'C') { + return make_tensor(gmem_ptr, make_layout(shape, make_stride(_1{}, r))); + } else { + return make_tensor(gmem_ptr, make_layout(shape, make_stride(c, _1{}))); + } +} + +template +CUTE_DEVICE void dense_gemm_s4_single_group( + ATensor const& A, + BTensor const& B, + const ElementS* Scales, + const ElementBI* Bias, + DTensor& C, + Coord blk_coord, + TiledMMA const& mma) { + using TA = typename ATensor::element_type; + using TB = typename BTensor::element_type; + static_assert(std::is_same_v, + "dense_gemm_s4_single_group: ElementB must be cutlass::uint4b_t"); + static constexpr int sg_local_range = 16; + + auto item = sycl::ext::oneapi::this_work_item::get_nd_item<3>(); + auto wg_m = get<0>(blk_coord); + auto wg_n = get<1>(blk_coord); + int local_id = item.get_local_linear_id(); + + Tensor cA = make_identity_tensor(A.shape()); + Tensor cB = make_identity_tensor(B.shape()); + Tensor cC = make_identity_tensor(C.shape()); + + auto wg_tile = mma.tile_mnk(); + auto wg_coord = make_coord(wg_m, wg_n, 0); + + 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, _)); + Tensor gC = local_tile(cC, wg_tile, wg_coord, Step<_1, _1, X>{}); + + auto copy_a = get_block_2d_copy_A(mma, A); + auto copy_b = get_block_2d_copy_B(mma, B); + auto copy_c = get_block_2d_copy_D(mma, C); + + auto thr_mma = mma.get_slice(local_id); + auto thr_copy_a = copy_a.get_slice(local_id); + auto thr_copy_b = copy_b.get_slice(local_id); + auto thr_copy_c = copy_c.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 tCrC = thr_mma.partition_sg_fragment_C(gC); + auto tCrC_out = thr_copy_c.partition_sg_fragment_S(gC); + auto tCgC = thr_copy_c.partition_D(gC); + + 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); + + const int prefetch_dist = 3; + constexpr auto barrier_scope = ScopeWorkgroup; + int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); + int k_tile_prefetch = 0; + + static constexpr auto ATOM_M = get<1>(typename TiledMMA::ThrLayoutVMNK{}.shape()); + static constexpr auto ATOM_N = get<2>(typename TiledMMA::ThrLayoutVMNK{}.shape()); + static constexpr auto tile_m = get<0>(wg_tile); + static constexpr auto tile_n = get<1>(wg_tile); + static constexpr auto SG_M = tile_m / ATOM_M; + static constexpr auto SG_N = tile_n / ATOM_N; + static constexpr int sg_n_strides = SG_N / sg_local_range; + + auto n_tile_start = wg_n * tile_n; + auto sg_local_n_coord = cutlass::get_sub_group_id() % ATOM_N; + int sg_local_id = cutlass::get_sub_group_local_id(); + int n_sg_start = sg_local_n_coord * SG_N; + + clear(tCrC); + + float sg_scale[sg_n_strides]; + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + int sg_local_n = sn * sg_local_range + sg_local_id; + sg_scale[sn] = static_cast(Scales[n_tile_start + n_sg_start + sg_local_n]); + } + + 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); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tCrB.size(); ++i) { + tCrB(i) = static_cast(static_cast(tCrB(i)) - 8.0f); + } + + cute::gemm(mma, tCrA, tCrB, tCrC); + + barrier_wait(barrier_scope); + } + + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + float s = sg_scale[sn]; + CUTLASS_PRAGMA_UNROLL + for (int sm = 0; sm < SG_M; ++sm) { + tCrC(sn * SG_M + sm) *= s; + } + } + + if (Bias != nullptr) { + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + int sg_local_n = sn * sg_local_range + sg_local_id; + float b_float = Bias[n_tile_start + n_sg_start + sg_local_n]; + CUTLASS_PRAGMA_UNROLL + for (int sm = 0; sm < SG_M; ++sm) { + tCrC(sn * SG_M + sm) += b_float; + } + } + } + + reorder(tCrC, tCrC_out); + copy(copy_c, tCrC_out, tCgC); +} + +template +CUTE_DEVICE void dense_gemm_s4_pergroup( + ATensor const& A, // (M,K) -- ElementA (bf16/fp16) + BTensor const& B, // (N,K) -- cutlass::uint4b_t (packed nibbles) + const ElementS* Scales, + const ElementBI* Bias, + DTensor& C, // (M,N) -- ElementA + Coord blk_coord, + TiledMMA const& mma) { + using TA = typename ATensor::element_type; + using TB = typename BTensor::element_type; + static_assert(std::is_same_v, + "dense_gemm_s4_pergroup: ElementB must be cutlass::uint4b_t (BestLA S4_CLIP)"); + static constexpr int group_size = GroupSize; + static constexpr int sg_local_range = 16; + auto item = sycl::ext::oneapi::this_work_item::get_nd_item<3>(); + auto wg_m = get<0>(blk_coord); + auto wg_n = get<1>(blk_coord); + int local_id = item.get_local_linear_id(); + + Tensor cA = make_identity_tensor(A.shape()); + Tensor cB = make_identity_tensor(B.shape()); + Tensor cC = make_identity_tensor(C.shape()); + + auto wg_tile = mma.tile_mnk(); + auto wg_coord = make_coord(wg_m, wg_n, 0); + + 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, _)); + Tensor gC = local_tile(cC, wg_tile, wg_coord, Step<_1, _1, X>{}); + + auto copy_a = get_block_2d_copy_A(mma, A); + auto copy_b = get_block_2d_copy_B(mma, B); + auto copy_c = get_block_2d_copy_D(mma, C); + + auto thr_mma = mma.get_slice(local_id); + auto thr_copy_a = copy_a.get_slice(local_id); + auto thr_copy_b = copy_b.get_slice(local_id); + auto thr_copy_c = copy_c.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 tCrC = thr_mma.partition_sg_fragment_C(gC); + auto tCrC_out = thr_copy_c.partition_sg_fragment_S(gC); + auto tCgC = thr_copy_c.partition_D(gC); + + 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); + + // Prefetch distance mirrors `xe_gemm_int_pergroup<>` for now. + // On-hardware perf tuning may want to grow `prefetch_dist` on the + // packed path since the B stream is half the bandwidth. + const int prefetch_dist = 3; + constexpr int prefetch_dist_scale = + ScaleGroupMajor ? (GroupSize == 32 ? 5 : 4) : 3; + constexpr auto barrier_scope = ScopeWorkgroup; + int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); + int k_tile_prefetch = 0; + + static constexpr auto ATOM_M = get<1>(typename TiledMMA::ThrLayoutVMNK{}.shape()); + static constexpr auto ATOM_N = get<2>(typename TiledMMA::ThrLayoutVMNK{}.shape()); + + static constexpr auto tile_m = get<0>(wg_tile); + static constexpr auto tile_n = get<1>(wg_tile); + static constexpr auto tile_k = get<2>(wg_tile); + static constexpr int tile_k_size = int(tile_k); + static constexpr int tiles_per_group = GroupSize / tile_k_size; + + static constexpr auto SG_M = tile_m / ATOM_M; + static constexpr auto SG_N = tile_n / ATOM_N; + + static constexpr int sg_n_strides = SG_N / sg_local_range; + + auto n_tile_start = wg_n * tile_n; + auto sg_local_n_coord = cutlass::get_sub_group_id() % ATOM_N; + int sg_local_id = cutlass::get_sub_group_local_id(); + int n_sg_start = sg_local_n_coord * SG_N; + int group_num = get<1>(A.shape()) / group_size; + int gemm_n = shape<0>(B); + + // Group-local accumulator: same fragment shape as `tCrC`, cleared at + // every scale-group boundary and folded into `tCrC` with a per-N-column + // scale before being reset. Mirrors the INT8 per-group path exactly. + auto tCrC_group = thr_mma.partition_sg_fragment_C(gC); + + clear(tCrC); + clear(tCrC_group); + + // Per-SG per-N scale cache. Same layout / semantics as the INT8 + // per-group path. + float sg_scale[sg_n_strides]; + + 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)); + } + CUTLASS_PRAGMA_UNROLL + for (int pg = 0; pg < prefetch_dist_scale; ++pg) { + if (pg * group_size < shape<1>(A)) { + if constexpr (ScaleGroupMajor) { + auto next_scales_tensor = make_tensor( + make_gmem_ptr(reinterpret_cast( + Scales + pg * gemm_n + n_tile_start + n_sg_start)), + make_layout(make_shape(Int{}, Int<1>{}), + make_stride(Int<1>{}, gemm_n))); + auto prefetch_scales = make_block_2d_prefetch<1>( + make_shape(Int{}, Int<1>{}), next_scales_tensor); + auto thr_prefetch_scales = prefetch_scales.get_slice(sg_local_id); + auto pSgS = thr_prefetch_scales.partition_S( + make_identity_tensor(make_shape(Int{}, Int<1>{}))); + prefetch(prefetch_scales, pSgS(_, 0, 0)); + } else { + auto next_scales_tensor = make_tensor( + make_gmem_ptr(reinterpret_cast( + Scales + (n_tile_start + n_sg_start) * group_num + pg)), + make_layout(make_shape(Int{}, Int<1>{}), + make_stride(group_num, Int<1>{}))); + auto prefetch_scales = make_block_2d_prefetch<1>( + make_shape(Int{}, Int<1>{}), next_scales_tensor); + auto thr_prefetch_scales = prefetch_scales.get_slice(sg_local_id); + auto pSgS = thr_prefetch_scales.partition_S( + make_identity_tensor(make_shape(Int{}, Int<1>{}))); + prefetch(prefetch_scales, pSgS(_, 0, 0)); + } + } + } + + 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); + + bool is_group_start; + int group_idx; + if constexpr (TileAlignedGroup) { + is_group_start = k_tile % tiles_per_group == 0; + group_idx = k_tile / tiles_per_group; + } else { + is_group_start = k_tile * tile_k % group_size == 0; + group_idx = (k_tile * tile_k) / group_size; + } + + if (is_group_start) { + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + int sg_local_n = sn * sg_local_range + sg_local_id; + if constexpr (ScaleGroupMajor) { + sg_scale[sn] = static_cast( + Scales[group_idx * gemm_n + n_tile_start + n_sg_start + sg_local_n]); + } else { + sg_scale[sn] = static_cast( + Scales[(n_tile_start + n_sg_start + sg_local_n) * group_num + group_idx]); + } + } + + if ((group_idx + prefetch_dist_scale) * group_size < shape<1>(A)) { + if constexpr (ScaleGroupMajor) { + auto next_scales_tensor = make_tensor( + make_gmem_ptr(reinterpret_cast( + Scales + (group_idx + prefetch_dist_scale) * gemm_n + + n_tile_start + n_sg_start)), + make_layout(make_shape(Int{}, Int<1>{}), + make_stride(Int<1>{}, gemm_n))); + auto prefetch_scales = make_block_2d_prefetch<1>( + make_shape(Int{}, Int<1>{}), next_scales_tensor); + auto thr_prefetch_scales = prefetch_scales.get_slice(sg_local_id); + auto pSgS = thr_prefetch_scales.partition_S( + make_identity_tensor(make_shape(Int{}, Int<1>{}))); + prefetch(prefetch_scales, pSgS(_, 0, 0)); + } else { + auto next_scales_tensor = make_tensor( + make_gmem_ptr(reinterpret_cast( + Scales + (n_tile_start + n_sg_start) * group_num + + group_idx + prefetch_dist_scale)), + make_layout(make_shape(Int{}, Int<1>{}), + make_stride(group_num, Int<1>{}))); + auto prefetch_scales = make_block_2d_prefetch<1>( + make_shape(Int{}, Int<1>{}), next_scales_tensor); + auto thr_prefetch_scales = prefetch_scales.get_slice(sg_local_id); + auto pSgS = thr_prefetch_scales.partition_S( + make_identity_tensor(make_shape(Int{}, Int<1>{}))); + prefetch(prefetch_scales, pSgS(_, 0, 0)); + } + } + } + + 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); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tCrB.size(); ++i) { + tCrB(i) = static_cast(static_cast(tCrB(i)) - 8.0f); + } + + // HOT MAINLOOP -- MMA accumulates into `tCrC_group`. Per-N scale + // is applied ONCE at the end of the group in the fold block below. + cute::gemm(mma, tCrA, tCrB, tCrC_group); + + bool is_group_end; + if constexpr (TileAlignedGroup) { + is_group_end = ((k_tile + 1) % tiles_per_group == 0) || + (k_tile + 1 == k_tile_count); + } else { + is_group_end = (((k_tile + 1) * tile_k) % group_size == 0) || + (k_tile + 1 == k_tile_count); + } + if (is_group_end) { + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + float s = sg_scale[sn]; + CUTLASS_PRAGMA_UNROLL + for (int sm = 0; sm < SG_M; ++sm) { + const int idx = sn * SG_M + sm; + tCrC(idx) += tCrC_group(idx) * s; + tCrC_group(idx) = 0.0f; + } + } + } + + barrier_wait(barrier_scope); + } + + if (Bias != nullptr) { + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + int sg_local_n = sn * sg_local_range + sg_local_id; + float b_float = Bias[n_tile_start + n_sg_start + sg_local_n]; + CUTLASS_PRAGMA_UNROLL + for (int sm = 0; sm < SG_M; ++sm) { + tCrC(sn * SG_M + sm) += b_float; + } + } + } + + reorder(tCrC, tCrC_out); + copy(copy_c, tCrC_out, tCgC); +} + +template +CUTE_DEVICE void DenseWoqS4GEMM(const ElementA* Activations, + const ElementB* Weights, + const ElementS* Scales, + const ElementBI* Bias, + ElementD* Outputs, + TiledMMA const& mma, + const int32_t gemm_m, + const int32_t gemm_n, + const int32_t gemm_k) { + constexpr char actual_layout_of_B = LayoutKindB ^ ('R' ^ 'C'); + + auto item = sycl::ext::oneapi::this_work_item::get_nd_item<3>(); + int wg_n = item.get_group(0); + int wg_m = item.get_group(1); + + auto A_tensor = make_dense_tensor( + const_cast(Activations), gemm_m, gemm_k); + auto B_tensor = make_dense_tensor( + const_cast(Weights), gemm_n, gemm_k); + auto D_tensor = make_dense_tensor(Outputs, gemm_m, + gemm_n); + auto tile_coord = make_coord(wg_m, wg_n, _, 0); + + if (gemm_k == GroupSize) { + dense_gemm_s4_single_group(A_tensor, B_tensor, Scales, + Bias, D_tensor, tile_coord, + mma); + } else if constexpr (GroupSize % 32 == 0) { + dense_gemm_s4_pergroup( + A_tensor, B_tensor, Scales, Bias, D_tensor, tile_coord, mma); + } else { + dense_gemm_s4_pergroup( + A_tensor, B_tensor, Scales, Bias, D_tensor, tile_coord, mma); + } +} + +template +void DenseWoqS4GEMMLauncherGroup(sycl::queue& stream, + const ElementA* activations, + const ElementB* weights, + const ElementS* scales, + const ElementBI* bias, + ElementD* outputs, + const int gemm_m, + const int gemm_n, + const int gemm_k) { + compat::set_default_queue(stream); + + using ElementA_non_CV = cutlass::platform::remove_cv_t; + auto op = XE_DPAS_TT<8, float, ElementA_non_CV, ElementA_non_CV>{}; + + using WGTile = typename policy::WGTile; + using SGLayout = typename policy::SGLayout; + using MMA = typename TiledMMAHelper, Layout, + SGLayout>::TiledMMA; + auto mma = MMA{}; + + auto wg_tile = mma.tile_mnk(); + const int tile_m = int(get<0>(wg_tile)); + const int tile_n = int(get<1>(wg_tile)); + const int m_tiles = (gemm_m + tile_m - 1) / tile_m; + const int n_tiles = (gemm_n + tile_n - 1) / tile_n; + + auto max_threads_per_workgroup = size(mma); + sycl::range<3> local(1, 1, max_threads_per_workgroup); + sycl::range<3> groups(n_tiles, m_tiles, 1); + + namespace syclex = sycl::ext::oneapi::experimental; + namespace intelex = sycl::ext::intel::experimental; + + syclex::properties kernel_props{syclex::sub_group_size<16>, + intelex::grf_size<256>}; + + using GmemTiledCopyA = typename policy::GmemTiledCopyA; + using GmemTiledCopyB = typename policy::GmemTiledCopyB; + using GmemTiledCopyD = typename policy::GmemTiledCopyD; + + stream.submit([&](sycl::handler& cgh) { + cgh.parallel_for>( + sycl::nd_range<3>{groups * local, local}, kernel_props, [=](auto) { + DenseWoqS4GEMM( + activations, weights, scales, bias, outputs, mma, gemm_m, gemm_n, + gemm_k); + }); + }); +} + +template +bool DenseWoqS4GEMMLauncherDispatchUpTo(sycl::queue& stream, + const ElementA* activations, + const ElementB* weights, + const ElementS* scales, + const ElementBI* bias, + ElementD* outputs, + const int gemm_m, + const int gemm_n, + const int gemm_k, + const int group_size) { + static_assert(MaxGroupSize >= kMinGroupSize && MaxGroupSize <= kMaxGroupSize && + (MaxGroupSize & (MaxGroupSize - 1)) == 0, + "DenseWoqS4GEMMLauncherDispatchUpTo: invalid max group size"); + if (group_size == GroupSize) { + DenseWoqS4GEMMLauncherGroup( + stream, activations, weights, scales, bias, outputs, gemm_m, gemm_n, + gemm_k); + return true; + } + if constexpr (GroupSize < MaxGroupSize) { + return DenseWoqS4GEMMLauncherDispatchUpTo( + stream, activations, weights, scales, bias, outputs, gemm_m, gemm_n, + gemm_k, group_size); + } + return false; +} + +template +bool DenseWoqS4GEMMLauncherDispatch(sycl::queue& stream, + const ElementA* activations, + const ElementB* weights, + const ElementS* scales, + const ElementBI* bias, + ElementD* outputs, + const int gemm_m, + const int gemm_n, + const int gemm_k, + const int group_size) { + return DenseWoqS4GEMMLauncherDispatchUpTo( + stream, activations, weights, scales, bias, outputs, gemm_m, gemm_n, + gemm_k, group_size); +} + +template +void DenseWoqS4GEMMLauncherUpTo(sycl::queue& stream, + const ElementA* activations, + const ElementB* weights, + const ElementS* scales, + const ElementBI* bias, + ElementD* outputs, + const int gemm_m, + const int gemm_n, + const int gemm_k, + const int group_size) { + if (!is_supported_group_size(group_size) || + group_size > MaxGroupSize || + !DenseWoqS4GEMMLauncherDispatchUpTo( + stream, activations, weights, scales, bias, outputs, gemm_m, gemm_n, + gemm_k, group_size)) { + throw std::runtime_error("dense_woq_s4_dpas: unsupported group size"); + } +} + +template +void DenseWoqS4GEMMLauncher(sycl::queue& stream, + const ElementA* activations, + const ElementB* weights, + const ElementS* scales, + const ElementBI* bias, + ElementD* outputs, + const int gemm_m, + const int gemm_n, + const int gemm_k, + const int group_size) { + DenseWoqS4GEMMLauncherUpTo( + stream, activations, weights, scales, bias, outputs, gemm_m, gemm_n, + gemm_k, group_size); +} + +} // namespace dense_woq_s4_dpas + +#endif // ARK_XPU && ARK_SYCL_TLA + +} // namespace ark diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_woq_s4_dpas_helpers.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_woq_s4_dpas_helpers.hpp new file mode 100644 index 000000000..13be8c3e7 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_woq_s4_dpas_helpers.hpp @@ -0,0 +1,41 @@ +#pragma once + +#include "sycl_tla_common.hpp" + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + +namespace ark { +namespace dense_woq_s4_dpas { + +inline constexpr int kMinGroupSize = 32; +inline constexpr int kMaxGroupSize = 4096; +inline constexpr int kPackedScaleMaxGroupSize = 128; + +inline bool is_supported_group_size(int group_size) { + return group_size >= kMinGroupSize && group_size <= kMaxGroupSize && + (group_size & (group_size - 1)) == 0; +} + +namespace detail { + +#define ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE(prefix) \ + void prefix##_group_n(sycl::queue* q, const void* activations, const void* weights, const void* scales, \ + const void* bias, void* outputs, int gemm_m, int gemm_n, int gemm_k, int group_size); \ + void prefix##_n_group(sycl::queue* q, const void* activations, const void* weights, const void* scales, \ + const void* bias, void* outputs, int gemm_m, int gemm_n, int gemm_k, int group_size) + +ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE(run_m4_n128); +ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE(run_m8_n128); +ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE(run_m16); +ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE(run_m32); +ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE(run_m32_n256); +ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE(run_m64_n256); +ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE(run_m128); + +#undef ARK_DECLARE_DENSE_WOQ_S4_DPAS_ROUTE + +} // namespace detail +} // namespace dense_woq_s4_dpas +} // namespace ark + +#endif \ No newline at end of file diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index 2ee3309ee..c033b7a77 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -855,4 +855,4 @@ inline bool moe_prefill_dpas_s4_pergroup_shape_ok(int N, int K, } // namespace moe_dpas_s4 } // namespace ark -#endif // ARK_XPU && ARK_SYCL_TLA +#endif // ARK_XPU && ARK_SYCL_TLA \ No newline at end of file diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp index f7f166a58..c7c0862ff 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp @@ -67,8 +67,8 @@ struct env_params { } static inline void env_i(const char* envstr, int& default_) { - const char* log_level_env = std::getenv(envstr); - if (log_level_env != nullptr) default_ = std::stoi(log_level_env); + const char* env_value = std::getenv(envstr); + if (env_value != nullptr) default_ = std::stoi(env_value); } }; @@ -567,4 +567,4 @@ struct QuantParam { inline int blks() { return k / blocksize; } }; -} // namespace ark \ No newline at end of file +} // namespace ark diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp index 7bfbc40bc..3f2eff1a3 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp @@ -25,6 +25,10 @@ #include "sycl_tla_common.hpp" #endif +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) +#include "sycl_tla_dense_woq_s4_dpas_helpers.hpp" +#endif + #if ARK_XPU #include "sycl_s8_wrapper.hpp" #endif @@ -62,6 +66,20 @@ class XpuWrapper { return p->blks() * p->n * bestla_dtype_bytes(p->scale_type); } + static inline bool use_dpas_scale_layout(QuantParam* p) { +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + return p->weight_type == BTLA_DTYPE::S4 && p->scale_type == BTLA_DTYPE::F16 && !p->asym && + dense_woq_s4_dpas::is_supported_group_size(p->blocksize) && + p->blocksize <= dense_woq_s4_dpas::kPackedScaleMaxGroupSize; +#else + return false; +#endif + } + + static inline size_t get_dpas_scale_size(QuantParam* p) { + return use_dpas_scale_layout(p) ? get_scale_size(p) : 0; + } + static inline size_t get_zp_size(QuantParam* p) { using namespace bestla::utils; if (!p->asym) return 0; @@ -78,7 +96,7 @@ class XpuWrapper { return nblk * p->n * bestla_dtype_bytes(p->scale_type); } - static inline size_t get_packw_size(QuantParam* p) { + static inline size_t get_packw_base_size(QuantParam* p) { size_t size = get_packw_qsize(p); size += get_scale_size(p); size += get_zp_size(p); @@ -86,6 +104,8 @@ class XpuWrapper { return size; } + static inline size_t get_packw_size(QuantParam* p) { return get_packw_base_size(p) + get_dpas_scale_size(p); } + static inline size_t get_scale_offset(QuantParam* p) { size_t size = get_packw_qsize(p); return size; @@ -103,6 +123,14 @@ class XpuWrapper { return size; } + static inline size_t get_dpas_scale_offset(QuantParam* p) { return get_packw_base_size(p); } + + static inline bool has_dpas_scale_layout(QuantParam* p, size_t blob_count) { + if (!use_dpas_scale_layout(p)) return false; + if (blob_count == 0) return false; + return blob_count >= get_dpas_scale_offset(p) + get_dpas_scale_size(p); + } + static bool can_comps8(QuantParam* p) { if (p->asym && p->weight_type == BTLA_DTYPE::S8) return false; if (p->blocksize == -1 || p->blocksize == p->k) return true; @@ -237,6 +265,9 @@ class XpuWrapper { }); }; q->submit(ker); + if (use_dpas_scale_layout(p)) { + q->memcpy(blobptr + get_dpas_scale_offset(p), scaleptr, get_dpas_scale_size(p)); + } if (rescale(p)) { #ifdef ARK_RESCALE auto scalext_ptr = (int8_t*)blobptr + get_scalext_offset(p); @@ -359,7 +390,7 @@ class XpuWrapper { static void unpackq(BTLA_DTYPE outt, int8_t* blob, void* optr, QuantParam* p, sycl::queue* q, size_t blob_count = 0) { if (blob_count > 0) { - auto expected = get_packw_size(p); + auto expected = get_packw_base_size(p); if (blob_count < expected) { throw std::runtime_error("Corrupt packed weight: blob size (" + std::to_string(blob_count) + ") less than expected (" + std::to_string(expected) + ")"); @@ -526,9 +557,91 @@ class XpuWrapper { } } +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + static constexpr size_t kWoqS4DpasMaxM = 128; + + static inline bool woq_s4_dpas_group_size_ok(int group_size) { + return dense_woq_s4_dpas::is_supported_group_size(group_size); + } + + static inline int woq_s4_dpas_tile_n(size_t m, QuantParam* p) { + if (m <= 16) return 128; + if (m <= 32) return p->blocksize == p->k ? 256 : 128; + return 256; + } + + static inline bool woq_s4_dpas_shape_ok(size_t m, QuantParam* p) { + if (m <= 1 || m > kWoqS4DpasMaxM) return false; + if (m > static_cast(std::numeric_limits::max())) return false; + if (p->blocksize <= 0 || p->n % woq_s4_dpas_tile_n(m, p) != 0 || (p->k & 1) != 0) return false; + if (p->k % p->blocksize != 0) return false; + return woq_s4_dpas_group_size_ok(p->blocksize); + } + + static bool woq_try_s4_dpas(sycl::queue* q, size_t m, QuantParam* p, const void* matA, const void* blobB, + void* matC, const void* bias, BTLA_DTYPE outt, size_t blob_count) { + if (p->weight_type != BTLA_DTYPE::S4 || p->scale_type != BTLA_DTYPE::F16 || outt != BTLA_DTYPE::F16 || p->asym) { + return false; + } + if (!woq_s4_dpas_shape_ok(m, p)) { + return false; + } + const bool use_dpas_scales = has_dpas_scale_layout(p, blob_count); + const auto scale_offset = use_dpas_scales ? get_dpas_scale_offset(p) : get_scale_offset(p); + const auto* scales_ptr = reinterpret_cast(blobB) + scale_offset; + +#define ARK_WOQ_S4_DPAS_LAUNCH(route, policy_name) \ + do { \ + if (env_params::Instance()->verbose <= 1) { \ + std::fprintf(stdout, \ + "[ARK_WOQ_S4_DPAS] launch:%s m=%zu n=%d k=%d blocksize=%d scale_layout=%s\n", \ + policy_name, m, p->n, p->k, p->blocksize, \ + use_dpas_scales ? "group_n" : "n_group"); \ + } \ + if (use_dpas_scales) { \ + ark::dense_woq_s4_dpas::detail::route##_group_n( \ + q, matA, blobB, scales_ptr, bias, matC, static_cast(m), p->n, \ + p->k, p->blocksize); \ + } else { \ + ark::dense_woq_s4_dpas::detail::route##_n_group( \ + q, matA, blobB, scales_ptr, bias, matC, static_cast(m), p->n, \ + p->k, p->blocksize); \ + } \ + } while (false); + + if (m <= 4) { + ARK_WOQ_S4_DPAS_LAUNCH(run_m4_n128, "dpas_w4a16_dense_policy_m_4_n128") + } else if (m <= 8) { + ARK_WOQ_S4_DPAS_LAUNCH(run_m8_n128, "dpas_w4a16_dense_policy_m_8_n128") + } else if (m <= 16) { + ARK_WOQ_S4_DPAS_LAUNCH(run_m16, "dpas_w4a16_dense_policy_m_16") + } else if (m <= 32) { + if (p->blocksize == p->k) { + ARK_WOQ_S4_DPAS_LAUNCH(run_m32_n256, "dpas_w4a16_dense_policy_m_32_n256") + } else { + ARK_WOQ_S4_DPAS_LAUNCH(run_m32, "dpas_w4a16_dense_policy_m_32") + } + } else if (m <= 64) { + ARK_WOQ_S4_DPAS_LAUNCH(run_m64_n256, "dpas_w4a16_dense_policy_m_64_n256") + } else if (m <= 128) { + ARK_WOQ_S4_DPAS_LAUNCH(run_m128, "dpas_w4a16_dense_policy_m_128") + } else { + return false; + } +#undef ARK_WOQ_S4_DPAS_LAUNCH + return true; + } +#endif + static int woq_gemv(sycl::queue* q, size_t m, QuantParam* p, const void* matA, const void* blobB, void* matC, - const void* bias, BTLA_DTYPE outt) { - if (m > 1) return -2; + const void* bias, BTLA_DTYPE outt, size_t blob_count = 0) { + if (m > 1) { +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + return woq_try_s4_dpas(q, m, p, matA, blobB, matC, bias, outt, blob_count) ? 0 : -1; +#else + return -1; +#endif + } using namespace bestla; using namespace bestla::sycl_prologue_b; auto qptr = (uint8_t*)blobB; @@ -703,15 +816,15 @@ class XpuWrapper { static void woq_gemm(int m, const void* a, const void* b, void* c, const void* bias, BTLA_DTYPE acdt, QuantParam* p, sycl::queue* q, size_t blob_count = 0) { if (blob_count > 0) { - auto expected = get_packw_size(p); + auto expected = get_packw_base_size(p); if (blob_count < expected) { throw std::runtime_error("Corrupt packed weight: blob size (" + std::to_string(blob_count) + ") less than expected (" + std::to_string(expected) + ")"); } } - auto ret = woq_gemv(q, m, p, a, b, c, bias, acdt); + auto ret = woq_gemv(q, m, p, a, b, c, bias, acdt, blob_count); if (ret) { - + check_compute_type(p); if (p->compute_type != BTLA_DTYPE::S8) { size_t elesize = bestla::utils::bestla_dtype_bytes(acdt); diff --git a/auto_round_extension/ark/test/test_woq_s4.py b/auto_round_extension/ark/test/test_woq_s4.py new file mode 100644 index 000000000..c66a6afc6 --- /dev/null +++ b/auto_round_extension/ark/test/test_woq_s4.py @@ -0,0 +1,357 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +import importlib +import sys +import time +import unittest +from importlib import metadata +from pathlib import Path + +import auto_round_kernel as ark +import torch +from ut_utils import gen_weis8 + +M_VALUES = [1, 2, 4, 8, 16, 32, 64, 128] +N = 16384 +K = 4096 +BLOCKSIZE = 32 +ACCURACY_N = 256 +ACCURACY_K = 256 +OUTPUT_GUARD_VALUE = -1234.0 +DTYPE = torch.float16 +DEVICE = "xpu" +COMPUTE_TYPE = "int8" +WEIGHT_TYPE = "int4" +SCALE_TYPE = "fp16" +ASYM = False +WARMUP_LIMIT = 1000 + + +def _print_config_types(): + config_values = { + "M_VALUES": M_VALUES, + "N": N, + "K": K, + "BLOCKSIZE": BLOCKSIZE, + "DTYPE": DTYPE, + "DEVICE": DEVICE, + "COMPUTE_TYPE": COMPUTE_TYPE, + "WEIGHT_TYPE": WEIGHT_TYPE, + "SCALE_TYPE": SCALE_TYPE, + } + print("\n=== Config types ===") + for name, value in config_values.items(): + print(f"{name}: {value}") + + +def _sync_xpu(): + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.synchronize() + + +def _is_power_of_two(value): + return value > 0 and (value & (value - 1)) == 0 + + +def _ark_expected_route(m, n=N, k=K, blocksize=BLOCKSIZE): + if m <= 1: + return "bestla_s4_gemv" + if m > 128: + return "woqgemm_s8(unpack_s4_to_s8)" + if n % 64 != 0 or (k & 1) != 0 or blocksize <= 0: + return "woqgemm_s8(unpack_s4_to_s8)" + if k % blocksize != 0: + return "woqgemm_s8(unpack_s4_to_s8)" + if blocksize < 32 or blocksize > 4096 or not _is_power_of_two(blocksize): + return "woqgemm_s8(unpack_s4_to_s8)" + return "woq_s4_dpas" + + +def _dense_s4_dpas_tile_m(m): + if m <= 4: + return 4 + if m <= 8: + return 8 + if m <= 16: + return 16 + if m <= 32: + return 32 + if m <= 64: + return 64 + if m <= 128: + return 128 + raise ValueError(f"m={m} is outside dense S4 DPAS coverage") + + +def _has_torch_int4_op(): + return hasattr(torch.ops, "_xpu_C") and hasattr(torch.ops._xpu_C, "int4_gemm_w4a16") + + +def _prepare_xpu_kernel_import_path(): + repo_root = Path(__file__).resolve().parents[4] + xpu_kernel_src = repo_root / "vllm-xpu-kernels" + package_dir = xpu_kernel_src / "vllm_xpu_kernels" + if not package_dir.exists(): + return + + try: + local_src = xpu_kernel_src.resolve() + except OSError: + return + + if any(package_dir.glob("_xpu_C*.so")): + local_src_str = str(local_src) + if local_src_str not in sys.path: + sys.path.insert(0, local_src_str) + return + + sys.path[:] = [entry for entry in sys.path if not entry or Path(entry).resolve() != local_src] + + +def _installed_xpu_kernel_extensions(): + try: + files = metadata.files("vllm-xpu-kernels") or [] + except metadata.PackageNotFoundError: + return "package metadata not found" + + extensions = sorted(str(file) for file in files if str(file).endswith((".so", ".pyd")) or "_xpu_C" in str(file)) + return ", ".join(extensions[:16]) if extensions else "no extension files found" + + +def _register_xpu_ops(): + if _has_torch_int4_op(): + return True + + _prepare_xpu_kernel_import_path() + import_errors = [] + for module_name in ("vllm_xpu_kernels._xpu_C", "vllm.platforms.xpu", "vllm._xpu_ops"): + try: + importlib.import_module(module_name) + except Exception as exc: + import_errors.append(f"{module_name}: {exc}") + if _has_torch_int4_op(): + return True + + print("\n[torch int4_gemm_w4a16] skip: cannot register torch.ops._xpu_C.int4_gemm_w4a16") + for error in import_errors: + print(f" {error}") + print(f" vllm-xpu-kernels extension files: {_installed_xpu_kernel_extensions()}") + return False + + +def _rand_packed_int4(size, dtype=torch.int32, device=DEVICE): + rand = torch.randint(-128, 128, [size // 2], device=device).to(torch.int8) + return rand.view(dtype=dtype) + + +def _torch_int4_case(n=N, k=K, blocksize=BLOCKSIZE, dtype=DTYPE, device=DEVICE): + weight = _rand_packed_int4(k * n, torch.int32, device).reshape(k // 8, n) + weight_nt = weight.transpose(0, 1).contiguous().transpose(0, 1) + scales = torch.rand([k // blocksize, n], device=device, dtype=dtype) / 300 + 0.002 + zero_points = torch.tensor([8], device=device, dtype=torch.int8) + bias = torch.randn(n, device=device, dtype=dtype) + return weight_nt, scales, zero_points, bias + + +def _ark_case(m, n=N, k=K, blocksize=BLOCKSIZE, dtype=DTYPE, device=DEVICE): + torch.manual_seed(0) + raw_s8_wei = gen_weis8(WEIGHT_TYPE, device, k, n) + scales = torch.rand(k // blocksize, n, dtype=dtype, device=device) / 300 + 0.002 + bias = torch.randn(1, n, dtype=dtype, device=device) + zp = torch.Tensor() + + packw = ark.repack_quantized_weight(raw_s8_wei, scales, zp, blocksize, COMPUTE_TYPE, WEIGHT_TYPE, SCALE_TYPE, ASYM) + revert_wei_t = ark.unpack_weight(packw, dtype, n, k, blocksize, COMPUTE_TYPE, WEIGHT_TYPE, SCALE_TYPE, ASYM) + revert_wei = revert_wei_t.t() + ref_weight = raw_s8_wei.to(dtype) * scales.repeat_interleave(repeats=blocksize, dim=0) + assert torch.allclose(revert_wei, ref_weight) + + activation = torch.randn(m, k, dtype=dtype, device=device) - 0.5 + ref_c = torch.matmul(activation, revert_wei) + bias + return activation, packw, bias, ref_c + + +def _runs_for_m(m): + return 1000 + + +def _batch_for_m(m): + return 64 if m == 1 else 8 + + +def _warmup_for_runs(runs): + return min(runs, WARMUP_LIMIT) + + +def _repeat_ark_blob(packw, batch): + return packw.unsqueeze(0).repeat(batch, 1) + + +def _repeat_activation(activation, batch): + return activation.unsqueeze(0).repeat(batch, 1, 1) + + +def _repeat_nt_weight(weight, batch): + weight_nk = weight.transpose(0, 1).contiguous() + return weight_nk.unsqueeze(0).repeat(batch, 1, 1).transpose(1, 2) + + +def _guarded_output(m, n, dtype=DTYPE, device=DEVICE): + guard_rows = _dense_s4_dpas_tile_m(m) - m + 1 + storage = torch.full((m + guard_rows, n), OUTPUT_GUARD_VALUE, dtype=dtype, device=device) + return storage[:m], storage[m:] + + +def _assert_ark_woqgemm_matches_reference_with_output_bounds(m, n=ACCURACY_N, k=ACCURACY_K, blocksize=BLOCKSIZE): + activation, packw, bias, ref_c = _ark_case(m, n=n, k=k, blocksize=blocksize) + output, guard = _guarded_output(m, n) + + actual = ark.woqgemm( + activation, + packw, + bias, + n, + k, + blocksize, + COMPUTE_TYPE, + WEIGHT_TYPE, + SCALE_TYPE, + ASYM, + out=output, + ) + _sync_xpu() + + assert actual.data_ptr() == output.data_ptr() + assert torch.allclose(actual, ref_c, rtol=0.1, atol=2.0) + assert torch.all(guard == OUTPUT_GUARD_VALUE) + + +def _memory_bytes(m, n=N, k=K, blocksize=BLOCKSIZE, dtype=DTYPE): + element_size = torch.empty((), dtype=dtype).element_size() + return m * k * element_size + m * n * element_size + n * k // 2 + (k // blocksize) * n * element_size + + +def _benchmark_loop(call, runs, warmup): + output = None + for i in range(warmup): + output = call(i) + _sync_xpu() + + start = time.perf_counter() + for i in range(runs): + output = call(i) + _sync_xpu() + return output, (time.perf_counter() - start) / runs + + +def _print_perf(m, batch, warmup, runs, op_name, dur, route=None): + ops = m * N * K * 2 + memsize = _memory_bytes(m) + route_text = f", route={route}" if route is not None else "" + print( + f"\n m={m}, n={N}, k={K}, blocksize={BLOCKSIZE}, batch={batch}, warmup={warmup}, runs={runs}, op={op_name}{route_text}" + ) + print(f"[Performance] Time: {dur * 1000:.4f} ms") + print(f" GFLOPS: {ops / dur / 1e9:.2f}") + print(f" Bandwidth: {memsize / dur / 1e9:.2f} GB/s") + + +def run_ark_woqgemm(): + print("\n=== ARK woqgemm ===") + print("Timed loops use the same warmup/runs/batch policy as the torch oneDNN path.") + for m in M_VALUES: + route = _ark_expected_route(m) + print(f"\n[ARK route] m={m}: {route}") + activation, packw, bias, ref_c = _ark_case(m) + runs = _runs_for_m(m) + batch = _batch_for_m(m) + warmup = _warmup_for_runs(runs) + activation_set = _repeat_activation(activation, batch) + packw_set = _repeat_ark_blob(packw, batch) + output_set = torch.empty(batch, m, N, dtype=activation.dtype, device=activation.device) + + def call(i): + idx = i % batch + return ark.woqgemm( + activation_set[idx], + packw_set[idx], + bias, + N, + K, + BLOCKSIZE, + COMPUTE_TYPE, + WEIGHT_TYPE, + SCALE_TYPE, + ASYM, + out=output_set[idx], + ) + + output, dur = _benchmark_loop(call, runs, warmup) + diff = abs(ref_c - output) + print( + f" Max Diff: {diff.max().item():.6f}, Mean Diff: {diff.mean().item():.6f}, " + f"ref mean:{ref_c.mean():.6f}, OUT mean:{output.mean():.6f}" + ) + assert torch.allclose(output, ref_c, rtol=0.1, atol=2.0) + _print_perf(m, batch, warmup, runs, "ark.woqgemm", dur, route) + + +def run_torch_int4_gemm_w4a16(): + print("\n=== torch.ops._xpu_C.int4_gemm_w4a16 ===") + if not hasattr(torch, "xpu") or not torch.xpu.is_available(): + print("[torch int4_gemm_w4a16] skip: no XPU device") + return + if not _register_xpu_ops(): + return + + torch.manual_seed(0) + weight, scales, zero_points, bias = _torch_int4_case() + for m in M_VALUES: + activation = torch.randn(m, K, dtype=DTYPE, device=DEVICE) - 0.5 + runs = _runs_for_m(m) + batch = _batch_for_m(m) + warmup = _warmup_for_runs(runs) + activation_set = _repeat_activation(activation, batch) + weight_set = _repeat_nt_weight(weight, batch) + + def call(i): + idx = i % batch + return torch.ops._xpu_C.int4_gemm_w4a16( + activation_set[idx], weight_set[idx], bias, scales, zero_points, BLOCKSIZE, None + ) + + try: + _, dur = _benchmark_loop(call, runs, warmup) + except Exception as exc: + print(f"\n[torch int4_gemm_w4a16] m={m} skip: {exc}") + continue + + _print_perf(m, batch, warmup, runs, "torch.ops._xpu_C.int4_gemm_w4a16", dur, "oneDNN_w4a16_int4") + + +def test_ark_woqgemm_sym_accuracy_and_output_bounds(): + if not hasattr(torch, "xpu") or not torch.xpu.is_available(): + raise unittest.SkipTest("No XPU Device") + + DENSE_S4_SYM_M_VALUES = [1, 2, 4, 8, 16, 32, 64, 128] + for m in DENSE_S4_SYM_M_VALUES: + expected_route = "bestla_s4_gemv" if m == 1 else "woq_s4_dpas" + assert _ark_expected_route(m, ACCURACY_N, ACCURACY_K, BLOCKSIZE) == expected_route + _assert_ark_woqgemm_matches_reference_with_output_bounds(m) + + +def test_dense_s4_dpas_non_tile_aligned_m_accuracy_and_output_bounds(): + if not hasattr(torch, "xpu") or not torch.xpu.is_available(): + raise unittest.SkipTest("No XPU Device") + + DENSE_S4_DPAS_PARTIAL_M_VALUES = [5, 6, 7, 9, 17, 33, 65] + for m in DENSE_S4_DPAS_PARTIAL_M_VALUES: + assert _ark_expected_route(m, ACCURACY_N, ACCURACY_K, BLOCKSIZE) == "woq_s4_dpas" + _assert_ark_woqgemm_matches_reference_with_output_bounds(m) + + +if __name__ == "__main__": + _print_config_types() + run_ark_woqgemm() + run_torch_int4_gemm_w4a16()