Skip to content

Commit 22a99f9

Browse files
tywuAMDassistant-librarian[bot]
tywuAMD
authored and
assistant-librarian[bot]
committed
[rocm-libraries] ROCm/rocm-libraries#7677 (commit 308af93)
[CK_Tile] Add scale16 Support for F4 WMMA in CK_Tile ## Motivation This PR adds CK Tile support for the scale16 F4 WMMA path on gfx1250 and improves warp GEMM unit test coverage/structure for gfx1250-specific cases. ## Technical Details - Scale16 support in warp GEMM dispatch and WMMA trait plumbing: added IsScale16 plumbing to warp GEMM dispatcher path - Warp GEMM test restructuring for gfx1250: added Warp GEMM gfx1250 coverage to verify all F4 WMMA paths ## Test Plan Run ./test_ck_tile_wg_32x16x128_fp4. ## Test Result ``` ./test_ck_tile_wg_32x16x128_fp4 [----------] Global test environment tear-down [==========] 3 tests from 1 test suite ran. (1751 ms total) [ PASSED ] 3 tests. ``` ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
1 parent 8d97265 commit 22a99f9

11 files changed

Lines changed: 508 additions & 235 deletions

‎include/ck_tile/core/numeric/mxfp_scale.hpp‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,101 @@ struct Packed4Scale
103103
}
104104
};
105105

106+
template <typename ScaleType>
107+
struct Packed8Scale
108+
{
109+
using scale_type = ScaleType;
110+
using raw_type = uint64_t;
111+
using raw_scale_type = typename ScaleType::raw_type;
112+
113+
static constexpr int num_pack = 8;
114+
union
115+
{
116+
raw_type data_;
117+
raw_scale_type scales_[num_pack]; // Direct byte/element access
118+
};
119+
120+
// Constructors
121+
CK_TILE_HOST_DEVICE constexpr Packed8Scale() = default;
122+
CK_TILE_HOST_DEVICE constexpr Packed8Scale(raw_type val) : data_(val) {}
123+
CK_TILE_HOST_DEVICE constexpr Packed8Scale(
124+
float s0, float s1, float s2, float s3, float s4, float s5, float s6, float s7)
125+
{
126+
set_scales_from_float(s0, s1, s2, s3, s4, s5, s6, s7);
127+
}
128+
129+
CK_TILE_HOST_DEVICE constexpr Packed8Scale(ScaleType s0,
130+
ScaleType s1,
131+
ScaleType s2,
132+
ScaleType s3,
133+
ScaleType s4,
134+
ScaleType s5,
135+
ScaleType s6,
136+
ScaleType s7)
137+
{
138+
set_scales(s0, s1, s2, s3, s4, s5, s6, s7);
139+
}
140+
141+
CK_TILE_HOST_DEVICE constexpr void set_scales_from_float(
142+
float s0, float s1, float s2, float s3, float s4, float s5, float s6, float s7)
143+
{
144+
set_scales(ScaleType(s0),
145+
ScaleType(s1),
146+
ScaleType(s2),
147+
ScaleType(s3),
148+
ScaleType(s4),
149+
ScaleType(s5),
150+
ScaleType(s6),
151+
ScaleType(s7));
152+
}
153+
154+
CK_TILE_HOST_DEVICE constexpr void set_scales(ScaleType s0,
155+
ScaleType s1,
156+
ScaleType s2,
157+
ScaleType s3,
158+
ScaleType s4,
159+
ScaleType s5,
160+
ScaleType s6,
161+
ScaleType s7)
162+
{
163+
data_ = 0;
164+
pack_scale(s0, 7);
165+
pack_scale(s1, 6);
166+
pack_scale(s2, 5);
167+
pack_scale(s3, 4);
168+
pack_scale(s4, 3);
169+
pack_scale(s5, 2);
170+
pack_scale(s6, 1);
171+
pack_scale(s7, 0);
172+
}
173+
174+
CK_TILE_HOST_DEVICE constexpr operator raw_type() const { return data_; }
175+
CK_TILE_HOST_DEVICE constexpr raw_type& data() [[clang::lifetimebound]] { return data_; }
176+
CK_TILE_HOST_DEVICE constexpr raw_type data() const { return data_; }
177+
178+
CK_TILE_HOST_DEVICE constexpr float unpack_to_float(int i) const
179+
{
180+
return static_cast<float>(unpack_scale(i));
181+
}
182+
183+
CK_TILE_HOST_DEVICE constexpr ScaleType unpack_scale(int i) const
184+
{
185+
return ScaleType(scales_[i]);
186+
}
187+
188+
CK_TILE_HOST_DEVICE constexpr void pack_from_float(float scale, int i)
189+
{
190+
pack_scale(ScaleType(scale), i);
191+
}
192+
193+
CK_TILE_HOST_DEVICE constexpr void pack_scale(ScaleType scale, int i)
194+
{
195+
scales_[i] = scale.get();
196+
}
197+
};
198+
106199
// Type alias for e8m0_t scales
107200
using Packed4Scale_E8M0 = Packed4Scale<e8m0_t>;
201+
using Packed8Scale_E8M0 = Packed8Scale<e8m0_t>;
108202

