ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-15 03:30:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

[CK] Add per-group zero-point support to wmma_cshuffle_v3_b_scale (asymmetric int4)

Browse Source 
Adds an optional asymmetric int4 dequant path to the wmma_cshuffle_v3
b_scale GEMM device-op. The existing path computes (nibble - 8) * scale
(symmetric uint4b8); the new path computes (nibble - zp) * scale per
group (matches AWQ / per-group-zero-point checkpoints). Implemented
via the algebraic identity:

  (nibble - zp) * scale  ==  (nibble - 8) * scale  -  (zp - 8) * scale

Callers precompute scaled_zp = (zp - 8) * scale once at weight load and
pass the [N, K/G] tensor alongside the existing scale tensor. The hot
loop pays one extra fp16 vector subtract per dequant pack.

Changes are additive and gated on optional `p_b_zero_point` being non-null
at the device-op level (defaults to nullptr) and `BZeroPointStruct = Empty`
at the gridwise/blockwise level (defaults to Empty). Symmetric callers see
no signature change and compile bit-identically.

Files:

* element/unary_element_wise_operation.hpp
    + `i4_to_half4_zp_scale()`: pk_i4 -> half4 dequant with per-group
      scaled_zp subtract (mirrors `i4_to_half4_scale`).
    + `DequantPack8WithZp`: element-op wrapping the above; mirrors
      `DequantPack8`. `is_pack8_invocable = true`.

* thread/threadwise_tensor_slice_transfer.hpp
    + `ThreadwiseTensorSliceTransfer_v4::Run()` overload that takes
      `scale + scaled_zp` instead of just `scale` and dispatches to
      `DequantPack8WithZp` on the pk_i4 path. Non-pk_i4 SrcData falls
      back to the plain dequant path with scaled_zp ignored.

* block/blockwise_gemm_pipeline_wmmaops_v1.hpp
    + `BZeroPointStruct = Empty` defaulted template arg on Run().
    + `if constexpr(HasBZp)` branch picks the new threadwise Run() overload
      and runs an extra GlobalLoad cycle for the zero-point struct,
      synchronized with the existing BScaleStruct cadence.

* block/blockwise_gemm_pipeline_wmmaops_v3.hpp
    + `BZeroPointStruct = Empty` defaulted template arg on Run() for
      signature compatibility with v1. v3 (Intrawave) does not implement
      the asymmetric path; passing a non-Empty BZeroPointStruct would
      compile but silently ignore it. Caller is responsible for routing
      asymmetric configs to v1 (Interwave).

* grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp
    + Forward `b_zero_point_struct` (defaulted Empty) from Run() to the
      blockwise pipeline's Run().

* grid/gridwise_gemm_wmma_cshuffle_v3_ab_scale.hpp
    + `p_b_zero_point_grid` field on Argument (defaulted nullptr).
    + Run() builds a BZeroPoint struct from `p_b_zero_point_grid` when
      non-null (mirroring the BScale struct construction) and threads it
      through to Base::Run; symmetric path stays bit-identical when
      p_b_zero_point_grid is nullptr.

* device/impl/device_gemm_wmma_cshuffle_v3_b_scale.hpp
    + Optional `p_b_zero_point = nullptr` trailing arg on `MakeArgument`.
      Stored on the Argument and propagated through the Run path. Existing
      symmetric callers see no signature change.

Backward compatibility verified: with p_b_zero_point=nullptr / BZeroPointStruct=Empty
the gridwise + blockwise dispatch produces bit-identical results to the
symmetric pre-change path. Tested via the standalone
example_gemm_wmma_fp16_pk_i4_v3_b_scale runner (CPU-reference verify
mode) on M=2048 N=19456 K=2560 (gate_up_proj prefill shape) and on
M=2048 across qkv (N=6144), o_proj (N=2560 K=4096), and down_proj
(N=2560 K=9728) on Strix Halo (gfx1151) — all pass.

Use case: vLLM W4A16 prefill GEMM dispatch on Strix Halo (gfx1151)
for AWQ-quantized models (e.g. Qwen/Qwen3-4B-AWQ).
JIRA: AIESW-32176.
This commit is contained in:
Marcus Rosen
2026-05-14 07:07:58 -07:00
parent fe2e29fa68
commit fe31dbe57e
7 changed files with 391 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp
View File

