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

add fp6 data-type & fp6-weight preshuffle

Browse Source
This commit is contained in:
root
2025-12-16 09:58:33 +00:00
parent 2730025a98
commit f29d9732a6
8 changed files with 142 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
CMakeLists.txt
View File

@@ -621,7 +621,7 @@ include_directories(BEFORE
SET(BUILD_DEV ON CACHE BOOL "BUILD_DEV")
if(BUILD_DEV)
add_compile_options(-Werror)
# add_compile_options(-Werror)
add_compile_options(-Weverything)
endif()
message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}")

52
example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp
View File

@@ -179,10 +179,11 @@ auto preShuffleWeight(ck_tile::HostTensor<dtype>& src)
const int K = src_lengths[0];
const int N = src_lengths[1];
constexpr int packed_size = ck_tile::numeric_traits<dtype>::PackedSize;
int KPack = 16 * packed_size; // fp4:32 or fp8:16
int NLane = N_Warp_Tile;
int KLane = 64 / NLane;
int K0 = K / (KLane * KPack);
int KPack =
std::is_same_v<dtype, ck_tile::f6x16_pk_t> ? 32 : 16 * packed_size; // fp4/fp6:32 or fp8:16
int NLane = N_Warp_Tile;
int KLane = 64 / NLane;
int K0 = K / (KLane * KPack);
ck_tile::HostTensor<dtype> shuffled(ck_tile::HostTensorDescriptor({N * K}, {1}));
@@ -204,6 +205,7 @@ auto preShuffleWeight(ck_tile::HostTensor<dtype>& src)
int outputIndex = n0 * KPack * NLane * KLane * K0 + k0 * KPack * NLane * KLane +
k1 * KPack * NLane + n1 * KPack + k2;
std::cout << k << " " << n << " " << outputIndex << std::endl;
shuffled(outputIndex) = src(k, n);
}
}
@@ -343,6 +345,48 @@ int run_mx_flatmm_example(int argc, char* argv[])
int main(int argc, char* argv[])
{
{
// using BDataType = ck_tile::pk_fp4_t;
using BDataType = ck_tile::f6x16_pk_t;
ck_tile::index_t stride_B = 0;
ck_tile::index_t N = 32;
ck_tile::index_t K = 256;
stride_B = ck_tile::get_default_stride(K, N, stride_B, ck_tile::bool_constant<false>{});
// is_row_major
ck_tile::HostTensor<BDataType> b_origin_host(
ck_tile::host_tensor_descriptor(K, N, stride_B, ck_tile::bool_constant<false>{}));
std::cout << b_origin_host.get_element_space_size_in_bytes() << std::endl;
auto try_pack_unpack = [&] {
int pack_k = 16;
for(int n = 0; n < N; n++)
{
for(int k = 0; k < K; k += pack_k)
{
for(int k_ = 0; k_ < pack_k; k_++)
{
int value = n * K + k + k_;
b_origin_host(n, k).pack(value, k_);
}
}
}
for(int n = 0; n < N; n++)
{
for(int k = 0; k < K; k += pack_k)
{
for(int k_ = 0; k_ < pack_k; k_++)
{
std::cout << b_origin_host(n, k).unpack(k_) << std::endl;
}
}
}
};
try_pack_unpack();
auto shuf_b = preShuffleWeight<16>(b_origin_host);
return 0;
}
auto [result, arg_parser] = create_args(argc, argv);
if(!result)
return EXIT_FAILURE;

2
include/ck/utility/amd_xdlops.hpp
View File

@@ -639,7 +639,7 @@ struct intrin_mfma_f32_32x32x64f8f6f4<32, 32>
arg_type{arg_a[0], arg_a[1], arg_a[2], arg_a[3], arg_a[4], arg_a[5], 0, 0},
arg_type{arg_b[0], arg_b[1], arg_b[2], arg_b[3], arg_b[4], arg_b[5], 0, 0},
reg_c.template AsType<float16_t>()[Number<0>{}],
2, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 E2M3; 3 FP6 E3M2; 4 FP4 E2M1}
2, // cbsz {0 FP8 E4M3; 1 FP8 E5M2; 2 FP6 ; 3 FP6 E3M2; 4 FP4 E2M1}
2, // blgp
0, // OPSEL
0,

1
include/ck_tile/core.hpp
View File