109203
} // namespace ck_tile

‎include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -234,12 +234,12 @@ struct WarpGemmAttributeWmma
234234
}
235235

236236
// c_vec += a_vec * b_vec
237-
template <typename... Params>
237+
template <typename... Params, typename AScaleType, typename BScaleType>
238238
CK_TILE_DEVICE void operator()(CVecType& c_vec,
239239
const AVecType& a_vec,
240-
const int32_t& a_scale,
240+
const AScaleType& a_scale,
241241
const BVecType& b_vec,
242-
const int32_t& b_scale) const
242+
const BScaleType& b_scale) const
243243
{
244244
if constexpr(kTransC)
245245
{
@@ -253,11 +253,11 @@ struct WarpGemmAttributeWmma
253253
}
254254

255255
// c_vec = a_vec * b_vec
256-
template <typename... Params>
256+
template <typename... Params, typename AScaleType, typename BScaleType>
257257
CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec,
258-
const int32_t& a_scale,
258+
const AScaleType& a_scale,
259259
const BVecType& b_vec,
260-
const int32_t& b_scale) const
260+
const BScaleType& b_scale) const
261261
{
262262
if constexpr(kTransC)
263263
{

‎include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ template <typename Arch,
1919
typename MXTypeEnable = void>
2020
struct WmmaTraits;
2121

22+
// Tag used to select scale16 WMMA traits specializations.
23+
struct WmmaScale16Tag
24+
{
25+
};
26+
2227
// Generic WMMA implementation using traits
2328
template <typename Traits>
2429
struct WarpGemmAttributeWmmaImpl
@@ -88,22 +93,22 @@ struct WarpGemmAttributeWmmaImpl
8893
Traits::template wmma_intrinsic<Params...>(a_vec, b_vec, CVecType{0.f}));
8994
}
9095

91-
template <typename... Params>
96+
template <typename... Params, typename AScaleType, typename BScaleType>
9297
CK_TILE_DEVICE void operator()(CVecType& c_vec,
9398
const AVecType& a_vec,
94-
const int32_t& a_scale,
99+
const AScaleType& a_scale,
95100
const BVecType& b_vec,
96-
const int32_t& b_scale) const
101+
const BScaleType& b_scale) const
97102
{
98103
c_vec = Traits::template wmma_intrinsic<Params...>(a_vec, a_scale, b_vec, b_scale, c_vec);
99104
}
100105

101106
// c_vec = a_vec * b_vec
102-
template <typename... Params>
107+
template <typename... Params, typename AScaleType, typename BScaleType>
103108
CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec,
104-
const int32_t& a_scale,
109+
const AScaleType& a_scale,
105110
const BVecType& b_vec,
106-
const int32_t& b_scale) const
111+
const BScaleType& b_scale) const
107112
{
108113
return bit_cast<CVecType>(Traits::template wmma_intrinsic<Params...>(
109114
a_vec, a_scale, b_vec, b_scale, CVecType{0.f}));
@@ -177,6 +182,9 @@ using WarpGemmAttributeWmmaImpl_f32_32x16x128_f4 =
177182
using WarpGemmAttributeWmmaImpl_f32_32x32x128_f4 =
178183
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 32, 128>>;
179184

185+
using WarpGemmAttributeWmmaImpl_f32_32x32x128_f4_scale16 = WarpGemmAttributeWmmaImpl<
186+
WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 32, 128, WmmaScale16Tag>>;
187+
180188
using WarpGemmAttributeWmmaImpl_f16_16x16x64_f8_f8 =
181189
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx125_t, fp8_t, fp8_t, fp16_t, 16, 16, 64>>;
182190

‎include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp‎

Lines changed: 55 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@
66
#include "warp_gemm_attribute_wmma_impl_base_traits.hpp"
77
#include "warp_gemm_params.hpp"
88
namespace ck_tile {
9+
10+
struct WmmaScale16Tag;
11+
912
// int8 specialization - GFX11
1013
template <>
1114
struct WmmaTraits<gfx11_t, int8_t, int8_t, int32_t, 16, 16, 16>
@@ -528,17 +531,18 @@ struct WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 16, 128>
528531
}
529532
};
530533