@@ -719,6 +719,11 @@ struct BlockwiseGemmWmmaops_pipeline_v1<BlockGemmPipelineScheduler::Interwave,
static constexpr index_t KRepeatPerNumScaleKBlock = 1;
};
// AIESW-32176: BZeroPointStruct (defaulted to Empty) carries per-group
// scaled_zp = (zp - 8) * scale precomputed at weight load. When non-Empty
// the dispatch picks the new threadwise Run overload that subtracts
// scaled_zp inside the dequant; symmetric callers (BZeroPointStruct=Empty
// by default) compile bit-identically to the original.
template <bool HasMainLoop,
TailNumber TailNum,
typename AGridDesc,
@@ -736,6 +741,7 @@ struct BlockwiseGemmWmmaops_pipeline_v1<BlockGemmPipelineScheduler::Interwave,
typename CThreadBuffer,
typename AScaleStruct,
typename BScaleStruct,
typename BZeroPointStruct = Empty,
typename enable_if<ck::is_same_v<AScaleStruct, Empty>, bool>::type = false>
__device__ void Run(const AGridDesc& a_grid_desc,
const ABlockDesc& a_block_desc,
@@ -753,8 +759,11 @@ struct BlockwiseGemmWmmaops_pipeline_v1<BlockGemmPipelineScheduler::Interwave,
AScaleStruct&,
BScaleStruct& b_scale_struct,
index_t num_loop,
index_t num_loop_per_scale) const
index_t num_loop_per_scale,
BZeroPointStruct b_zero_point_struct = {}) const
{
constexpr bool HasBZp = !ck::is_same_v<BZeroPointStruct, Empty>;
constexpr index_t KPerWaveBlock = wmma_gemm.GetKPerWaveBlk();
auto a_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, ComputeTypeA>(
@@ -771,6 +780,10 @@ struct BlockwiseGemmWmmaops_pipeline_v1<BlockGemmPipelineScheduler::Interwave,
// Scales global load
b_scale_struct.template GlobalLoad<0>(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<BlockGemmPipelineScheduler::Interwave,
b_thread_buf);
});
}
else
else if constexpr(!HasBZp)
{
static_for<0, NRepeat, 1>{}([&](auto n0) {
b_thread_copy_.Run(
@@ -817,6 +830,30 @@ struct BlockwiseGemmWmmaops_pipeline_v1<BlockGemmPipelineScheduler::Interwave,
b_thread_buf);
});
}
else
{
// AIESW-32176: scale + scaled_zp dispatch. The
// BZeroPointStruct mirrors BScaleStruct exactly (same
// [num_scale_k_block, num_scale_krepeat] shape and same
// GlobalLoad cadence) — they're built from the same
// BScale template type in the gridwise.
static_for<0, NRepeat, 1>{}([&](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<n0 * BZeroPointStruct::num_scale_k_block +
(k0_offset + k0_inner) /
BZeroPointStruct::num_scale_krepeat>{}],
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<BlockGemmPipelineScheduler::Interwave,
blockwise_gemm_func();
b_scale_struct.template GlobalLoad<0>((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<BScaleStruct, Empty>::value == false)
{
block_sync_lds();

11
include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp
View File

@@ -336,6 +336,10 @@ struct BlockwiseGemmWmmaops_pipeline_v3<BlockGemmPipelineScheduler::Intrawave,
});
}
// AIESW-32176: BZeroPointStruct is accepted as a defaulted-Empty trailing
// arg for signature compatibility with the v1 pipeline. v3 doesn't
// implement asymmetric dequant; only Empty is supported here. If a real
// BZeroPointStruct is passed, this static_assert fires.
template <bool HasMainLoop,
TailNumber TailNum,
typename AGridDesc,
@@ -353,6 +357,7 @@ struct BlockwiseGemmWmmaops_pipeline_v3<BlockGemmPipelineScheduler::Intrawave,
typename CThreadBuffer,
typename AScaleStruct,
typename BScaleStruct,
typename BZeroPointStruct = Empty,
typename enable_if<ck::is_same_v<AScaleStruct, Empty>, bool>::type = false>
__device__ void Run(const AGridDesc& a_grid_desc,
const ABlockDesc& a_block_desc,
@@ -370,8 +375,12 @@ struct BlockwiseGemmWmmaops_pipeline_v3<BlockGemmPipelineScheduler::Intrawave,
AScaleStruct&,
BScaleStruct& b_scale_struct,
index_t num_loop,
index_t num_loop_per_scale) const
index_t num_loop_per_scale,
BZeroPointStruct = {}) const
{
// AIESW-32176: zero-point arg is ignored by v3 (asymmetric is only
// implemented in the v1 Interwave path). Caller is responsible for
// routing asymmetric configs to v1.
__builtin_amdgcn_sched_barrier(0);
constexpr index_t KPerWaveBlock = wmma_gemm.GetKPerWaveBlk();

10
include/ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3_b_scale.hpp
View File

@@ -182,6 +182,9 @@ struct DeviceGemm_BScale_Wmma_CShuffleV3 : public DeviceGemmV2BScale<ALayout,
bool GetPermuteB() override { return PermuteB; }
// AIESW-32176: optional p_b_zero_point trailing arg (defaults to nullptr).
// When non-null the gridwise/blockwise pipelines pick up the asymmetric
// dequant path. Existing symmetric callers see no signature change.
static auto MakeArgument(const ADataType* p_a,
const BDataType* p_b,
CDataType* p_c,
@@ -196,7 +199,8 @@ struct DeviceGemm_BScale_Wmma_CShuffleV3 : public DeviceGemmV2BScale<ALayout,
index_t KBatch,
AElementwiseOperation a_element_op,
BElementwiseOperation b_element_op,
CElementwiseOperation cde_element_op)
CElementwiseOperation cde_element_op,
const BScaleDataType* p_b_zero_point = nullptr)
{
return Argument{std::array<const void*, 1>{p_a},
std::array<const void*, 1>{p_b},
@@ -216,7 +220,9 @@ struct DeviceGemm_BScale_Wmma_CShuffleV3 : public DeviceGemmV2BScale<ALayout,
KBatch,
a_element_op,
b_element_op,
cde_element_op};
cde_element_op,
/*is_reduce=*/false,
p_b_zero_point};
}
static auto MakeInvoker() { return Invoker{}; }