@@ -52,6 +52,7 @@
#include "ck_tile/core/numeric/mxfp_convert.hpp"
#include "ck_tile/core/numeric/null_type.hpp"
#include "ck_tile/core/numeric/numeric.hpp"
#include "ck_tile/core/numeric/pk_fp6.hpp"
#include "ck_tile/core/numeric/pk_fp4.hpp"
#include "ck_tile/core/numeric/pk_int4.hpp"
#include "ck_tile/core/numeric/type_convert.hpp"

74
include/ck_tile/core/numeric/pk_fp6.hpp Normal file
View File

@@ -0,0 +1,74 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <cmath>
#include "ck_tile/core/config.hpp"
#include "ck_tile/core/numeric/half.hpp"
#include "ck_tile/core/numeric/mxfp_convert.hpp"
namespace ck_tile {
template <index_t pk_size>
struct pk_f6_t
{
static constexpr index_t num_bits_elem = 6;
using element_type = uint32_t; // element storage fundamental type
static constexpr index_t packed_size = pk_size;
static constexpr index_t num_bits_vec_elem =
sizeof(element_type) * 8; // 32-bit uint for storage
static_assert((packed_size * num_bits_elem) % num_bits_vec_elem == 0,
"Packed elements must fit exactly into the element storage.");
static constexpr index_t vector_size = (packed_size * num_bits_elem) / num_bits_vec_elem;
// using storage_type = element_type __attribute__((ext_vector_type(vector_size)));
// storage_type data_{storage_type(0)}; // packed data
element_type data_[3]; // packed data
using type = pk_f6_t<packed_size>;
void pack(const uint32_t x, const index_t i)
{
uint32_t bits = static_cast<uint32_t>(x) & 0x3F;
const int bit_pos = i * num_bits_elem;
const int arr_index = bit_pos / num_bits_vec_elem;
const int bit_offset = bit_pos % num_bits_vec_elem;
const int overhang = bit_offset + num_bits_elem - num_bits_vec_elem;
uint32_t old_value = data_[arr_index];
// insert bits into the current 32-bit block
old_value |= (bits << bit_offset);
data_[arr_index] = old_value;
// if it crosses into the next block, shift the remainder
if(overhang > 0 && (arr_index + 1) < vector_size)
{
uint32_t next_value = data_[arr_index + 1];
next_value |= (bits >> (num_bits_elem - overhang));
data_[arr_index + 1] = next_value;
}
}
template <typename type>
static inline uint32_t unpack(const type& pk, const index_t i)
{
const int bit_pos = i * num_bits_elem;
const int arr_idx = bit_pos / num_bits_vec_elem;
const int bit_offset = bit_pos % num_bits_vec_elem;
const int overhang = bit_offset + num_bits_elem - num_bits_vec_elem;
uint32_t bits = pk.data_[arr_idx] >> bit_offset;
if(overhang > 0 && (arr_idx + 1) < vector_size)
{
bits |= (pk.data_[arr_idx + 1] & ((1u << overhang) - 1)) << (num_bits_elem - overhang);
}
return bits & 0x3F;
}
inline uint32_t unpack(const index_t i) const { return unpack(*this, i); }
};
using f6x16_pk_t = pk_f6_t<16>;
template <>
struct numeric_traits<f6x16_pk_t>
{
static constexpr int PackedSize = 16;
};
} // namespace ck_tile

4
include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp
View File