531-
template <>
532-
struct WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 32, 128>
534+
template <bool IsScale16>
535+
struct WmmaTraitsGfx125PkFp4F32_32x32x128
533536
: WmmaTraitsBase<gfx12_t, pk_fp4_t, pk_fp4_t, float, 128, false, 32, 32>
534537
{
535-
using ArchType = gfx125_t;
538+
using ArchType = gfx125_t;
539+
using ScaleType = std::conditional_t<IsScale16, int64_t, int32_t>;
536540

537541
template <typename... Params>
538542
CK_TILE_DEVICE static CVecType wmma_intrinsic(const AVecType& a_vec,
539-
const int32_t& a_scale,
543+
const ScaleType& a_scale,
540544
const BVecType& b_vec,
541-
const int32_t& b_scale,
545+
const ScaleType& b_scale,
542546
const CVecType& c_vec)
543547
{
544548
#ifdef __gfx125__
@@ -569,19 +573,38 @@ struct WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 32, 128>
569573
const auto& b_slice = b_buffer.template get_as<BSliceType>()[n];
570574
auto& c_slice = c_result.template get_as<CSliceType>()[n];
571575

572-
c_slice = __builtin_amdgcn_wmma_scale_f32_32x16x128_f4(
573-
bit_cast<int32x16_t>(a_slice),
574-
bit_cast<int32x8_t>(b_slice),
575-
0,
576-
c_slice,
577-
1, // OPSEL[0] - fixed to 1 for F4
578-
P::scale_a, // OPSEL_HI[0] - scale data type for A
579-
a_scale,
580-
n.value, // OPSEL[1] - select B scale (iterates over N blocks)
581-
P::scale_b, // OPSEL_HI[1] - scale data type for B
582-
b_scale,
583-
0, // NEG
584-
0); // NEG_HI
576+
if constexpr(IsScale16)
577+
{
578+
c_slice = __builtin_amdgcn_wmma_scale16_f32_32x16x128_f4(
579+
bit_cast<int32x16_t>(a_slice),
580+
bit_cast<int32x8_t>(b_slice),
581+
0,
582+
c_slice,
583+
1, // OPSEL[0] - fixed to 1 for F4
584+
P::scale_a, // OPSEL_HI[0] - scale data type for A
585+
a_scale,
586+
n.value, // OPSEL[1] - select B scale (iterates over N blocks)
587+
P::scale_b, // OPSEL_HI[1] - scale data type for B
588+
b_scale,
589+
0, // NEG
590+
0); // NEG_HI
591+
}
592+
else
593+
{
594+
c_slice = __builtin_amdgcn_wmma_scale_f32_32x16x128_f4(
595+
bit_cast<int32x16_t>(a_slice),
596+
bit_cast<int32x8_t>(b_slice),
597+
0,
598+
c_slice,
599+
1, // OPSEL[0] - fixed to 1 for F4
600+
P::scale_a, // OPSEL_HI[0] - scale data type for A
601+
a_scale,
602+
n.value, // OPSEL[1] - select B scale (iterates over N blocks)
603+
P::scale_b, // OPSEL_HI[1] - scale data type for B
604+
b_scale,
605+
0, // NEG
606+
0); // NEG_HI
607+
}
585608
});
586609

587610
return bit_cast<CVecType>(c_result);
@@ -602,7 +625,8 @@ struct WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 32, 128>
602625
#ifdef __gfx125__
603626
// Pass default scale values 1.0f
604627
Packed4Scale_E8M0 pkscale(1.0f, 1.0f, 1.0f, 1.0f);
605-
return wmma_intrinsic(a_vec, pkscale, b_vec, pkscale, c_vec);
628+
const auto default_scale = static_cast<ScaleType>(pkscale);
629+
return wmma_intrinsic(a_vec, default_scale, b_vec, default_scale, c_vec);
606630
#else
607631
ck_tile::ignore = a_vec;
608632
ck_tile::ignore = b_vec;
@@ -612,6 +636,18 @@ struct WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 32, 128>
612636
}
613637
};
614638