82
include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp
View File

@@ -75,6 +75,43 @@ __device__ inline half4_t i4_to_half4_scale(int q, const ck::half2_t& scale)
return res.template AsType<half4_t>()[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<half_t, 4> res;
res.template AsType<half2_t>()(Number<0>{}) = bit_cast<half2_t>(lo) + bit_cast<half2_t>(SUB);
res.template AsType<half2_t>()(Number<1>{}) =
bit_cast<half2_t>(hi) * bit_cast<half2_t>(MUL) + bit_cast<half2_t>(ADD);
res.template AsType<half2_t>()(Number<0>{}) =
res.template AsType<half2_t>()(Number<0>{}) * scale - scaled_zp;
res.template AsType<half2_t>()(Number<1>{}) =
res.template AsType<half2_t>()(Number<1>{}) * scale - scaled_zp;
return res.template AsType<half4_t>()[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 <typename Y, typename X, typename S, typename ZP>
__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<half_t, 8> result;
result.template AsType<half4_t>()(Number<0>{}) =
i4_to_half4_zp_scale(bit_cast<int>(x), s, scaled_zp);
result.template AsType<half4_t>()(Number<1>{}) =
i4_to_half4_zp_scale(bit_cast<int>(x) >> 8, s, scaled_zp);
y = result.template AsType<half8_t>()[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<half_t, 8> dst;
vector_type<pk_i4_t, 4> src{x};
static_for<0, 4, 1>{}([&](auto i) {
auto v = type_convert<half2_t>(src.template AsType<pk_i4_t>()[i]);
dst.template AsType<half2_t>()(i) = v * s - scaled_zp;
});
y = dst.template AsType<half8_t>()[Number<0>{}];
#endif
}
constexpr const static bool is_pack8_invocable = true;
};
struct PassThroughPack2
{
static constexpr const char* name = "PassThroughPack2";

124
include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_ab_scale.hpp
View File

@@ -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 <bool HasMainKBlockLoop,
InMemoryDataOperationEnum EGlobalMemoryDataOperation,
TailNumber TailNum,
@@ -723,8 +731,9 @@ struct GridwiseGemm_wmma_cshuffle_v3_ab_scale
BElementwiseOperation b_element_op,
CDEElementwiseOperation cde_element_op,
EpilogueArgument& epilogue_args,
const index_t A_k_id = 0,
const index_t B_k_id = 0)
const index_t A_k_id = 0,
const index_t B_k_id = 0,
const BScaleType* p_b_zero_point_grid = nullptr)
{
const auto as_grid_desc_ak0_m_ak1 = MakeAsGridDescriptor_AK0_M_AK1(
problem.M, problem.MPadded, problem.K, problem.KPadded, problem.StrideAs, problem.AK0);
@@ -765,37 +774,82 @@ struct GridwiseGemm_wmma_cshuffle_v3_ab_scale
// BScale struct
auto b_scale_struct = MakeBScale<1>(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<decltype(as_grid_desc_ak0_m_ak1),
decltype(bs_grid_desc_bk0_n_bk1),
decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock),
decltype(e_grid_desc_mblock_mperblock_nblock_nperblock),
decltype(a_scale_struct),
decltype(b_scale_struct),
decltype(epilogue_args),
HasMainKBlockLoop,
EGlobalMemoryDataOperation,
TailNum>(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<decltype(as_grid_desc_ak0_m_ak1),
decltype(bs_grid_desc_bk0_n_bk1),
decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock),
decltype(e_grid_desc_mblock_mperblock_nblock_nperblock),
decltype(a_scale_struct),
decltype(b_scale_struct),
decltype(epilogue_args),
HasMainKBlockLoop,
EGlobalMemoryDataOperation,
TailNum>(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<decltype(as_grid_desc_ak0_m_ak1),
decltype(bs_grid_desc_bk0_n_bk1),
decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock),
decltype(e_grid_desc_mblock_mperblock_nblock_nperblock),
decltype(a_scale_struct),
decltype(b_scale_struct),
decltype(epilogue_args),
HasMainKBlockLoop,
EGlobalMemoryDataOperation,
TailNum,
decltype(b_zero_point_struct)>(
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);
}
};