@@ -314,7 +314,9 @@ struct MXFlatmmKernel : FlatmmKernel<TilePartitioner_, MXFlatmmPipeline_, Epilog
const auto& b_flat_pad_view = views.at(I1);
const auto& ds_pad_view = views.at(I2);
const auto& e_pad_view = views.at(I3);
// printf("%d %d %d\n",TilePartitioner::MPerBlock,TilePartitioner::KPerBlock,TilePartitioner::NPerBlock);
// printf("==============\n");
// printf("%d %d\n",MXFlatmmPipeline::flatNPerWarp,MXFlatmmPipeline::flatKPerWarp);
const auto& a_block_window = [&]() {
static_assert(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>,
"A tensor for mx must be RowMajor");

13
include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1.hpp
View File

@@ -115,7 +115,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Problem
static constexpr index_t NWarp = config.template at<2>();
static constexpr index_t MIterPerWarp = kMPerBlock / (MWarp * WG::kM);
static constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WG::kN);
static constexpr index_t NIterPerWarp = kNPerBlock / (NWarp * WG::kN);//512/(4*16)
static constexpr index_t KIterPerWarp = kKPerBlock / WG::kK;
static constexpr index_t KFlatBytesPerBlockPerIter = flatKPerWarp / BPackedSize;
@@ -569,7 +569,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Problem
constexpr auto packed_n_idx = nIter / number<NXdlPack>{};
constexpr auto packed_n_rank = nIter % number<NXdlPack>{};
return b_flat_dram_window.get_load_offset(
tuple<number<packed_n_idx * NXdlPack * NFlatPerBlockPerIter>,
tuple<number<packed_n_idx * NXdlPack * NFlatPerBlockPerIter/*4*/>,
number<0>>{}) +
b_flat_dram_window.get_load_offset(
tuple<number<packed_n_rank>, number<0>>{});
@@ -659,7 +659,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Problem
});
});
// move Scale B window to next K
move_tile_window(scale_b_dram_window, {0, kKPerBlock / (32 * KXdlPack)});
move_tile_window(scale_b_dram_window, {0, kKPerBlock / (32/*32个k占一个scale*/ * KXdlPack)});
__builtin_amdgcn_sched_barrier(0);
// Prefetch A1
@@ -738,7 +738,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Problem
operator()<ikxdl * MXdlPack + imxdl, ikxdl * NXdlPack + inxdl>(
c_warp_tensors(number<m_iter>{})(number<n_iter>{}),
bit_cast<typename WG::AWarpTensor>(
a_warp_tensor(number<AwarpIter>{})),
a_warp_tensor(number<AwarpIter>{})),//为什么这里不通过miter和kiter索引？
bit_cast<typename WG::BWarpTensor>(
b_warp_tensor_ping(number<n_iter>{})(number<k_iter>{})),
scale_a_tile_tensor_ping(mIter_pack)(kIter_pack)
@@ -754,6 +754,7 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Problem
{
constexpr auto AmIter = addr % 2 + addr / 4 * 2;
constexpr auto AkIter = addr / 2 % 2;
// if(blockIdx.x==0 && threadIdx.x==0)
a_warp_tensor(number<AwarpIter>{}) = load_tile_with_offset(
a_warp_window_ping,
tuple<number<AmIter * WG::kM>, number<AkIter * WG::kK>>{});
@@ -951,6 +952,8 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Problem
{
constexpr auto AmIter = addr % 2 + addr / 4 * 2;
constexpr auto AkIter = addr / 2 % 2;
// if(blockIdx.x==0 && threadIdx.x==0)
// printf("%d %d\n",AmIter,AkIter);
a_warp_tensor(number<AwarpIter>{}) = load_tile_with_offset(
a_warp_window_ping,
tuple<number<AmIter * WG::kM>, number<AkIter * WG::kK>>{});
@@ -1052,6 +1055,8 @@ struct MXFlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Problem
{
constexpr auto AmIter = addr % 2 + addr / 4 * 2;
constexpr auto AkIter = addr / 2 % 2;
// if(blockIdx.x==0 && threadIdx.x==0)
// printf("%d %d\n",AmIter,AkIter);
a_warp_tensor(number<AwarpIter>{}) = load_tile_with_offset(
a_warp_window_ping,
tuple<number<AmIter * WG::kM>, number<AkIter * WG::kK>>{});

10
include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp
View File

@@ -212,8 +212,8 @@ struct MXFlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy
if constexpr(K_Thread == AK1)
return make_static_tile_distribution(
tile_distribution_encoding< //
sequence<NWarps>,
tuple<sequence<MWarps, MXdlPack, MPerXdl>, sequence<K_Lane, AK1>>,
sequence<NWarps>,//4
tuple<sequence<MWarps, MXdlPack, MPerXdl>, sequence<K_Lane, AK1>>,//1,2,16 | 4,32
tuple<sequence<1, 0>, sequence<2, 1>>,
tuple<sequence<0, 0>, sequence<0, 2>>,
sequence<2>,
@@ -339,9 +339,9 @@ struct MXFlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy
CK_TILE_HOST_DEVICE static constexpr auto MakeMX_ScaleA_FlatDramTileDistribution()
{
return make_static_tile_distribution(
tile_distribution_encoding<sequence<NWarps>, // ?
tuple<sequence<MWarps, MPerXdl>, // second direction
sequence<K_Lane, 1>>, // first direction
tile_distribution_encoding<sequence<NWarps>,//4 // ?
tuple<sequence<MWarps, MPerXdl>,//1,16 // second direction
sequence<K_Lane, 1>>,//4,1 // first direction
tuple<sequence<1, 0>, sequence<2, 1>>, // which direction
tuple<sequence<0, 0>, sequence<0, 1>>, // which index
// <repeat, vec_load>