diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp index ee36f75164..b82570668d 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp @@ -719,6 +719,11 @@ struct BlockwiseGemmWmmaops_pipeline_v1, bool>::type = false> __device__ void Run(const AGridDesc& a_grid_desc, const ABlockDesc& a_block_desc, @@ -753,8 +759,11 @@ struct BlockwiseGemmWmmaops_pipeline_v1; + constexpr index_t KPerWaveBlock = wmma_gemm.GetKPerWaveBlk(); auto a_thread_buf = make_static_buffer( @@ -771,6 +780,10 @@ struct BlockwiseGemmWmmaops_pipeline_v1(num_loop_per_scale == 1); + if constexpr(HasBZp) + { + b_zero_point_struct.template GlobalLoad<0>(num_loop_per_scale == 1); + } // Local prefill 1 a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); @@ -802,7 +815,7 @@ struct BlockwiseGemmWmmaops_pipeline_v1{}([&](auto n0) { b_thread_copy_.Run( @@ -817,6 +830,30 @@ struct BlockwiseGemmWmmaops_pipeline_v1{}([&](auto n0) { + b_thread_copy_.Run( + b_block_desc_k0_n0_n1_n2_k1, + make_tuple(I0, n0, k0_offset + k0_inner, I0, I0, I0, I0), + b_block_buf, + b_scale_struct.scale_thread_bufs(I0)[Number< + n0 * BScaleStruct::num_scale_k_block + + (k0_offset + k0_inner) / BScaleStruct::num_scale_krepeat>{}], + b_zero_point_struct.scale_thread_bufs( + I0)[Number{}], + b_thread_desc_, + make_tuple(I0, n0, k0_inner, I0, I0, I0, I0), + b_thread_buf); + }); + } }); __builtin_amdgcn_sched_barrier(0); @@ -919,6 +956,10 @@ struct BlockwiseGemmWmmaops_pipeline_v1((i + 2) % num_loop_per_scale == 0); + if constexpr(HasBZp) + { + b_zero_point_struct.template GlobalLoad<0>((i + 2) % num_loop_per_scale == 0); + } if constexpr(ck::is_same::value == false) { block_sync_lds(); diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp index 03146f22ee..48e131a1ea 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp @@ -336,6 +336,10 @@ struct BlockwiseGemmWmmaops_pipeline_v3, bool>::type = false> __device__ void Run(const AGridDesc& a_grid_desc, const ABlockDesc& a_block_desc, @@ -370,8 +375,12 @@ struct BlockwiseGemmWmmaops_pipeline_v3{p_a}, std::array{p_b}, @@ -216,7 +220,9 @@ struct DeviceGemm_BScale_Wmma_CShuffleV3 : public DeviceGemmV2BScale()[Number<0>{}]; } +// AIESW-32176: asymmetric int4 dequant. +// With symmetric (uint4b8) the existing path computes (nibble - 8) * scale. +// With asymmetric (zero point zp per group), the desired result is +// (nibble - zp) * scale. Equivalently: +// (nibble - 8) * scale - (zp - 8) * scale == result_sym - scaled_zp +// so we precompute scaled_zp = (zp - 8) * scale once at load time and +// subtract it after the existing scale multiply. One extra fp16 vector +// subtract per dequant pack; everything else is identical to +// i4_to_half4_scale, including the magic SUB/MUL/ADD constants. +__device__ inline half4_t +i4_to_half4_zp_scale(int q, const ck::half2_t& scale, const ck::half2_t& scaled_zp) +{ + const int LO = 0x000f000f; + const int HI = 0x00f000f0; + const int EX = 0x64006400; + + int lo = (q & LO) | EX; + int hi = (q & HI) | EX; + + const int SUB = 0xE408E408; // half2 {-1032, -1032} + const int MUL = 0x2c002c00; // half2 {1 / 16, 1 / 16} + const int ADD = 0xd480d480; // half2 {-72, -72} + + vector_type res; + + res.template AsType()(Number<0>{}) = bit_cast(lo) + bit_cast(SUB); + res.template AsType()(Number<1>{}) = + bit_cast(hi) * bit_cast(MUL) + bit_cast(ADD); + + res.template AsType()(Number<0>{}) = + res.template AsType()(Number<0>{}) * scale - scaled_zp; + res.template AsType()(Number<1>{}) = + res.template AsType()(Number<1>{}) * scale - scaled_zp; + + return res.template AsType()[Number<0>{}]; +} + __device__ inline f8x4_t i4_to_f8x4(int q) { const int LO = 0x000f000f; @@ -299,6 +336,51 @@ struct DequantPack8 constexpr const static bool is_pack8_invocable = true; }; +// AIESW-32176: asymmetric variant of DequantPack8. +// Same as DequantPack8 but takes an additional `scaled_zp` per-group +// parameter (= (zp - 8) * scale precomputed at weight load) and +// subtracts it from the result so the dequant matches AWQ's +// (nibble - zp) * scale formula. See i4_to_half4_zp_scale. +struct DequantPack8WithZp +{ + static constexpr const char* name = "DequantPack8WithZp"; + + template + __host__ __device__ void operator()(Y& y, const X& x, const S& s, const ZP& zp) const; + + __host__ __device__ constexpr void operator()(ck::half8_t& y, + const ck::pk_i4x4_t& x, + const ck::half2_t& s, + const ck::half2_t& scaled_zp) const + { +#if CK_USE_PK4_LAYOUT_SHUFFLE + vector_type result; + + result.template AsType()(Number<0>{}) = + i4_to_half4_zp_scale(bit_cast(x), s, scaled_zp); + result.template AsType()(Number<1>{}) = + i4_to_half4_zp_scale(bit_cast(x) >> 8, s, scaled_zp); + + y = result.template AsType()[Number<0>{}]; +#else + // Reference (slow) path for the !CK_USE_PK4_LAYOUT_SHUFFLE config: + // do the unscaled dequant via existing type_convert and then apply + // (* scale) - scaled_zp manually. + vector_type dst; + vector_type src{x}; + + static_for<0, 4, 1>{}([&](auto i) { + auto v = type_convert(src.template AsType()[i]); + dst.template AsType()(i) = v * s - scaled_zp; + }); + + y = dst.template AsType()[Number<0>{}]; +#endif + } + + constexpr const static bool is_pack8_invocable = true; +}; + struct PassThroughPack2 { static constexpr const char* name = "PassThroughPack2"; diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_ab_scale.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_ab_scale.hpp index 92561d00d4..f64c388557 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_ab_scale.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_ab_scale.hpp @@ -333,7 +333,8 @@ struct GridwiseGemm_wmma_cshuffle_v3_ab_scale AElementwiseOperation a_element_op_, BElementwiseOperation b_element_op_, CDEElementwiseOperation cde_element_op_, - bool is_reduce_ = false) + bool is_reduce_ = false, + const BScaleType* p_b_zero_point_grid_ = nullptr) : Problem{M_, N_, K_, @@ -350,6 +351,7 @@ struct GridwiseGemm_wmma_cshuffle_v3_ab_scale p_e_grid{p_e_grid_}, p_a_scale_grid{p_a_scale_grid_}, p_b_scale_grid{p_b_scale_grid_}, + p_b_zero_point_grid{p_b_zero_point_grid_}, a_element_op{a_element_op_}, b_element_op{b_element_op_}, cde_element_op{cde_element_op_}, @@ -395,6 +397,9 @@ struct GridwiseGemm_wmma_cshuffle_v3_ab_scale const AScaleType* p_a_scale_grid; const BScaleType* p_b_scale_grid; + // AIESW-32176: optional per-group precomputed (zp - 8) * scale for + // asymmetric int4. nullptr for symmetric — kernel skips zp subtract. + const BScaleType* p_b_zero_point_grid; const AElementwiseOperation a_element_op; const BElementwiseOperation b_element_op; const CDEElementwiseOperation cde_element_op; @@ -707,6 +712,9 @@ struct GridwiseGemm_wmma_cshuffle_v3_ab_scale } } + // AIESW-32176: p_b_zero_point_grid (nullptr for symmetric callers) selects + // the asymmetric dequant path. When non-null the caller has precomputed + // (zp - 8) * scale per group at weight load. template (problem, p_b_scale_grid, block_n_id); + // AIESW-32176: BZeroPoint struct (mirror of BScale; same shape, same + // type, same load cadence). When p_b_zero_point_grid is nullptr we keep + // BZeroPoint as Empty so the Base::Run path stays bit-identical to + // symmetric. const index_t num_k_block_per_scale = GetKBlockPerScale(); - Base::template Run(p_as_grid, - p_bs_grid, - p_ds_grid, - p_e_grid, - p_shared, - as_grid_desc_ak0_m_ak1, - bs_grid_desc_bk0_n_bk1, - ds_grid_desc_mblock_mperblock_nblock_nperblock, - e_grid_desc_mblock_mperblock_nblock_nperblock, - a_element_op, - b_element_op, - cde_element_op, - block_m_id, - block_n_id, - num_k_block_per_scale, - a_scale_struct, - b_scale_struct, - epilogue_args, - A_k_id, - B_k_id); + if(p_b_zero_point_grid == nullptr) + { + Base::template Run(p_as_grid, + p_bs_grid, + p_ds_grid, + p_e_grid, + p_shared, + as_grid_desc_ak0_m_ak1, + bs_grid_desc_bk0_n_bk1, + ds_grid_desc_mblock_mperblock_nblock_nperblock, + e_grid_desc_mblock_mperblock_nblock_nperblock, + a_element_op, + b_element_op, + cde_element_op, + block_m_id, + block_n_id, + num_k_block_per_scale, + a_scale_struct, + b_scale_struct, + epilogue_args, + A_k_id, + B_k_id); + } + else + { + auto b_zero_point_struct = MakeBScale<1>(problem, p_b_zero_point_grid, block_n_id); + + Base::template Run( + p_as_grid, + p_bs_grid, + p_ds_grid, + p_e_grid, + p_shared, + as_grid_desc_ak0_m_ak1, + bs_grid_desc_bk0_n_bk1, + ds_grid_desc_mblock_mperblock_nblock_nperblock, + e_grid_desc_mblock_mperblock_nblock_nperblock, + a_element_op, + b_element_op, + cde_element_op, + block_m_id, + block_n_id, + num_k_block_per_scale, + a_scale_struct, + b_scale_struct, + epilogue_args, + A_k_id, + B_k_id, + /*k_batch=*/1, + b_zero_point_struct); + } } // NOTE: Wrapper function to have __global__ function in common @@ -861,7 +915,11 @@ struct GridwiseGemm_wmma_cshuffle_v3_ab_scale karg.cde_element_op, epilogue_args, A_k_id, - B_k_id); + B_k_id, + // AIESW-32176: optional zp pointer; nullptr keeps the symmetric path. + // No splitk offset applied because zp shares the scale layout — the + // splitk K offset is absorbed into the same scale stride. + karg.p_b_zero_point_grid); } }; diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp index a843a94777..517e24d869 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp @@ -1364,6 +1364,8 @@ struct GridwiseGemm_wmma_cshuffle_v3_base // Note: arguments k_batch and k_id should be set if splitk is used // with implicit gemm (no pointer shift but shift using tensor descriptors) + // AIESW-32176: BZeroPointStruct (defaulted to Empty) is forwarded to the + // blockwise pipeline's Run() to pick up the asymmetric dequant path. template + TailNumber TailNum = TailNumber::Odd, + typename BZeroPointStruct = typename BlockwiseGemmPipe::Empty> __device__ static void Run(AsGridPointer p_as_grid, BsGridPointer p_bs_grid, DsGridPointer p_ds_grid, @@ -1394,9 +1397,10 @@ struct GridwiseGemm_wmma_cshuffle_v3_base AScaleStruct& a_scale_struct, BScaleStruct& b_scale_struct, EpilogueArgument& epilogue_args, - const index_t A_k_id = 0, - const index_t B_k_id = 0, - const index_t k_batch = 1) + const index_t A_k_id = 0, + const index_t B_k_id = 0, + const index_t k_batch = 1, + BZeroPointStruct b_zero_point_struct = {}) { const auto as_grid_buf = generate_tuple( [&](auto i) { @@ -1481,7 +1485,8 @@ struct GridwiseGemm_wmma_cshuffle_v3_base a_scale_struct, b_scale_struct, num_k_block_main_loop, - num_k_block_per_scale); + num_k_block_per_scale, + b_zero_point_struct); // AIESW-32176: defaults to Empty; v1+v3 both accept // Epilogue: // - CShuffle / direct store diff --git a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp index 91bf2d5832..0e9da3e75e 100644 --- a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp +++ b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp @@ -1679,6 +1679,153 @@ struct ThreadwiseTensorSliceTransfer_v4 }); } + // AIESW-32176: scale + per-group zero point overload. + // Same as the "Fuse scale" Run() above but additionally takes a scaled_zp + // (= (zp - 8) * scale precomputed at weight load) and uses + // DequantPack8WithZp on the pk_i4 path so the dequant matches AWQ's + // (nibble - zp) * scale. All non-pk_i4 SrcData paths fall back to the + // plain dequant + scale (scaled_zp ignored — only meaningful for pk_i4). + template + __device__ void Run(const SrcDesc&, + const SrcRefToOriginDisplacement&, + const SrcBuffer& src_buf, + const DstData& scale, + const DstData& scaled_zp, + const DstDesc&, + const DstOriginIdx&, + DstBuffer& dst_buf) const + { + static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(), + "wrong! SrcDesc and DstDesc need to known at compile-time"); + static_assert( + is_same, remove_cvref_t>::value && + is_same, remove_cvref_t>::value, + "wrong! SrcBuffer or DstBuffer data type is wrong"); + static_assert(DstBuffer::IsStaticBuffer(), "wrong! DstBuffer need to be StaticBuffer"); + static_assert(is_known_at_compile_time>::value && + is_known_at_compile_time>::value, + "wrong! SrcOriginToRefDistance and DstOriginToRefDistance need to be known " + "at compile-time"); + + constexpr auto src_desc = remove_cvref_t{}; + constexpr auto dst_desc = remove_cvref_t{}; + + constexpr auto src_ref_to_origin_disp_idx = to_multi_index(SrcRefToOriginDisplacement{}); + constexpr auto dst_origin_idx = to_multi_index(DstOriginIdx{}); + + constexpr auto src_scalar_per_access = generate_sequence_v2( + [&](auto i) constexpr { + if constexpr(i == SrcVectorDim) + return Number{}; + else + return Number<1>{}; + }, + Number{}); + + constexpr auto src_scalar_step_in_vector = generate_sequence_v2( + [&](auto i) constexpr { + if constexpr(i == SrcVectorDim) + return Number<1>{}; + else + return Number<0>{}; + }, + Number{}); + + constexpr auto access_lengths = SliceLengths{} / src_scalar_per_access; + constexpr auto dim_access_order = DimAccessOrder{}; + constexpr auto ordered_access_lengths = + container_reorder_given_new2old(access_lengths, dim_access_order); + + static_ford{}([&](auto ordered_access_idx) { + constexpr auto data_to_origin_disp_idx = + ordered_access_idx.ReorderGivenOld2New(dim_access_order) * src_scalar_per_access; + + constexpr auto src_ref_to_data_disp_idx = + src_ref_to_origin_disp_idx + data_to_origin_disp_idx; + + constexpr auto src_ref_to_data_disp_coord_step = + make_tensor_coordinate_step(src_desc, src_ref_to_data_disp_idx); + + auto src_data_coord = src_ref_coord_; + move_tensor_coordinate(src_desc, src_data_coord, src_ref_to_data_disp_coord_step); + + vector_type_maker_t src_tmp_vector; + using src_vector_t = typename decltype(src_tmp_vector)::type; + + const bool is_src_valid = coordinate_has_valid_offset_assuming_visible_index_is_valid( + src_desc, src_data_coord); + + if constexpr(SrcBuffer::IsDynamicBuffer()) + { + src_tmp_vector.template AsType()(Number<0>{}) = + src_buf.template Get(src_data_coord.GetOffset() / PackedSize, + is_src_valid); + } + else if constexpr(SrcBuffer::IsStaticBuffer()) + { + static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { + constexpr index_t src_offset = src_desc.CalculateOffset( + src_ref_to_origin_disp_idx + data_to_origin_disp_idx + + i * src_scalar_step_in_vector); + src_tmp_vector.template AsType()(i) = src_buf[Number{}]; + }); + } + + if constexpr(is_same, pk_i4_t>::value) + { + vector_type_maker_t dst_tmp_vector; + vector_type scale_vector; + scale_vector.template AsType()(Number<0>{}) = scale; + scale_vector.template AsType()(Number<1>{}) = scale; + vector_type scaled_zp_vector; + scaled_zp_vector.template AsType()(Number<0>{}) = scaled_zp; + scaled_zp_vector.template AsType()(Number<1>{}) = scaled_zp; + + constexpr index_t pack_size = 8; + static_assert(SrcScalarPerVector % pack_size == 0, ""); + + using src_v_t = typename vector_type_maker_t::type; + using dst_v_t = typename vector_type_maker_t::type; + using scale_v_t = typename vector_type_maker_t::type; + + static_for<0, SrcScalarPerVector / pack_size, 1>{}([&](auto i) { + ck::tensor_operation::element_wise::DequantPack8WithZp{}( + dst_tmp_vector.template AsType()(i), + src_tmp_vector.template AsType()[i], + scale_vector.template AsType()[Number<0>{}], + scaled_zp_vector.template AsType()[Number<0>{}]); + }); + + static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { + constexpr index_t dst_offset = dst_desc.CalculateOffset( + dst_origin_idx + data_to_origin_disp_idx + i * src_scalar_step_in_vector); + dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; + }); + } + else + { + // Non-pk_i4 SrcData: scaled_zp has no defined meaning here; + // fall back to plain type_convert + scale. Asserts on dtype + // mismatch are deliberately omitted to keep this overload usable + // wherever the scale-only overload is. + vector_type_maker_t dst_tmp_vector; + static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { + dst_tmp_vector.template AsType()(i) = + type_convert(src_tmp_vector.template AsType()[i]) * scale; + }); + + static_for<0, SrcScalarPerVector, 1>{}([&](auto i) { + constexpr index_t dst_offset = dst_desc.CalculateOffset( + dst_origin_idx + data_to_origin_disp_idx + i * src_scalar_step_in_vector); + dst_buf(Number{}) = dst_tmp_vector.template AsType()[i]; + }); + } + }); + } + template __device__ void MoveSrcSliceWindow(const SrcDesc&, const SrcSliceMoveStepIdx& src_slice_move_step_idx)