已合并
feat(gmm): expand qgmm mx grouped examples #220
zhangzhizhuo创建于 27 天前
feat(gmm): expand qgmm mx grouped examples #220
已合并
共 6 个文件变更+307-129
| @@ -55,15 +55,24 @@ M 轴分组场景分别计算各组矩阵乘并按 M 轴拼接结果。 | |||
| 55 | | `k` | K 维度 | | 55 | | `k` | K 维度 | |
| 56 | | `n` | N 维度 | | 56 | | `n` | N 维度 | |
| 57 | | `dtype` | MX 数据类型 | | 57 | | `dtype` | MX 数据类型 | |
| 58 | -| `transA` | A 矩阵是否转置，当前样例配置为 `false` | | 58 | +| `transA` | A 矩阵是否转置 | |
| 59 | | `transB` | B 矩阵是否转置 | | 59 | | `transB` | B 矩阵是否转置 | |
| 60 | -| `format` | A、B 矩阵的数据格式及对应布局 | | 60 | +| `format` | A、B 矩阵的数据格式，支持 ND/NZ；转置由 `transA`/`transB` 表示 | |
| 61 | | `weight_mode` | 权重存储模式，支持 `single` 和 `multi` | | 61 | | `weight_mode` | 权重存储模式，支持 `single` 和 `multi` | |
| 62 | | `bias` | 是否启用偏置 | | 62 | | `bias` | 是否启用偏置 | |
| 63 | | `group_list_type` | Group List 类型 | | 63 | | `group_list_type` | Group List 类型 | |
| 64 | -| `l1_buffer_stage` | L1 buffer数 | | 64 | +| `group_list` | GroupList 原始值，以分号分隔；Length 填各组长度，Offset 填累计偏移，Sparse 填索引/长度对 | |
| 65 | +| `base_k` | L0 计算块的 K 大小 | | ||
| 66 | +| `tile_k_l1` | A/B 在 L1 中的 K 切分大小 | | ||
| 67 | +| `scale_k_l1` | ScaleA/ScaleB 在 L1 中的 K 切分大小 | | ||
| 68 | +| `l1_buffers` | L1 buffer 数量，支持 2 或 3 | | ||
| 69 | +| `db_l0c` | L0C buffer 数量 | | ||
| 70 | +| `a_full_load` | A 是否全载入；用于选择 kernel 编译期策略 | | ||
| 65 | 71 | ||
| 66 | -当前样例覆盖 MX 量化的 M 轴分组。GroupList 中的数值表示 M 轴上各组的大小或累计偏移。 | 72 | +`transA` 表示 A 矩阵的转置状态，Group List 的分组轴用于区分 M 轴分组和 K 轴分组，两者的参数语义相互独立。当前 QGMM MX 样例根据 `transA` 选择分组路径：`transA=false` 对应 M 轴分组，`transA=true` 对应 K 轴分组。 |
| 73 | + | ||
| 74 | +当前样例覆盖 MX 量化的 M 轴和 K 轴分组。Tiling 参数由每条 CSV case 独立配置，不在 kernel 中按分组类型固定。 | ||
| 75 | +`group_list` 中的分组长度必须为正，并完整覆盖分组轴；M 轴分组的总长度为 `e*m`，K 轴分组的总长度为 `k`。 | ||
| 67 | 76 | ||
| 68 | ## 编译和运行 | 77 | ## 编译和运行 |
| 69 | 78 | ||
| @@ -24,39 +24,64 @@ def _build_command(executable, case): | |||
| 24 | trans_b = case["transB"].strip().lower() == "true" | 24 | trans_b = case["transB"].strip().lower() == "true" |
| 25 | with_bias = case["bias"].strip().lower() == "true" | 25 | with_bias = case["bias"].strip().lower() == "true" |
| 26 | group_list_type = case["group_list_type"].strip().lower() | 26 | group_list_type = case["group_list_type"].strip().lower() |
| 27 | - l1_buffer_stage = int(case["l1_buffer_stage"]) | 27 | + base_k = int(case["base_k"]) |
| 28 | + tile_k_l1 = int(case["tile_k_l1"]) | ||
| 29 | + scale_k_l1 = int(case["scale_k_l1"]) | ||
| 30 | + l1_buffers = int(case["l1_buffers"]) | ||
| 31 | + db_l0c = int(case["db_l0c"]) | ||
| 32 | + a_full_load = case["a_full_load"].strip().lower() == "true" | ||
| 28 | layout_pair = [ | 33 | layout_pair = [ |
| 29 | item.strip() for item in case["format"].strip().strip('"()').lower().split(",") | 34 | item.strip() for item in case["format"].strip().strip('"()').lower().split(",") |
| 30 | ] | 35 | ] |
| 31 | - expected_a = "dn" if trans_a else "nd" | 36 | + expected_b = layout_pair[-1] |
| 32 | - expected_b = ( | 37 | + kernel_layout_b = ( |
| 33 | ("zn" if trans_b else "nz") | 38 | ("zn" if trans_b else "nz") |
| 34 | - if layout_pair[-1] in ("nz", "zn") | 39 | + if expected_b == "nz" |
| 35 | else ("dn" if trans_b else "nd") | 40 | else ("dn" if trans_b else "nd") |
| 36 | ) | 41 | ) |
| 37 | supported = ( | 42 | supported = ( |
| 38 | case["dtype"].strip() in SUPPORTED_DTYPES | 43 | case["dtype"].strip() in SUPPORTED_DTYPES |
| 39 | and len(layout_pair) == 2 | 44 | and len(layout_pair) == 2 |
| 40 | - and layout_pair == [expected_a, expected_b] | 45 | + and layout_pair == ["nd", expected_b] |
| 46 | + and expected_b in ("nd", "nz") | ||
| 41 | and case["weight_mode"].strip() in ("single", "multi") | 47 | and case["weight_mode"].strip() in ("single", "multi") |
| 42 | and (not with_bias or case["dtype"].strip().startswith("mxfp4")) | 48 | and (not with_bias or case["dtype"].strip().startswith("mxfp4")) |
| 43 | and (case["dtype"].strip() != "mxfp4_e1m2" or layout_pair[1] in ("nz", "zn")) | 49 | and (case["dtype"].strip() != "mxfp4_e1m2" or layout_pair[1] in ("nz", "zn")) |
| 44 | and group_list_type in SUPPORTED_GROUP_LISTS | 50 | and group_list_type in SUPPORTED_GROUP_LISTS |
| 45 | - and l1_buffer_stage in (2, 3) | 51 | + and base_k > 0 |
| 52 | + and tile_k_l1 > 0 | ||
| 53 | + and scale_k_l1 > 0 | ||
| 54 | + and l1_buffers in (2, 3) | ||
| 55 | + and db_l0c in (1, 2) | ||
| 46 | and all(int(case[dim]) > 0 for dim in ("e", "m", "n", "k")) | 56 | and all(int(case[dim]) > 0 for dim in ("e", "m", "n", "k")) |
| 57 | + and ( | ||
| 58 | + not trans_a | ||
| 59 | + or ( | ||
| 60 | + case["dtype"].strip() in ("mxfp8_e4m3", "mxfp8_e5m2") | ||
| 61 | + and layout_pair == ["nd", "nd"] | ||
| 62 | + and case["weight_mode"].strip() == "single" | ||
| 63 | + and not with_bias | ||
| 64 | + and group_list_type in ("length", "offset") | ||
| 65 | + ) | ||
| 66 | + ) | ||
| 47 | ) | 67 | ) |
| 48 | if not supported: | 68 | if not supported: |
| 49 | return None | 69 | return None |
| 50 | return [ | 70 | return [ |
| 51 | executable, | 71 | executable, |
| 52 | case["dtype"].strip(), | 72 | case["dtype"].strip(), |
| 53 | - layout_pair[1], | 73 | + kernel_layout_b, |
| 54 | case["weight_mode"].strip(), | 74 | case["weight_mode"].strip(), |
| 55 | *(case[dim].strip() for dim in ("e", "m", "n", "k")), | 75 | *(case[dim].strip() for dim in ("e", "m", "n", "k")), |
| 56 | str(trans_a).lower(), | 76 | str(trans_a).lower(), |
| 57 | str(with_bias).lower(), | 77 | str(with_bias).lower(), |
| 58 | group_list_type, | 78 | group_list_type, |
| 59 | - str(l1_buffer_stage), | 79 | + str(base_k), |
| 80 | + str(tile_k_l1), | ||
| 81 | + str(scale_k_l1), | ||
| 82 | + str(l1_buffers), | ||
| 83 | + str(db_l0c), | ||
| 84 | + str(a_full_load).lower(), | ||
| 60 | ] | 85 | ] |
| 61 | 86 | ||
| 62 | 87 | ||
| @@ -73,7 +98,7 @@ def _print_details(completed): | |||
| 73 | return details | 98 | return details |
| 74 | 99 | ||
| 75 | 100 | ||
| 76 | -def _generation_command(scripts_dir, command, case_dir): | 101 | +def _generation_command(scripts_dir, command, case_dir, case): |
| 77 | generation = [ | 102 | generation = [ |
| 78 | sys.executable, | 103 | sys.executable, |
| 79 | os.path.join(scripts_dir, "gen_data.py"), | 104 | os.path.join(scripts_dir, "gen_data.py"), |
| @@ -91,6 +116,8 @@ def _generation_command(scripts_dir, command, case_dir): | |||
| 91 | command[7], | 116 | command[7], |
| 92 | "--group-list-type", | 117 | "--group-list-type", |
| 93 | command[10], | 118 | command[10], |
| 119 | + "--group-list", | ||
| 120 | + (case.get("group_list") or "").strip(), | ||
| 94 | "--output-dir", | 121 | "--output-dir", |
| 95 | case_dir, | 122 | case_dir, |
| 96 | ] | 123 | ] |
| @@ -117,7 +144,7 @@ def _run_case(executable, scripts_dir, output_root, case): | |||
| 117 | case_dir = os.path.join(output_root, case_name) | 144 | case_dir = os.path.join(output_root, case_name) |
| 118 | shutil.rmtree(case_dir, ignore_errors=True) | 145 | shutil.rmtree(case_dir, ignore_errors=True) |
| 119 | os.makedirs(case_dir) | 146 | os.makedirs(case_dir) |
| 120 | - completed = _run(_generation_command(scripts_dir, command, case_dir)) | 147 | + completed = _run(_generation_command(scripts_dir, command, case_dir, case)) |
| 121 | details = _print_details(completed) | 148 | details = _print_details(completed) |
| 122 | if completed.returncode != 0: | 149 | if completed.returncode != 0: |
| 123 | return _failure(case_name, "gen_data", details or "data generation failed") | 150 | return _failure(case_name, "gen_data", details or "data generation failed") |
| @@ -132,6 +159,12 @@ def _run_case(executable, scripts_dir, output_root, case): | |||
| 132 | os.path.join(scripts_dir, "verify_result.py"), | 159 | os.path.join(scripts_dir, "verify_result.py"), |
| 133 | os.path.join(case_dir, "golden_c.bin"), | 160 | os.path.join(case_dir, "golden_c.bin"), |
| 134 | output_path, | 161 | output_path, |
| 162 | + "--groups", | ||
| 163 | + command[4], | ||
| 164 | + "--m", | ||
| 165 | + command[5], | ||
| 166 | + "--n", | ||
| 167 | + command[6], | ||
| 135 | ] | 168 | ] |
| 136 | ) | 169 | ) |
| 137 | details = _print_details(completed) | 170 | details = _print_details(completed) |
| @@ -15,7 +15,7 @@ | |||
| 15 | 15 | ||
| 16 | /** | 16 | /** |
| 17 | * @file qgmm_mx.cpp | 17 | * @file qgmm_mx.cpp |
| 18 | - * @brief QGMM MX grouped-matmul example covering MXFP4/MXFP8 and ND/DN/NZ/ZN weights. | 18 | + * @brief QGMM MX grouped-matmul example covering MXFP4/MXFP8 and ND/NZ weights. |
| 19 | */ | 19 | */ |
| 20 | 20 | ||
| 21 | 21 | ||
| @@ -45,16 +45,16 @@ using NdLayout = AscendC::Te::NDExtLayoutPtn; | |||
| 45 | template <typename T> | 45 | template <typename T> |
| 46 | inline constexpr bool IS_FP4_TYPE = std::is_same_v<T, fp4x2_e2m1_t> || std::is_same_v<T, fp4x2_e1m2_t>; | 46 | inline constexpr bool IS_FP4_TYPE = std::is_same_v<T, fp4x2_e2m1_t> || std::is_same_v<T, fp4x2_e1m2_t>; |
| 47 | 47 | ||
| 48 | -template <typename AType, typename BType, typename LayoutA, typename LayoutB> | 48 | +template <typename AType, typename BType, typename LayoutA, typename LayoutB, uint64_t FullLoadMode> |
| 49 | __global__ __aicore__ void qgmm_mx_kernel(GM_ADDR a, GM_ADDR b, GM_ADDR scaleA, GM_ADDR scaleB, GM_ADDR c, GM_ADDR bias, | 49 | __global__ __aicore__ void qgmm_mx_kernel(GM_ADDR a, GM_ADDR b, GM_ADDR scaleA, GM_ADDR scaleB, GM_ADDR c, GM_ADDR bias, |
| 50 | GM_ADDR groupList, uint32_t groupNum, int64_t m, int64_t n, int64_t k, | 50 | GM_ADDR groupList, uint32_t groupNum, int64_t m, int64_t n, int64_t k, |
| 51 | - uint8_t singleW, uint8_t groupListType, uint8_t l1BufferStage, | 51 | + uint32_t baseK, uint32_t tileKL1, uint32_t scaleKL1, uint8_t l1BufferStage, |
| 52 | - uint8_t withBias) | 52 | + uint8_t dbL0C, uint8_t singleW, uint8_t groupListType, uint8_t withBias) |
| 53 | { | 53 | { |
| 54 | KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); | 54 | KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); |
| 55 | AscendC::InitSocState(); | 55 | AscendC::InitSocState(); |
| 56 | using ProblemShape = AscendC::Te::Shape<int64_t, int64_t, int64_t, int64_t>; | 56 | using ProblemShape = AscendC::Te::Shape<int64_t, int64_t, int64_t, int64_t>; |
| 57 | - using Policy = Blaze::Gemm::GroupedMatmulWithScaleMx<0>; | 57 | + using Policy = Blaze::Gemm::GroupedMatmulWithScaleMx<FullLoadMode>; |
| 58 | using Mmad = Blaze::Gemm::Block::BlockMmad<Policy, AType, LayoutA, BType, LayoutB, float, NdLayout, float, | 58 | using Mmad = Blaze::Gemm::Block::BlockMmad<Policy, AType, LayoutA, BType, LayoutB, float, NdLayout, float, |
| 59 | NdLayout>; | 59 | NdLayout>; |
| 60 | using Kernel = Blaze::Gemm::Kernel::GemmUniversal<ProblemShape, Mmad, Blaze::Gemm::Block::BlockEpilogueEmpty, | 60 | using Kernel = Blaze::Gemm::Kernel::GemmUniversal<ProblemShape, Mmad, Blaze::Gemm::Block::BlockEpilogueEmpty, |
| @@ -68,7 +68,16 @@ __global__ __aicore__ void qgmm_mx_kernel(GM_ADDR a, GM_ADDR b, GM_ADDR scaleA, | |||
| 68 | p.mmadParams.scaleBGmAddr = scaleB; | 68 | p.mmadParams.scaleBGmAddr = scaleB; |
| 69 | p.mmadParams.biasGmAddr = bias; | 69 | p.mmadParams.biasGmAddr = bias; |
| 70 | p.groupListGmAddr = groupList; | 70 | p.groupListGmAddr = groupList; |
| 71 | - p.gmmParams = {groupNum, m, n, k, 16, 64, 64, 64, 64, 2, 2, withBias, 1, l1BufferStage, 0, groupListType, singleW}; | 71 | + constexpr uint8_t groupType = std::is_same_v<LayoutA, AscendC::Te::DNExtLayoutPtn> ? 2 : 0; |
| 72 | + p.gmmParams = {groupNum, m, | ||
| 73 | + n, k, | ||
| 74 | + 16, 64, | ||
| 75 | + baseK, tileKL1, | ||
| 76 | + tileKL1, scaleKL1, | ||
| 77 | + scaleKL1, withBias, | ||
| 78 | + dbL0C, l1BufferStage, | ||
| 79 | + groupType, groupListType, | ||
| 80 | + singleW}; | ||
| 72 | Kernel kernel; | 81 | Kernel kernel; |
| 73 | kernel(p); | 82 | kernel(p); |
| 74 | } | 83 | } |
| @@ -120,6 +129,14 @@ struct DeviceBuffers { | |||
| 120 | } | 129 | } |
| 121 | }; | 130 | }; |
| 122 | 131 | ||
| 132 | +struct QgmmTilingConfig { | ||
| 133 | + uint32_t baseK; | ||
| 134 | + uint32_t tileKL1; | ||
| 135 | + uint32_t scaleKL1; | ||
| 136 | + uint8_t l1BufferStage; | ||
| 137 | + uint8_t dbL0C; | ||
| 138 | +}; | ||
| 139 | + | ||
| 123 | void AllocateCommonBuffers(DeviceBuffers& device, const CaseBytes& bytes, size_t biasBytes, size_t groupListBytes) | 140 | void AllocateCommonBuffers(DeviceBuffers& device, const CaseBytes& bytes, size_t biasBytes, size_t groupListBytes) |
| 124 | { | 141 | { |
| 125 | ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&device.a), bytes.a, ACL_MEM_MALLOC_HUGE_FIRST)); | 142 | ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&device.a), bytes.a, ACL_MEM_MALLOC_HUGE_FIRST)); |
| @@ -144,15 +161,17 @@ std::vector<uint8_t> ReadBinary(const std::string& path, size_t size) | |||
| 144 | return data; | 161 | return data; |
| 145 | } | 162 | } |
| 146 | 163 | ||
| 147 | -template <bool MultiTensor> | 164 | +template <bool MultiTensor, bool KGrouped> |
| 148 | void PrepareWeightBuffers(DeviceBuffers& device, const CaseBytes& bytes, uint32_t groupNum, const std::string& dataDir) | 165 | void PrepareWeightBuffers(DeviceBuffers& device, const CaseBytes& bytes, uint32_t groupNum, const std::string& dataDir) |
| 149 | { | 166 | { |
| 150 | if constexpr (!MultiTensor) { | 167 | if constexpr (!MultiTensor) { |
| 151 | - ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&device.b), groupNum * bytes.bGroup, ACL_MEM_MALLOC_HUGE_FIRST)); | 168 | + const size_t tensorCount = KGrouped ? 1U : groupNum; |
| 152 | - ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&device.scaleB), groupNum * bytes.scaleBGroup, | 169 | + ACL_CHECK( |
| 170 | + aclrtMalloc(reinterpret_cast<void**>(&device.b), tensorCount * bytes.bGroup, ACL_MEM_MALLOC_HUGE_FIRST)); | ||
| 171 | + ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&device.scaleB), tensorCount * bytes.scaleBGroup, | ||
| 153 | ACL_MEM_MALLOC_HUGE_FIRST)); | 172 | ACL_MEM_MALLOC_HUGE_FIRST)); |
| 154 | - const auto allWeights = ReadBinary(dataDir + "/input_b.bin", groupNum * bytes.bGroup); | 173 | + const auto allWeights = ReadBinary(dataDir + "/input_b.bin", tensorCount * bytes.bGroup); |
| 155 | - const auto allScales = ReadBinary(dataDir + "/scale_b.bin", groupNum * bytes.scaleBGroup); | 174 | + const auto allScales = ReadBinary(dataDir + "/scale_b.bin", tensorCount * bytes.scaleBGroup); |
| 156 | ACL_CHECK( | 175 | ACL_CHECK( |
| 157 | aclrtMemcpy(device.b, allWeights.size(), allWeights.data(), allWeights.size(), ACL_MEMCPY_HOST_TO_DEVICE)); | 176 | aclrtMemcpy(device.b, allWeights.size(), allWeights.data(), allWeights.size(), ACL_MEMCPY_HOST_TO_DEVICE)); |
| 158 | ACL_CHECK(aclrtMemcpy(device.scaleB, allScales.size(), allScales.data(), allScales.size(), | 177 | ACL_CHECK(aclrtMemcpy(device.scaleB, allScales.size(), allScales.data(), allScales.size(), |
| @@ -211,15 +230,16 @@ void WriteOutput(const std::string& outputPath, const std::vector<float>& output | |||
| 211 | } | 230 | } |
| 212 | } | 231 | } |
| 213 | 232 | ||
| 214 | -template <typename AType, typename BType, typename LayoutA, typename LayoutB, bool MultiTensor> | 233 | +template <typename AType, typename BType, typename LayoutA, typename LayoutB, bool MultiTensor, uint64_t FullLoadMode> |
| 215 | void LaunchKernel(DeviceBuffers& device, const CaseBytes& bytes, uint32_t groupNum, int64_t m, int64_t n, int64_t k, | 234 | void LaunchKernel(DeviceBuffers& device, const CaseBytes& bytes, uint32_t groupNum, int64_t m, int64_t n, int64_t k, |
| 216 | - uint8_t groupListType, uint8_t l1BufferStage, bool withBias, const std::string& outputPath) | 235 | + uint8_t groupListType, const QgmmTilingConfig& tiling, bool withBias, const std::string& outputPath) |
| 217 | { | 236 | { |
| 218 | aclrtStream stream = nullptr; | 237 | aclrtStream stream = nullptr; |
| 219 | ACL_CHECK(aclrtCreateStream(&stream)); | 238 | ACL_CHECK(aclrtCreateStream(&stream)); |
| 220 | - qgmm_mx_kernel<AType, BType, LayoutA, LayoutB><<<static_cast<uint32_t>(GetAicCoreNum()), 0, stream>>>( | 239 | + qgmm_mx_kernel<AType, BType, LayoutA, LayoutB, FullLoadMode><<<static_cast<uint32_t>(GetAicCoreNum()), 0, stream>>>( |
| 221 | device.a, device.b, device.scaleA, device.scaleB, device.c, device.bias, device.groupList, groupNum, m, n, k, | 240 | device.a, device.b, device.scaleA, device.scaleB, device.c, device.bias, device.groupList, groupNum, m, n, k, |
| 222 | - MultiTensor ? 0 : 1, groupListType, l1BufferStage, withBias ? 1 : 0); | 241 | + tiling.baseK, tiling.tileKL1, tiling.scaleKL1, tiling.l1BufferStage, tiling.dbL0C, MultiTensor ? 0 : 1, |
| 242 | + groupListType, withBias ? 1 : 0); | ||
| 223 | ACL_CHECK(aclrtSynchronizeStream(stream)); | 243 | ACL_CHECK(aclrtSynchronizeStream(stream)); |
| 224 | std::vector<float> output(bytes.c / sizeof(float), -1.0f); | 244 | std::vector<float> output(bytes.c / sizeof(float), -1.0f); |
| 225 | ACL_CHECK(aclrtMemcpy(output.data(), bytes.c, device.c, bytes.c, ACL_MEMCPY_DEVICE_TO_HOST)); | 245 | ACL_CHECK(aclrtMemcpy(output.data(), bytes.c, device.c, bytes.c, ACL_MEMCPY_DEVICE_TO_HOST)); |
| @@ -240,8 +260,8 @@ std::vector<int64_t> MakeGroupList(uint32_t groupNum, int64_t splitValue, uint8_ | |||
| 240 | return groupList; | 260 | return groupList; |
| 241 | } | 261 | } |
| 242 | 262 | ||
| 243 | -template <typename AType, typename BType, typename LayoutA, typename LayoutB, bool MultiTensor> | 263 | +template <typename AType, typename BType, typename LayoutA, typename LayoutB, bool MultiTensor, uint64_t FullLoadMode> |
| 244 | -int RunCase(uint32_t groupNum, int64_t m, int64_t n, int64_t k, uint8_t groupListType, uint8_t l1BufferStage, | 264 | +int RunCase(uint32_t groupNum, int64_t m, int64_t n, int64_t k, uint8_t groupListType, const QgmmTilingConfig& tiling, |
| 245 | bool withBias, const std::string& dataDir, const std::string& outputPath) | 265 | bool withBias, const std::string& dataDir, const std::string& outputPath) |
| 246 | { | 266 | { |
| 247 | constexpr bool transA = std::is_same_v<LayoutA, AscendC::Te::DNExtLayoutPtn>; | 267 | constexpr bool transA = std::is_same_v<LayoutA, AscendC::Te::DNExtLayoutPtn>; |
| @@ -253,62 +273,66 @@ int RunCase(uint32_t groupNum, int64_t m, int64_t n, int64_t k, uint8_t groupLis | |||
| 253 | const size_t storedK = weightNz ? (transB ? ((k + c0 - 1U) / c0 * c0) : ((k + 15U) / 16U * 16U)) : k; | 273 | const size_t storedK = weightNz ? (transB ? ((k + c0 - 1U) / c0 * c0) : ((k + 15U) / 16U * 16U)) : k; |
| 254 | const size_t storedN = weightNz ? (transB ? ((n + 15U) / 16U * 16U) : ((n + c0 - 1U) / c0 * c0)) : n; | 274 | const size_t storedN = weightNz ? (transB ? ((n + 15U) / 16U * 16U) : ((n + c0 - 1U) / c0 * c0)) : n; |
| 255 | const size_t scaleK = static_cast<size_t>((k + 63) / 64) * 2U; | 275 | const size_t scaleK = static_cast<size_t>((k + 63) / 64) * 2U; |
| 256 | - const size_t scaleFactor = transA ? 2U : 1U; | 276 | + const size_t aElements = transA ? static_cast<size_t>(m * k) : static_cast<size_t>(groupNum) * m * k; |
| 257 | - const CaseBytes bytes = {MxBytes<AType>(static_cast<size_t>(groupNum) * m * k), MxBytes<BType>(storedK * storedN), | 277 | + const size_t bElements = transA ? static_cast<size_t>(k * n) : storedK * storedN; |
| 258 | - static_cast<size_t>(groupNum) * m * scaleK * scaleFactor * sizeof(ScaleType), | 278 | + const size_t scaleAElements = transA ? (static_cast<size_t>(k / 64) + groupNum) * m * 2U : |
| 259 | - static_cast<size_t>(n) * scaleK * scaleFactor * sizeof(ScaleType), | 279 | + static_cast<size_t>(groupNum) * m * scaleK; |
| 260 | - static_cast<size_t>(transA ? m : groupNum * m) * n * sizeof(float)}; | 280 | + const size_t scaleBElements = transA ? (static_cast<size_t>(k / 64) + groupNum) * n * 2U : |
| 281 | + static_cast<size_t>(n) * scaleK; | ||
| 282 | + const CaseBytes bytes = {MxBytes<AType>(aElements), MxBytes<BType>(bElements), scaleAElements * sizeof(ScaleType), | ||
| 283 | + scaleBElements * sizeof(ScaleType), static_cast<size_t>(groupNum) * m * n * sizeof(float)}; | ||
| 261 | const std::vector<int64_t> groupList = MakeGroupList(groupNum, transA ? k : m, groupListType); | 284 | const std::vector<int64_t> groupList = MakeGroupList(groupNum, transA ? k : m, groupListType); |
| 262 | DeviceBuffers device(groupNum); | 285 | DeviceBuffers device(groupNum); |
| 263 | AllocateCommonBuffers(device, bytes, static_cast<size_t>(groupNum) * n * sizeof(float), | 286 | AllocateCommonBuffers(device, bytes, static_cast<size_t>(groupNum) * n * sizeof(float), |
| 264 | groupList.size() * sizeof(int64_t)); | 287 | groupList.size() * sizeof(int64_t)); |
| 265 | - PrepareWeightBuffers<MultiTensor>(device, bytes, groupNum, dataDir); | 288 | + PrepareWeightBuffers<MultiTensor, transA>(device, bytes, groupNum, dataDir); |
| 266 | CopyCommonInputs(device, bytes, static_cast<size_t>(groupNum) * n * sizeof(float), | 289 | CopyCommonInputs(device, bytes, static_cast<size_t>(groupNum) * n * sizeof(float), |
| 267 | groupList.size() * sizeof(int64_t), dataDir); | 290 | groupList.size() * sizeof(int64_t), dataDir); |
| 268 | - LaunchKernel<AType, BType, LayoutA, LayoutB, MultiTensor>(device, bytes, groupNum, m, n, k, groupListType, | 291 | + LaunchKernel<AType, BType, LayoutA, LayoutB, MultiTensor, FullLoadMode>( |
| 269 | - l1BufferStage, withBias, outputPath); | 292 | + device, bytes, groupNum, m, n, k, groupListType, tiling, withBias, outputPath); |
| 270 | return 0; | 293 | return 0; |
| 271 | } | 294 | } |
| 272 | 295 | ||
| 273 | -template <typename T, typename LayoutA, typename LayoutB> | 296 | +template <typename T, typename LayoutA, typename LayoutB, uint64_t FullLoadMode> |
| 274 | -int DispatchMulti(bool multi, uint32_t e, int64_t m, int64_t n, int64_t k, uint8_t groupListType, uint8_t l1BufferStage, | 297 | +int DispatchMulti(bool multi, uint32_t e, int64_t m, int64_t n, int64_t k, uint8_t groupListType, |
| 275 | - bool withBias, const std::string& dataDir, const std::string& outputPath) | 298 | + const QgmmTilingConfig& tiling, bool withBias, const std::string& dataDir, |
| 299 | + const std::string& outputPath) | ||
| 276 | { | 300 | { |
| 277 | - return multi ? RunCase<T, T, LayoutA, LayoutB, true>(e, m, n, k, groupListType, l1BufferStage, withBias, dataDir, | 301 | + return multi ? RunCase<T, T, LayoutA, LayoutB, true, FullLoadMode>(e, m, n, k, groupListType, tiling, withBias, |
| 278 | - outputPath) : | 302 | + dataDir, outputPath) : |
| 279 | - RunCase<T, T, LayoutA, LayoutB, false>(e, m, n, k, groupListType, l1BufferStage, withBias, dataDir, | 303 | + RunCase<T, T, LayoutA, LayoutB, false, FullLoadMode>(e, m, n, k, groupListType, tiling, withBias, |
| 280 | - outputPath); | 304 | + dataDir, outputPath); |
| 281 | } | 305 | } |
| 282 | 306 | ||
| 283 | -template <typename T, typename LayoutA> | 307 | +template <typename T, typename LayoutA, uint64_t FullLoadMode> |
| 284 | int DispatchFormat(const std::string& format, bool multi, uint32_t e, int64_t m, int64_t n, int64_t k, | 308 | int DispatchFormat(const std::string& format, bool multi, uint32_t e, int64_t m, int64_t n, int64_t k, |
| 285 | - uint8_t groupListType, uint8_t l1BufferStage, bool withBias, const std::string& dataDir, | 309 | + uint8_t groupListType, const QgmmTilingConfig& tiling, bool withBias, const std::string& dataDir, |
| 286 | const std::string& outputPath) | 310 | const std::string& outputPath) |
| 287 | { | 311 | { |
| 288 | if (format == "nd") | 312 | if (format == "nd") |
| 289 | - return DispatchMulti<T, LayoutA, NdLayout>(multi, e, m, n, k, groupListType, l1BufferStage, withBias, dataDir, | 313 | + return DispatchMulti<T, LayoutA, NdLayout, FullLoadMode>(multi, e, m, n, k, groupListType, tiling, withBias, |
| 290 | - outputPath); | 314 | + dataDir, outputPath); |
| 291 | if (format == "dn") | 315 | if (format == "dn") |
| 292 | - return DispatchMulti<T, LayoutA, AscendC::Te::DNExtLayoutPtn>(multi, e, m, n, k, groupListType, l1BufferStage, | 316 | + return DispatchMulti<T, LayoutA, AscendC::Te::DNExtLayoutPtn, FullLoadMode>( |
| 293 | - withBias, dataDir, outputPath); | 317 | + multi, e, m, n, k, groupListType, tiling, withBias, dataDir, outputPath); |
| 294 | if (format == "nz") | 318 | if (format == "nz") |
| 295 | - return DispatchMulti<T, LayoutA, AscendC::Te::NZLayoutPtn>(multi, e, m, n, k, groupListType, l1BufferStage, | 319 | + return DispatchMulti<T, LayoutA, AscendC::Te::NZLayoutPtn, FullLoadMode>(multi, e, m, n, k, groupListType, |
| 296 | - withBias, dataDir, outputPath); | 320 | + tiling, withBias, dataDir, outputPath); |
| 297 | if (format == "zn") | 321 | if (format == "zn") |
| 298 | - return DispatchMulti<T, LayoutA, AscendC::Te::ZNLayoutPtn>(multi, e, m, n, k, groupListType, l1BufferStage, | 322 | + return DispatchMulti<T, LayoutA, AscendC::Te::ZNLayoutPtn, FullLoadMode>(multi, e, m, n, k, groupListType, |
| 299 | - withBias, dataDir, outputPath); | 323 | + tiling, withBias, dataDir, outputPath); |
| 300 | return 2; | 324 | return 2; |
| 301 | } | 325 | } |
| 302 | 326 | ||
| 303 | -template <typename T> | 327 | +template <typename T, uint64_t FullLoadMode> |
| 304 | int DispatchTransA(const std::string& format, bool multi, bool transA, uint32_t e, int64_t m, int64_t n, int64_t k, | 328 | int DispatchTransA(const std::string& format, bool multi, bool transA, uint32_t e, int64_t m, int64_t n, int64_t k, |
| 305 | - uint8_t groupListType, uint8_t l1BufferStage, bool withBias, const std::string& dataDir, | 329 | + uint8_t groupListType, const QgmmTilingConfig& tiling, bool withBias, const std::string& dataDir, |
| 306 | const std::string& outputPath) | 330 | const std::string& outputPath) |
| 307 | { | 331 | { |
| 308 | - return transA ? DispatchFormat<T, AscendC::Te::DNExtLayoutPtn>(format, multi, e, m, n, k, groupListType, | 332 | + return transA ? DispatchFormat<T, AscendC::Te::DNExtLayoutPtn, FullLoadMode>( |
| 309 | - l1BufferStage, withBias, dataDir, outputPath) : | 333 | + format, multi, e, m, n, k, groupListType, tiling, withBias, dataDir, outputPath) : |
| 310 | - DispatchFormat<T, NdLayout>(format, multi, e, m, n, k, groupListType, l1BufferStage, withBias, | 334 | + DispatchFormat<T, NdLayout, FullLoadMode>(format, multi, e, m, n, k, groupListType, tiling, |
| 311 | - dataDir, outputPath); | 335 | + withBias, dataDir, outputPath); |
| 312 | } | 336 | } |
| 313 | 337 | ||
| 314 | struct QgmmCaseConfig { | 338 | struct QgmmCaseConfig { |
| @@ -323,7 +347,8 @@ struct QgmmCaseConfig { | |||
| 323 | bool withBias; | 347 | bool withBias; |
| 324 | std::string groupList; | 348 | std::string groupList; |
| 325 | uint8_t groupListType; | 349 | uint8_t groupListType; |
| 326 | - uint8_t l1BufferStage; | 350 | + QgmmTilingConfig tiling; |
| 351 | + bool aFullLoad; | ||
| 327 | std::string dataDir; | 352 | std::string dataDir; |
| 328 | std::string outputPath; | 353 | std::string outputPath; |
| 329 | }; | 354 | }; |
| @@ -332,7 +357,8 @@ void PrintUsage() | |||
| 332 | { | 357 | { |
| 333 | std::cerr << "Usage: qgmm_mx <mxfp4_e2m1|mxfp4_e1m2|mxfp8_e4m3|mxfp8_e5m2> " | 358 | std::cerr << "Usage: qgmm_mx <mxfp4_e2m1|mxfp4_e1m2|mxfp8_e4m3|mxfp8_e5m2> " |
| 334 | "<nd|dn|nz|zn> <single|multi> <e> <m> <n> <k> <transA> <bias> " | 359 | "<nd|dn|nz|zn> <single|multi> <e> <m> <n> <k> <transA> <bias> " |
| 335 | - "<length|offset|sparse> <l1BufferStage> <dataDir> <outputPath>" | 360 | + "<length|offset|sparse> <baseK> <tileKL1> <scaleKL1> <l1Buffers> <dbL0C> <aFullLoad> " |
| 361 | + "<dataDir> <outputPath>" | ||
| 336 | << std::endl; | 362 | << std::endl; |
| 337 | } | 363 | } |
| 338 | 364 | ||
| @@ -350,9 +376,14 @@ QgmmCaseConfig ParseConfig(char** argv) | |||
| 350 | config.withBias = std::string(argv[9]) == "true"; | 376 | config.withBias = std::string(argv[9]) == "true"; |
| 351 | config.groupList = argv[10]; | 377 | config.groupList = argv[10]; |
| 352 | config.groupListType = config.groupList == "offset" ? 0 : (config.groupList == "length" ? 1 : 2); | 378 | config.groupListType = config.groupList == "offset" ? 0 : (config.groupList == "length" ? 1 : 2); |
| 353 | - config.l1BufferStage = static_cast<uint8_t>(std::stoul(argv[11])); | 379 | + config.tiling.baseK = static_cast<uint32_t>(std::stoul(argv[11])); |
| 354 | - config.dataDir = argv[12]; | 380 | + config.tiling.tileKL1 = static_cast<uint32_t>(std::stoul(argv[12])); |
| 355 | - config.outputPath = argv[13]; | 381 | + config.tiling.scaleKL1 = static_cast<uint32_t>(std::stoul(argv[13])); |
| 382 | + config.tiling.l1BufferStage = static_cast<uint8_t>(std::stoul(argv[14])); | ||
| 383 | + config.tiling.dbL0C = static_cast<uint8_t>(std::stoul(argv[15])); | ||
| 384 | + config.aFullLoad = std::string(argv[16]) == "true"; | ||
| 385 | + config.dataDir = argv[17]; | ||
| 386 | + config.outputPath = argv[18]; | ||
| 356 | return config; | 387 | return config; |
| 357 | } | 388 | } |
| 358 | 389 | ||
| @@ -361,17 +392,25 @@ bool IsValidConfig(const QgmmCaseConfig& config) | |||
| 361 | const bool validShape = config.groupNum > 0 && config.m > 0 && config.n > 0 && config.k > 0; | 392 | const bool validShape = config.groupNum > 0 && config.m > 0 && config.n > 0 && config.k > 0; |
| 362 | const bool validGroupList = config.groupList == "offset" || config.groupList == "length" || | 393 | const bool validGroupList = config.groupList == "offset" || config.groupList == "length" || |
| 363 | config.groupList == "sparse"; | 394 | config.groupList == "sparse"; |
| 364 | - const bool validL1Buffer = config.l1BufferStage == 2 || config.l1BufferStage == 3; | 395 | + const bool validTiling = config.tiling.baseK > 0 && config.tiling.tileKL1 > 0 && config.tiling.scaleKL1 > 0 && |
| 365 | - return validShape && validGroupList && validL1Buffer; | 396 | + (config.tiling.l1BufferStage == 2 || config.tiling.l1BufferStage == 3) && |
| 397 | + (config.tiling.dbL0C == 1 || config.tiling.dbL0C == 2); | ||
| 398 | + return validShape && validGroupList && validTiling; | ||
| 366 | } | 399 | } |
| 367 | 400 | ||
| 368 | -template <typename T> | 401 | +template <typename T, uint64_t FullLoadMode> |
| 369 | int DispatchConfig(const QgmmCaseConfig& config) | 402 | int DispatchConfig(const QgmmCaseConfig& config) |
| 370 | { | 403 | { |
| 371 | const bool multi = config.weightMode == "multi"; | 404 | const bool multi = config.weightMode == "multi"; |
| 372 | - return DispatchTransA<T>(config.format, multi, config.transA, config.groupNum, config.m, config.n, config.k, | 405 | + return DispatchTransA<T, FullLoadMode>(config.format, multi, config.transA, config.groupNum, config.m, config.n, |
| 373 | - config.groupListType, config.l1BufferStage, config.withBias, config.dataDir, | 406 | + config.k, config.groupListType, config.tiling, config.withBias, |
| 374 | - config.outputPath); | 407 | + config.dataDir, config.outputPath); |
| 408 | +} | ||
| 409 | + | ||
| 410 | +template <typename T> | ||
| 411 | +int DispatchFullLoad(const QgmmCaseConfig& config) | ||
| 412 | +{ | ||
| 413 | + return config.aFullLoad ? DispatchConfig<T, 1>(config) : DispatchConfig<T, 0>(config); | ||
| 375 | } | 414 | } |
| 376 | 415 | ||
| 377 | int RunConfiguredCase(const QgmmCaseConfig& config) | 416 | int RunConfiguredCase(const QgmmCaseConfig& config) |
| @@ -379,13 +418,13 @@ int RunConfiguredCase(const QgmmCaseConfig& config) | |||
| 379 | if (config.weightMode != "single" && config.weightMode != "multi") | 418 | if (config.weightMode != "single" && config.weightMode != "multi") |
| 380 | return 2; | 419 | return 2; |
| 381 | if (config.dtype == "mxfp8_e4m3") | 420 | if (config.dtype == "mxfp8_e4m3") |
| 382 | - return DispatchConfig<fp8_e4m3fn_t>(config); | 421 | + return DispatchFullLoad<fp8_e4m3fn_t>(config); |
| 383 | if (config.dtype == "mxfp8_e5m2") | 422 | if (config.dtype == "mxfp8_e5m2") |
| 384 | - return DispatchConfig<fp8_e5m2_t>(config); | 423 | + return DispatchFullLoad<fp8_e5m2_t>(config); |
| 385 | if (config.dtype == "mxfp4_e2m1") | 424 | if (config.dtype == "mxfp4_e2m1") |
| 386 | - return DispatchConfig<fp4x2_e2m1_t>(config); | 425 | + return DispatchFullLoad<fp4x2_e2m1_t>(config); |
| 387 | if (config.dtype == "mxfp4_e1m2") | 426 | if (config.dtype == "mxfp4_e1m2") |
| 388 | - return DispatchConfig<fp4x2_e1m2_t>(config); | 427 | + return DispatchFullLoad<fp4x2_e1m2_t>(config); |
| 389 | return 2; | 428 | return 2; |
| 390 | } | 429 | } |
| 391 | 430 | ||
| @@ -403,7 +442,7 @@ void PrintResult(int ret, const QgmmCaseConfig& config) | |||
| 403 | 442 | ||
| 404 | int main(int argc, char** argv) | 443 | int main(int argc, char** argv) |
| 405 | { | 444 | { |
| 406 | - if (argc != 14) { | 445 | + if (argc != 19) { |
| 407 | PrintUsage(); | 446 | PrintUsage(); |
| 408 | return 2; | 447 | return 2; |
| 409 | } | 448 | } |
| @@ -1,14 +1,28 @@ | |||
| 1 | -casename,e,m,k,n,dtype,transA,transB,format,weight_mode,bias,group_list_type,l1_buffer_stage | 1 | +casename,e,m,k,n,dtype,transA,transB,format,weight_mode,bias,group_list_type,group_list,base_k,tile_k_l1,scale_k_l1,l1_buffers,db_l0c,a_full_load |
| 2 | -qgmm_mgroup_e1_m16_k64_n64_fp8_e4m3_weight_nd_single_length_l1_2,1,16,64,64,mxfp8_e4m3,false,false,"(ND,ND)",single,false,length,2 | 2 | +qgmm_mgroup_e1_m16_k64_n64_fp8_e4m3_weight_nd_single_length_l1_2,1,16,64,64,mxfp8_e4m3,false,false,"(ND,ND)",single,false,length,16,64,64,2,2,1,false |
| 3 | -qgmm_mgroup_e2_m24_k64_n48_fp8_e5m2_weight_nd_single_tail_length_l1_2,2,24,64,48,mxfp8_e5m2,false,false,"(ND,ND)",single,false,length,2 | 3 | +qgmm_mgroup_e2_m24_k64_n48_fp8_e5m2_weight_nd_single_tail_length_l1_2,2,24,64,48,mxfp8_e5m2,false,false,"(ND,ND)",single,false,length,24;24,64,64,2,2,1,false |
| 4 | -qgmm_mgroup_e4_m32_k128_n64_fp4_e2m1_weight_nd_single_length_l1_2,4,32,128,64,mxfp4_e2m1,false,false,"(ND,ND)",single,false,length,2 | 4 | +qgmm_mgroup_e4_m32_k128_n64_fp4_e2m1_weight_nd_single_length_l1_2,4,32,128,64,mxfp4_e2m1,false,false,"(ND,ND)",single,false,length,32;32;32;32,64,64,2,2,1,false |
| 5 | -qgmm_mgroup_e2_m16_k64_n64_fp8_e5m2_weight_nd_transpose_single_length_l1_2,2,16,64,64,mxfp8_e5m2,false,true,"(ND,DN)",single,false,length,2 | 5 | +qgmm_mgroup_e2_m16_k64_n64_fp8_e5m2_weight_nd_transpose_single_length_l1_2,2,16,64,64,mxfp8_e5m2,false,true,"(ND,ND)",single,false,length,16;16,64,64,2,2,1,false |
| 6 | -qgmm_mgroup_e2_m24_k64_n48_fp8_e4m3_weight_nz_single_length_l1_2,2,24,64,48,mxfp8_e4m3,false,false,"(ND,NZ)",single,false,length,2 | 6 | +qgmm_mgroup_e2_m24_k64_n48_fp8_e4m3_weight_nz_single_length_l1_2,2,24,64,48,mxfp8_e4m3,false,false,"(ND,NZ)",single,false,length,24;24,64,64,2,2,1,false |
| 7 | -qgmm_mgroup_e2_m16_k64_n64_fp8_e4m3_weight_nz_transpose_single_length_l1_2,2,16,64,64,mxfp8_e4m3,false,true,"(ND,ZN)",single,false,length,2 | 7 | +qgmm_mgroup_e2_m16_k64_n64_fp8_e4m3_weight_nz_transpose_single_length_l1_2,2,16,64,64,mxfp8_e4m3,false,true,"(ND,NZ)",single,false,length,16;16,64,64,2,2,1,false |
| 8 | -qgmm_mgroup_e2_m16_k64_n64_fp4_e2m1_weight_nd_transpose_single_length_l1_2,2,16,64,64,mxfp4_e2m1,false,true,"(ND,DN)",single,false,length,2 | 8 | +qgmm_mgroup_e2_m16_k64_n64_fp4_e2m1_weight_nd_transpose_single_length_l1_2,2,16,64,64,mxfp4_e2m1,false,true,"(ND,ND)",single,false,length,16;16,64,64,2,2,1,false |
| 9 | -qgmm_mgroup_e2_m16_k64_n64_fp4_e2m1_weight_nd_single_bias_offset_l1_2,2,16,64,64,mxfp4_e2m1,false,false,"(ND,ND)",single,true,offset,2 | 9 | +qgmm_mgroup_e2_m16_k64_n64_fp4_e2m1_weight_nd_single_bias_offset_l1_2,2,16,64,64,mxfp4_e2m1,false,false,"(ND,ND)",single,true,offset,16;32,64,64,2,2,1,false |
| 10 | -qgmm_mgroup_e3_m16_k64_n64_fp8_e4m3_weight_nz_multi_length_l1_2,3,16,64,64,mxfp8_e4m3,false,false,"(ND,NZ)",multi,false,length,2 | 10 | +qgmm_mgroup_e3_m16_k64_n64_fp8_e4m3_weight_nz_multi_length_l1_2,3,16,64,64,mxfp8_e4m3,false,false,"(ND,NZ)",multi,false,length,16;16;16,64,64,2,2,1,false |
| 11 | -qgmm_mgroup_e2_m32_k128_n64_fp4_e2m1_weight_nz_multi_length_l1_2,2,32,128,64,mxfp4_e2m1,false,false,"(ND,NZ)",multi,false,length,2 | 11 | +qgmm_mgroup_e2_m32_k128_n64_fp4_e2m1_weight_nz_multi_length_l1_2,2,32,128,64,mxfp4_e2m1,false,false,"(ND,NZ)",multi,false,length,32;32,64,64,2,2,1,false |
| 12 | -qgmm_mgroup_e2_m32_k128_n64_fp4_e1m2_weight_nz_multi_length_l1_2,2,32,128,64,mxfp4_e1m2,false,false,"(ND,NZ)",multi,false,length,2 | 12 | +qgmm_mgroup_e2_m32_k128_n64_fp4_e1m2_weight_nz_multi_length_l1_2,2,32,128,64,mxfp4_e1m2,false,false,"(ND,NZ)",multi,false,length,32;32,64,64,2,2,1,false |
| 13 | -qgmm_mgroup_e2_m16_k64_n64_fp8_e4m3_weight_nd_single_sparse_l1_2,2,16,64,64,mxfp8_e4m3,false,false,"(ND,ND)",single,false,sparse,2 | 13 | +qgmm_mgroup_e2_m16_k64_n64_fp8_e4m3_weight_nd_single_sparse_l1_2,2,16,64,64,mxfp8_e4m3,false,false,"(ND,ND)",single,false,sparse,0;16;1;16,64,64,2,2,1,false |
| 14 | -qgmm_mgroup_e2_m64_k192_n128_fp8_e4m3_weight_nd_single_length_l1_3,2,64,192,128,mxfp8_e4m3,false,false,"(ND,ND)",single,false,length,3 | 14 | +qgmm_mgroup_e2_m64_k192_n128_fp8_e4m3_weight_nd_single_length_l1_3,2,64,192,128,mxfp8_e4m3,false,false,"(ND,ND)",single,false,length,64;64,64,64,2,3,1,false |
| 15 | +qgmm_kgroup_e2_m16_k64_n64_fp8_e4m3_weight_nd_single_length_l1_2,2,16,64,64,mxfp8_e4m3,true,false,"(ND,ND)",single,false,length,29;35,128,128,4,2,2,false | ||
| 16 | +qgmm_kgroup_e2_m16_k64_n64_fp8_e5m2_weight_nd_single_offset_l1_2,2,16,64,64,mxfp8_e5m2,true,false,"(ND,ND)",single,false,offset,29;64,128,128,4,2,2,false | ||
| 17 | +qgmm_mgroup_e3_m17_k96_n33_fp8_e4m3_weight_nd_single_offset_tail,3,17,96,33,mxfp8_e4m3,false,false,"(ND,ND)",single,false,offset,13;31;51,64,64,2,2,1,false | ||
| 18 | +qgmm_mgroup_e3_m32_k128_n80_fp8_e5m2_weight_nd_transpose_single_offset,3,32,128,80,mxfp8_e5m2,false,true,"(ND,ND)",single,false,offset,32;64;96,64,64,2,2,1,false | ||
| 19 | +qgmm_mgroup_e4_m16_k128_n64_fp8_e4m3_weight_nz_multi_offset,4,16,128,64,mxfp8_e4m3,false,false,"(ND,NZ)",multi,false,offset,16;32;48;64,64,64,2,2,1,false | ||
| 20 | +qgmm_mgroup_e2_m48_k192_n128_fp8_e5m2_weight_nz_transpose_single_length,2,48,192,128,mxfp8_e5m2,false,true,"(ND,NZ)",single,false,length,48;48,64,64,2,3,1,false | ||
| 21 | +qgmm_mgroup_e3_m24_k96_n66_fp4_e2m1_weight_nd_single_length_tail,3,24,96,66,mxfp4_e2m1,false,false,"(ND,ND)",single,false,length,24;24;24,64,64,2,2,1,false | ||
| 22 | +qgmm_mgroup_e2_m32_k128_n64_fp4_e2m1_weight_nz_transpose_single_offset,2,32,128,64,mxfp4_e2m1,false,true,"(ND,NZ)",single,false,offset,32;64,64,64,2,2,1,false | ||
| 23 | +qgmm_mgroup_e3_m16_k128_n64_fp4_e1m2_weight_nz_transpose_multi_offset,3,16,128,64,mxfp4_e1m2,false,true,"(ND,NZ)",multi,false,offset,16;32;48,64,64,2,2,1,false | ||
| 24 | +qgmm_mgroup_e4_m16_k64_n64_fp8_e5m2_weight_nz_single_sparse,4,16,64,64,mxfp8_e5m2,false,false,"(ND,NZ)",single,false,sparse,0;16;1;16;2;16;3;16,64,64,2,2,1,false | ||
| 25 | +qgmm_mgroup_e2_m16_k64_n64_fp8_e4m3_weight_nd_single_length_a_full_load,2,16,64,64,mxfp8_e4m3,false,false,"(ND,ND)",single,false,length,16;16,64,64,2,2,1,true | ||
| 26 | +qgmm_mgroup_e2_m64_k256_n128_fp8_e4m3_weight_nd_single_length_large_k,2,64,256,128,mxfp8_e4m3,false,false,"(ND,ND)",single,false,length,64;64,64,64,2,2,1,false | ||
| 27 | +qgmm_kgroup_e3_m17_k128_n33_fp8_e4m3_weight_nd_single_length_tail,3,17,128,33,mxfp8_e4m3,true,false,"(ND,ND)",single,false,length,32;32;64,128,128,4,2,2,false | ||
| 28 | +qgmm_kgroup_e3_m32_k128_n64_fp8_e5m2_weight_nd_single_offset_tail,3,32,128,64,mxfp8_e5m2,true,false,"(ND,ND)",single,false,offset,17;64;128,128,128,4,2,2,false | ||
| @@ -83,15 +83,34 @@ def _format_weight(codes, weight_format, k, n, c0): | |||
| 83 | ) | 83 | ) |
| 84 | 84 | ||
| 85 | 85 | ||
| 86 | -def _make_group_list(group_num, split_size, list_type): | 86 | +def _parse_group_list(text, group_num, total_size, list_type): |
| 87 | - lengths = np.full(group_num, split_size, np.int64) | 87 | + values = np.asarray([int(value) for value in text.split(";")], np.int64) |
| 88 | - if list_type == "offset": | ||
| 89 | - return np.cumsum(lengths, dtype=np.int64) | ||
| 90 | if list_type == "sparse": | 88 | if list_type == "sparse": |
| 91 | - return np.column_stack((np.arange(group_num, dtype=np.int64), lengths)).reshape( | 89 | + if values.size != group_num * 2: |
| 92 | - -1 | 90 | + raise ValueError("sparse group-list must contain e index/length pairs") |
| 93 | - ) | 91 | + pairs = values.reshape(group_num, 2) |
| 94 | - return lengths | 92 | + if sorted(pairs[:, 0].tolist()) != list(range(group_num)): |
| 93 | + raise ValueError( | ||
| 94 | + "sparse group-list indices must be a permutation of [0, e)" | ||
| 95 | + ) | ||
| 96 | + lengths = pairs[:, 1] | ||
| 97 | + elif list_type == "offset": | ||
| 98 | + if ( | ||
| 99 | + values.size != group_num | ||
| 100 | + or np.any(values < 0) | ||
| 101 | + or np.any(np.diff(values) < 0) | ||
| 102 | + ): | ||
| 103 | + raise ValueError("offset group-list must contain e nondecreasing offsets") | ||
| 104 | + lengths = np.diff(np.concatenate((np.zeros(1, np.int64), values))) | ||
| 105 | + else: | ||
| 106 | + if values.size != group_num: | ||
| 107 | + raise ValueError("length group-list must contain e lengths") | ||
| 108 | + lengths = values | ||
| 109 | + if np.any(lengths <= 0): | ||
| 110 | + raise ValueError("group-list lengths must be positive") | ||
| 111 | + if int(lengths.sum()) != total_size: | ||
| 112 | + raise ValueError("group-list must cover the complete grouped axis") | ||
| 113 | + return values, lengths | ||
| 95 | 114 | ||
| 96 | 115 | ||
| 97 | def _write_weights(output_dir, encoded, scales, multi_tensor): | 116 | def _write_weights(output_dir, encoded, scales, multi_tensor): |
| @@ -108,46 +127,96 @@ def generate(args): | |||
| 108 | os.makedirs(args.output_dir, exist_ok=True) | 127 | os.makedirs(args.output_dir, exist_ok=True) |
| 109 | values, codes, c0 = TYPE_INFO[args.dtype] | 128 | values, codes, c0 = TYPE_INFO[args.dtype] |
| 110 | rng = np.random.default_rng(args.seed) | 129 | rng = np.random.default_rng(args.seed) |
| 111 | - a_indices = rng.integers(0, values.size, size=(args.e, args.m, args.k)) | ||
| 112 | - b_indices = rng.integers(0, values.size, size=(args.e, args.k, args.n)) | ||
| 113 | - a_values = values[a_indices] | ||
| 114 | - b_values = values[b_indices] | ||
| 115 | - | ||
| 116 | if args.trans_a: | 130 | if args.trans_a: |
| 117 | - input_a = _encode(codes[a_indices].transpose(0, 2, 1), args.dtype) | 131 | + group_list, k_lengths = _parse_group_list( |
| 118 | - golden = sum(a_values[index] @ b_values[index] for index in range(args.e)) | 132 | + args.group_list, args.e, args.k, args.group_list_type |
| 133 | + ) | ||
| 134 | + a_indices = [ | ||
| 135 | + rng.integers(0, values.size, size=(args.m, int(group_k))) | ||
| 136 | + for group_k in k_lengths | ||
| 137 | + ] | ||
| 138 | + b_indices = [ | ||
| 139 | + rng.integers(0, values.size, size=(int(group_k), args.n)) | ||
| 140 | + for group_k in k_lengths | ||
| 141 | + ] | ||
| 142 | + input_a = _encode( | ||
| 143 | + np.concatenate([codes[index].T.reshape(-1) for index in a_indices]), | ||
| 144 | + args.dtype, | ||
| 145 | + ) | ||
| 146 | + encoded_b = [ | ||
| 147 | + _encode(codes[index].reshape(-1), args.dtype) for index in b_indices | ||
| 148 | + ] | ||
| 149 | + scale_slots = args.k // 64 + args.e | ||
| 150 | + scale_a = np.full((scale_slots, args.m, 2), MX_SCALE_ONE, np.uint8) | ||
| 151 | + scale_b_storage = np.full((scale_slots, args.n, 2), MX_SCALE_ONE, np.uint8) | ||
| 152 | + golden_groups = [] | ||
| 153 | + cumulative_k = 0 | ||
| 154 | + for group_index, (group_k, a_index, b_index) in enumerate( | ||
| 155 | + zip(k_lengths, a_indices, b_indices) | ||
| 156 | + ): | ||
| 157 | + scale_count = (int(group_k) + 31) // 32 | ||
| 158 | + scale_a_codes = rng.integers( | ||
| 159 | + 0x7D, 0x80, size=(args.m, scale_count), dtype=np.uint8 | ||
| 160 | + ) | ||
| 161 | + scale_b_codes = rng.integers( | ||
| 162 | + 0x7D, 0x80, size=(scale_count, args.n), dtype=np.uint8 | ||
| 163 | + ) | ||
| 164 | + scale_start = cumulative_k // 64 + group_index | ||
| 165 | + for scale_index in range(scale_count): | ||
| 166 | + slot = scale_start + scale_index // 2 | ||
| 167 | + lane = scale_index % 2 | ||
| 168 | + scale_a[slot, :, lane] = scale_a_codes[:, scale_index] | ||
| 169 | + scale_b_storage[slot, :, lane] = scale_b_codes[scale_index, :] | ||
| 170 | + a_scale = np.exp2(scale_a_codes.astype(np.int16) - 127) | ||
| 171 | + b_scale = np.exp2(scale_b_codes.astype(np.int16) - 127) | ||
| 172 | + k_scale_index = np.arange(int(group_k)) // 32 | ||
| 173 | + scaled_a = values[a_index] * a_scale[:, k_scale_index] | ||
| 174 | + scaled_b = values[b_index] * b_scale[k_scale_index, :] | ||
| 175 | + golden_groups.append(scaled_a @ scaled_b) | ||
| 176 | + cumulative_k += int(group_k) | ||
| 177 | + golden = np.stack(golden_groups) | ||
| 178 | + scale_a = scale_a.reshape(-1) | ||
| 179 | + scale_b = [scale_b_storage.reshape(-1)] | ||
 | |||
