blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_selector.hpp Source File

blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_selector.hpp Source File#

Composable Kernel: blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_selector.hpp Source File
blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_selector.hpp
Go to the documentation of this file.
1// SPDX-License-Identifier: MIT
2// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
3
4#pragma once
5
6#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v1.hpp"
7#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v1.hpp"
8#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp"
9#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v3.hpp"
10namespace ck {
11
12template <BlockGemmPipelineVersion BlkGemmPipelineVer,
13 BlockGemmPipelineScheduler BlkGemmPipeSche,
14 index_t BlockSize,
15 typename ADataType,
16 typename BDataType,
17 typename ComputeDataType,
18 typename AccDataType,
19 typename ATileDesc,
20 typename BTileDesc,
21 typename AMmaTileDesc,
22 typename BMmaTileDesc,
23 index_t ABlockTransferSrcScalarPerVector,
24 index_t BBlockTransferSrcScalarPerVector,
25 index_t MPerBlock,
26 index_t NPerBlock,
27 index_t KPerBlock,
28 index_t MScaleBlock,
29 index_t NScaleBlock,
30 index_t KScaleBlock,
31 index_t MPerXDL,
32 index_t NPerXDL,
33 index_t MRepeat,
34 index_t NRepeat,
35 index_t KPack,
36 bool GUFusion = false>
37constexpr auto BlockGemmBlockMoeScaleBPreshufflePipeline_Selector()
38{
39 if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
40 {
41 if constexpr(GUFusion)
42 {
43 return BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v1<
44 BlkGemmPipeSche,
45 BlockSize,
46 ADataType,
47 BDataType,
48 ComputeDataType,
49 AccDataType,
50 ATileDesc,
51 BTileDesc,
52 AMmaTileDesc,
53 BMmaTileDesc,
54 ABlockTransferSrcScalarPerVector,
55 BBlockTransferSrcScalarPerVector,
56 MPerBlock,
57 NPerBlock,
58 KPerBlock,
59 MScaleBlock,
60 NScaleBlock,
61 KScaleBlock,
62 MPerXDL,
63 NPerXDL,
64 MRepeat,
65 NRepeat,
66 KPack>{};
67 }
68 else
69 {
70 return BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1<
71 BlkGemmPipeSche,
72 BlockSize,
73 ADataType,
74 BDataType,
75 ComputeDataType,
76 AccDataType,
77 ATileDesc,
78 BTileDesc,
79 AMmaTileDesc,
80 BMmaTileDesc,
81 ABlockTransferSrcScalarPerVector,
82 BBlockTransferSrcScalarPerVector,
83 MPerBlock,
84 NPerBlock,
85 KPerBlock,
86 MScaleBlock,
87 NScaleBlock,
88 KScaleBlock,
89 MPerXDL,
90 NPerXDL,
91 MRepeat,
92 NRepeat,
93 KPack>{};
94 }
95 }
96#if 0
97 else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2)
98 {
99 return BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v2<
100 BlkGemmPipeSche,
101 BlockSize,
102 ADataType,
103 BDataType,
104 ComputeDataType,
105 AccDataType,
106 ATileDesc,
107 BTileDesc,
108 AMmaTileDesc,
109 BMmaTileDesc,
110 ABlockTransferSrcScalarPerVector,
111 BBlockTransferSrcScalarPerVector,
112 MPerBlock,
113 NPerBlock,
114 KPerBlock,
115 MPerXDL,
116 NPerXDL,
117 MRepeat,
118 NRepeat,
119 KPack>{};
120 }
121#endif
122 else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
123 {
124 static_assert(MRepeat >= 4, "MRepeat should at least be 4 in BlockGemmPipelineVersion::v3");
125 if constexpr(GUFusion)
126 {
127 return BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3<
128 BlkGemmPipeSche,
129 BlockSize,
130 ADataType,
131 BDataType,
132 ComputeDataType,
133 AccDataType,
134 ATileDesc,
135 BTileDesc,
136 AMmaTileDesc,
137 BMmaTileDesc,
138 ABlockTransferSrcScalarPerVector,
139 BBlockTransferSrcScalarPerVector,
140 MPerBlock,
141 NPerBlock,
142 KPerBlock,
143 MScaleBlock,
144 NScaleBlock,
145 KScaleBlock,
146 MPerXDL,
147 NPerXDL,
148 MRepeat,
149 NRepeat,
150 KPack>{};
151 }
152 else
153 {
154 return BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3<
155 BlkGemmPipeSche,
156 BlockSize,
157 ADataType,
158 BDataType,
159 ComputeDataType,
160 AccDataType,
161 ATileDesc,
162 BTileDesc,
163 AMmaTileDesc,
164 BMmaTileDesc,
165 ABlockTransferSrcScalarPerVector,
166 BBlockTransferSrcScalarPerVector,
167 MPerBlock,
168 NPerBlock,
169 KPerBlock,
170 MScaleBlock,
171 NScaleBlock,
172 KScaleBlock,
173 MPerXDL,
174 NPerXDL,
175 MRepeat,
176 NRepeat,
177 KPack>{};
178 }
179 }
180 else
181 {
182 std::cerr << "BlockGemmPipeline configuration is not available" << std::endl;
183 }
184}
185
186} // namespace ck
blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v1.hpp
blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp
blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v1.hpp
blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v3.hpp
ck
Definition ck.hpp:268
ck::index_t
int32_t index_t
Definition ck.hpp:299
ck::BlockGemmPipelineVersion
BlockGemmPipelineVersion
Definition blkgemmpipe_scheduler.hpp:12
ck::BlockGemmPipelineVersion::v2
@ v2
Definition blkgemmpipe_scheduler.hpp:15
ck::BlockGemmPipelineVersion::v3
@ v3
Definition blkgemmpipe_scheduler.hpp:16
ck::BlockGemmPipelineVersion::v1
@ v1
Definition blkgemmpipe_scheduler.hpp:14
ck::BlockGemmPipelineScheduler
BlockGemmPipelineScheduler
Definition blkgemmpipe_scheduler.hpp:25
ck::BlockGemmBlockMoeScaleBPreshufflePipeline_Selector
constexpr auto BlockGemmBlockMoeScaleBPreshufflePipeline_Selector()
Definition blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_selector.hpp:37
ck::BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v1
Definition blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v1.hpp:40
ck::BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3
Definition blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp:40
ck::BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v1
Definition blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v1.hpp:40
ck::BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_v3
Definition blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_v3.hpp:40