639+
template <>
640+
struct WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 32, 128>
641+
: WmmaTraitsGfx125PkFp4F32_32x32x128<false>
642+
{
643+
};
644+
645+
template <>
646+
struct WmmaTraits<gfx125_t, pk_fp4_t, pk_fp4_t, float, 32, 32, 128, WmmaScale16Tag>
647+
: WmmaTraitsGfx125PkFp4F32_32x32x128<true>
648+
{
649+
};
650+
615651
// f8f6f4 specialization - GFX125
616652
enum F8F6F4OpDataTypeEnum
617653
{

‎include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp‎

Lines changed: 35 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ template <typename AType,
3333
bool UseStructuredSparsity = false,
3434
WGAttrNumAccessEnum AttrNumAccessA = ESingle,
3535
WGAttrNumAccessEnum AttrNumAccessB = AttrNumAccessA,
36+
bool IsScale16 = false,
3637
typename Enable = void>
3738
struct Dispatcher;
3839

@@ -178,10 +179,10 @@ template<> struct Dispatcher<bf8_t, bf8_t, float, 32, 32, 16, true> { using Ty
178179

179180
#if !defined(__gfx125__)
180181
// scale mfma based f8f6f4
181-
template<typename A, typename B, WGAttrNumAccessEnum I>
182-
struct Dispatcher<A, B, float, 16, 16, 128, false, false, false, I, I, std::enable_if_t<I != EDefault>> { using Type = WarpGemmMfma_f32_16x16x128_f8f6f4<A, B, I>; };
183-
template<typename A, typename B, WGAttrNumAccessEnum I>
184-
struct Dispatcher<A, B, float, 16, 16, 128, true, false, false, I, I, std::enable_if_t<I != EDefault>> { using Type = WarpGemmMfma_f32_16x16x128_f8f6f4_CTransposed<A, B, I>; };
182+
template<typename A, typename B, WGAttrNumAccessEnum I, bool IsScale16>
183+
struct Dispatcher<A, B, float, 16, 16, 128, false, false, false, I, I, IsScale16, std::enable_if_t<I != EDefault>> { using Type = WarpGemmMfma_f32_16x16x128_f8f6f4<A, B, I>; };
184+
template<typename A, typename B, WGAttrNumAccessEnum I, bool IsScale16>
185+
struct Dispatcher<A, B, float, 16, 16, 128, true, false, false, I, I, IsScale16, std::enable_if_t<I != EDefault>> { using Type = WarpGemmMfma_f32_16x16x128_f8f6f4_CTransposed<A, B, I>; };
185186
#endif
186187

187188
template<> struct Dispatcher<fp8_t, fp8_t, float, 32, 32, 64, false> { using Type = WarpGemmMfma_f32_32x32x64_fp8_fp8<>; };
@@ -224,7 +225,7 @@ template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<f
224225
template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<bf8_t, fp8_t, float, 16, 16, 64, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type = WarpGemmWmma_f32_16x16x64_bf8_f8<TransposeC, AttrNumAccess>; };
225226

226227
template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<pk_fp4_t, pk_fp4_t, float, 32, 16, 128, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type = WarpGemmWmma_f32_32x16x128_f4<TransposeC, AttrNumAccess>; };
227-
template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<pk_fp4_t, pk_fp4_t, float, 32, 32, 128, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type = WarpGemmWmma_f32_32x32x128_f4<TransposeC, AttrNumAccess>; };
228+
template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess, bool IsScale16> struct Dispatcher<pk_fp4_t, pk_fp4_t, float, 32, 32, 128, TransposeC, false, false, AttrNumAccess, AttrNumAccess, IsScale16> : WmmaTag { using Type = WarpGemmWmma_f32_32x32x128_f4<TransposeC, AttrNumAccess, IsScale16>; };
228229