15
include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp
View File

@@ -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 <typename AGridDesc_AK0_M_K1,
typename BGridDesc_BK0_N_K1,
typename DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock,
@@ -1373,7 +1375,8 @@ struct GridwiseGemm_wmma_cshuffle_v3_base
typename EpilogueArgument,
bool HasMainKBlockLoop,
InMemoryDataOperationEnum EGlobalMemoryDataOperation,
TailNumber TailNum = TailNumber::Odd>
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

147
include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp
View File

@@ -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 <typename SrcRefToOriginDisplacement,
typename DstOriginIdx,
typename SrcBuffer,
typename DstBuffer>
__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<typename SrcBuffer::type>, remove_cvref_t<SrcData>>::value &&
is_same<remove_cvref_t<typename DstBuffer::type>, remove_cvref_t<DstData>>::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<remove_cvref_t<SrcRefToOriginDisplacement>>::value &&
is_known_at_compile_time<remove_cvref_t<DstOriginIdx>>::value,
"wrong! SrcOriginToRefDistance and DstOriginToRefDistance need to be known "
"at compile-time");
constexpr auto src_desc = remove_cvref_t<SrcDesc>{};
constexpr auto dst_desc = remove_cvref_t<DstDesc>{};
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<SrcScalarPerVector>{};
else
return Number<1>{};
},
Number<nDim>{});
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<nDim>{});
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<decltype(ordered_access_lengths)>{}([&](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<SrcData, SrcScalarPerVector / PackedSize> 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<src_vector_t>()(Number<0>{}) =
src_buf.template Get<src_vector_t>(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<SrcData>()(i) = src_buf[Number<src_offset>{}];
});
}
if constexpr(is_same<remove_cvref_t<SrcData>, pk_i4_t>::value)
{
vector_type_maker_t<DstData, SrcScalarPerVector> dst_tmp_vector;
vector_type<DstData, 2> scale_vector;
scale_vector.template AsType<DstData>()(Number<0>{}) = scale;
scale_vector.template AsType<DstData>()(Number<1>{}) = scale;
vector_type<DstData, 2> scaled_zp_vector;
scaled_zp_vector.template AsType<DstData>()(Number<0>{}) = scaled_zp;
scaled_zp_vector.template AsType<DstData>()(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<SrcData, pack_size / PackedSize>::type;
using dst_v_t = typename vector_type_maker_t<DstData, pack_size>::type;
using scale_v_t = typename vector_type_maker_t<DstData, 2>::type;
static_for<0, SrcScalarPerVector / pack_size, 1>{}([&](auto i) {
ck::tensor_operation::element_wise::DequantPack8WithZp{}(
dst_tmp_vector.template AsType<dst_v_t>()(i),
src_tmp_vector.template AsType<src_v_t>()[i],
scale_vector.template AsType<scale_v_t>()[Number<0>{}],
scaled_zp_vector.template AsType<scale_v_t>()[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_offset>{}) = dst_tmp_vector.template AsType<DstData>()[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<DstData, SrcScalarPerVector> dst_tmp_vector;
static_for<0, SrcScalarPerVector, 1>{}([&](auto i) {
dst_tmp_vector.template AsType<DstData>()(i) =
type_convert<DstData>(src_tmp_vector.template AsType<SrcData>()[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_offset>{}) = dst_tmp_vector.template AsType<DstData>()[i];
});
}
});
}
template <typename SrcSliceMoveStepIdx>
__device__ void MoveSrcSliceWindow(const SrcDesc&,
const SrcSliceMoveStepIdx& src_slice_move_step_idx)