| 119 | else: | 180 | else: |
| 120 | - input_a = _encode(codes[a_indices], args.dtype) | 181 | + group_list, m_groups = _parse_group_list( |
| 182 | + args.group_list, args.e, args.e * args.m, args.group_list_type | ||
| 183 | + ) | ||
| 184 | + a_indices = [ | ||
| 185 | + rng.integers(0, values.size, size=(int(group_m), args.k)) | ||
| 186 | + for group_m in m_groups | ||
| 187 | + ] | ||
| 188 | + b_indices = rng.integers(0, values.size, size=(args.e, args.k, args.n)) | ||
| 189 | + a_values = [values[index] for index in a_indices] | ||
| 190 | + b_values = values[b_indices] | ||
| 191 | + input_a = _encode( | ||
| 192 | + np.concatenate([codes[index].reshape(-1) for index in a_indices]), | ||
| 193 | + args.dtype, | ||
| 194 | + ) | ||
| 121 | golden = np.concatenate( | 195 | golden = np.concatenate( |
| 122 | [a_values[index] @ b_values[index] for index in range(args.e)], axis=0 | 196 | [a_values[index] @ b_values[index] for index in range(args.e)], axis=0 |
| 123 | ) | 197 | ) |
| 124 | - | 198 | + encoded_b = [ |
| 125 | - encoded_b = [ | 199 | + _encode( |
| 126 | - _encode( | 200 | + _format_weight( |
| 127 | - _format_weight( | 201 | + codes[b_indices[index]], args.weight_format, args.k, args.n, c0 |
| 128 | - codes[b_indices[index]], args.weight_format, args.k, args.n, c0 | 202 | + ), |
| 129 | - ), | 203 | + args.dtype, |
| 130 | - args.dtype, | 204 | + ) |
| 131 | - ) | 205 | + for index in range(args.e) |
| 132 | - for index in range(args.e) | 206 | + ] |
| 133 | - ] | 207 | + scale_k = _align_up(args.k, 64) // 32 |
| 134 | - scale_k = _align_up(args.k, 64) // 32 | 208 | + scale_a = np.full(args.e * args.m * scale_k, MX_SCALE_ONE, np.uint8) |
| 135 | - scale_factor = 2 if args.trans_a else 1 | 209 | + scale_b = [ |
| 136 | - scale_a = np.full(args.e * args.m * scale_k * scale_factor, MX_SCALE_ONE, np.uint8) | 210 | + np.full(args.n * scale_k, MX_SCALE_ONE, np.uint8) for _ in range(args.e) |
| 137 | - scale_b = [ | 211 | + ] |
| 138 | - np.full(args.n * scale_k * scale_factor, MX_SCALE_ONE, np.uint8) | ||
| 139 | - for _ in range(args.e) | ||
| 140 | - ] | ||
| 141 | bias = np.full((args.e, args.n), 0.25 if args.with_bias else 0.0, np.float32) | 212 | bias = np.full((args.e, args.n), 0.25 if args.with_bias else 0.0, np.float32) |
| 142 | if args.with_bias: | 213 | if args.with_bias: |
| 143 | - golden += bias[0] if args.trans_a else np.repeat(bias, args.m, axis=0) | 214 | + golden += bias[0] if args.trans_a else np.repeat(bias, m_groups, axis=0) |
| 144 | 215 | ||
| 145 | input_a.tofile(os.path.join(args.output_dir, "input_a.bin")) | 216 | input_a.tofile(os.path.join(args.output_dir, "input_a.bin")) |
| 146 | scale_a.tofile(os.path.join(args.output_dir, "scale_a.bin")) | 217 | scale_a.tofile(os.path.join(args.output_dir, "scale_a.bin")) |
| 147 | bias.tofile(os.path.join(args.output_dir, "bias.bin")) | 218 | bias.tofile(os.path.join(args.output_dir, "bias.bin")) |
| 148 | - _make_group_list( | 219 | + group_list.tofile(os.path.join(args.output_dir, "group_list.bin")) |
| 149 | - args.e, args.k if args.trans_a else args.m, args.group_list_type | ||
| 150 | - ).tofile(os.path.join(args.output_dir, "group_list.bin")) | ||
| 151 | _write_weights(args.output_dir, encoded_b, scale_b, args.multi_tensor) | 220 | _write_weights(args.output_dir, encoded_b, scale_b, args.multi_tensor) |
| 152 | golden.astype(np.float32).tofile(os.path.join(args.output_dir, "golden_c.bin")) | 221 | golden.astype(np.float32).tofile(os.path.join(args.output_dir, "golden_c.bin")) |
| 153 | 222 | ||
| @@ -165,6 +234,7 @@ def main(): | |||
| 165 | parser.add_argument( | 234 | parser.add_argument( |
| 166 | "--group-list-type", choices=("length", "offset", "sparse"), required=True | 235 | "--group-list-type", choices=("length", "offset", "sparse"), required=True |
| 167 | ) | 236 | ) |
| 237 | + parser.add_argument("--group-list", required=True) | ||
| 168 | parser.add_argument("--trans-a", action="store_true") | 238 | parser.add_argument("--trans-a", action="store_true") |
| 169 | parser.add_argument("--multi-tensor", action="store_true") | 239 | parser.add_argument("--multi-tensor", action="store_true") |
| 170 | parser.add_argument("--with-bias", action="store_true") | 240 | parser.add_argument("--with-bias", action="store_true") |
| @@ -22,6 +22,9 @@ def main(): | |||
| 22 | parser.add_argument("actual") | 22 | parser.add_argument("actual") |
| 23 | parser.add_argument("--rtol", type=float, default=1e-3) | 23 | parser.add_argument("--rtol", type=float, default=1e-3) |
| 24 | parser.add_argument("--atol", type=float, default=1e-3) | 24 | parser.add_argument("--atol", type=float, default=1e-3) |
| 25 | + parser.add_argument("--groups", type=int, default=1) | ||
| 26 | + parser.add_argument("--m", type=int) | ||
| 27 | + parser.add_argument("--n", type=int) | ||
| 25 | args = parser.parse_args() | 28 | args = parser.parse_args() |
| 26 | golden = np.fromfile(args.golden, dtype=np.float32) | 29 | golden = np.fromfile(args.golden, dtype=np.float32) |
| 27 | actual = np.fromfile(args.actual, dtype=np.float32) | 30 | actual = np.fromfile(args.actual, dtype=np.float32) |
| @@ -30,6 +33,16 @@ def main(): | |||
| 30 | f"output size mismatch: actual={actual.size}, golden={golden.size}" | 33 | f"output size mismatch: actual={actual.size}, golden={golden.size}" |
| 31 | ) | 34 | ) |
| 32 | close = np.isclose(actual, golden, rtol=args.rtol, atol=args.atol, equal_nan=False) | 35 | close = np.isclose(actual, golden, rtol=args.rtol, atol=args.atol, equal_nan=False) |
| 36 | + if args.m and args.n and actual.size == args.groups * args.m * args.n: | ||
| 37 | + grouped_actual = actual.reshape(args.groups, args.m, args.n) | ||
| 38 | + grouped_golden = golden.reshape(args.groups, args.m, args.n) | ||
| 39 | + for group_index in range(args.groups): | ||
| 40 | + group_error = float( | ||
| 41 | + np.max( | ||
| 42 | + np.abs(grouped_actual[group_index] - grouped_golden[group_index]) | ||
| 43 | + ) | ||
| 44 | + ) | ||
| 45 | + print(f"[INFO] group {group_index}: max_abs_error={group_error}") | ||
| 33 | if not np.all(close): | 46 | if not np.all(close): |
| 34 | mismatch = np.flatnonzero(~close) | 47 | mismatch = np.flatnonzero(~close) |
| 35 | index = int(mismatch[0]) | 48 | index = int(mismatch[0]) |
🟡 Medium Priority
在
trans_a分支中：encoded_b是长度为e（group 数量）的列表，每个元素是一个 k-group 的 weight tensor；当
multi_tensor=True时，_write_weights()（第 98-101 行）使用zip(encoded, scales)并行迭代写入文件。由于zip以较短列表为准，只会迭代 1 次，仅第一个 weight tensor 和唯一的 scale tensor 被写入磁盘，其余e-1个 weight tensor 被静默丢弃。触发条件：同时指定
--trans-a --multi-tensor（e > 1）。失效模式：输出目录中只有
input_b_0.bin和scale_b_0.bin，缺少input_b_1.bin等文件。C++ 算子找不到期望的 weight 文件，测试失败或读取到错误数据。当前影响范围：CSV 中所有
transA=true用例的weight_mode=single，未触发此路径。但命令行直接调用时存在隐患。建议：当 trans_a 且 multi_tensor 时，需要确保 scale_b 列表长度与 encoded_b 一致。最简单的修复：将 scale_b_storage 复制 e 份，使 scale_b = [scale_b_storage.reshape(-1)] * args.e。或者如果语义上 trans_a 的 scale_b 只有一份共享的 scale，则需修改 _write_weights 的逻辑来适配这种不对称情况。