229230
#if defined(__gfx125__)
230231
template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<fp8_t, fp8_t, float, 16, 16, 64, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type = WarpGemmWmma_f32_16x16x64_f8_f8<TransposeC, AttrNumAccess>; };
@@ -244,8 +245,27 @@ template<> struct Dispatcher<fp8_t, fp8_t, float, 16, 16, 64, true> { using Typ
244245
template<> struct Dispatcher<bf8_t, bf8_t, float, 16, 16, 64, true> { using Type = WarpGemmMfma_f32_16x16x64_bf8_bf8_CTransposed; };
245246
#endif
246247

247-
template<typename A, typename B, bool TransposeC, WGAttrNumAccessEnum AttrNumAccessA, WGAttrNumAccessEnum AttrNumAccessB>
248-
struct Dispatcher<A, B, float, 32, 32, 128, TransposeC, false, false, AttrNumAccessA, AttrNumAccessB> : WmmaTag { using Type = WarpGemmWmma_f32_32x32x128_f8f6f4<A, B, TransposeC, AttrNumAccessA, AttrNumAccessB>; };
248+
template<typename A,
249+
typename B,
250+
bool TransposeC,
251+
WGAttrNumAccessEnum AttrNumAccessA,
252+
WGAttrNumAccessEnum AttrNumAccessB,
253+
bool IsScale16>
254+
struct Dispatcher<A,
255+
B,
256+
float,
257+
32,
258+
32,
259+
128,
260+
TransposeC,
261+
false,
262+
false,
263+
AttrNumAccessA,
264+
AttrNumAccessB,
265+
IsScale16> : WmmaTag
266+
{
267+
using Type = WarpGemmWmma_f32_32x32x128_f8f6f4<A, B, TransposeC, AttrNumAccessA, AttrNumAccessB>;
268+
};
249269

250270
template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<fp8_t, fp8_t, half_t, 16, 16, 64, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type =WarpGemmWmma_f16_16x16x64_f8_f8<TransposeC, AttrNumAccess>; };
251271
template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<bf8_t, bf8_t, half_t, 16, 16, 64, TransposeC, false, false, AttrNumAccess, AttrNumAccess> : WmmaTag { using Type =WarpGemmWmma_f16_16x16x64_bf8_bf8<TransposeC, AttrNumAccess>; };
@@ -265,12 +285,12 @@ template<bool TransposeC, WGAttrNumAccessEnum AttrNumAccess> struct Dispatcher<u
265285

266286
template <typename AType, typename BType, typename AccType,
267287
index_t M, index_t N, index_t K,
268-
bool TransposeC, bool SA, bool SS>
288+
bool TransposeC, bool SA, bool SS, bool IsScale16>
269289
struct Dispatcher<AType, BType, AccType, M, N, K, TransposeC, SA, SS,
270-
EDefault, EDefault,
290+
EDefault, EDefault, IsScale16,
271291
std::enable_if_t<!std::is_base_of_v<WmmaTag,
272-
Dispatcher<AType, BType, AccType, M, N, K, TransposeC, SA, SS, ESingle, ESingle, void>>>>
273-
: Dispatcher<AType, BType, AccType, M, N, K, TransposeC, SA, SS, ESingle, ESingle, void> {};
292+
Dispatcher<AType, BType, AccType, M, N, K, TransposeC, SA, SS, ESingle, ESingle, IsScale16, void>>>>
293+
: Dispatcher<AType, BType, AccType, M, N, K, TransposeC, SA, SS, ESingle, ESingle, IsScale16, void> {};
274294

275295
// clang-format on
276296
} // namespace warp_gemm_dispatcher
@@ -286,7 +306,8 @@ template <typename AType,
286306
bool SwizzleA = false,
287307
bool UseStructuredSparsity = false,
288308
WGAttrNumAccessEnum AttrNumAccessA = WGAttrNumAccessEnum::Default,
289-
WGAttrNumAccessEnum AttrNumAccessB = AttrNumAccessA>
309+
WGAttrNumAccessEnum AttrNumAccessB = AttrNumAccessA,
310+
bool IsScale16 = false>
290311
using WarpGemmDispatcher = typename impl::warp_gemm_dispatcher::Dispatcher< //
291312
AType,
292313
BType,
@@ -298,6 +319,7 @@ using WarpGemmDispatcher = typename impl::warp_gemm_dispatcher::Dispatcher< //
298319
SwizzleA,
299320
UseStructuredSparsity,
300321
AttrNumAccessA,
301-
AttrNumAccessB>::Type;
322+
AttrNumAccessB,
323+
IsScale16>::Type;
302324

303325
} // namespace ck_tile

0 commit comments

Comments
 (0)