mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-11 17:51:40 +00:00
b0f200713a
feat(ck-tile): add block-scale GEMM operators (aquant, bquant, abquant) (#8519) JIRA ID - AICK-1289 Motivation Adds three new block-scale quantized GEMM operators to the CK Tile Engine for FP8/BF8 inference workloads. Technical Details gemm_aquant: A-matrix quantized GEMM with per-row-group scale tensor [M, K/group_size_k] gemm_bquant: B-matrix quantized GEMM with per-column-group scale tensor [K/group_size_k, N] gemm_abquant: Both A and B quantized with independent group-scale tensors Each operator includes CMakeLists, Python instance builder with tier sampling, C++ benchmark/profiler with host reference verification, and config JSONs. Supporting changes to gemm_instance_builder.py, gemm_validation_utils.py, sampling infra, and the operation support matrix. Test Plan Build and run all three operators with fp8/bf8 on gfx942/gfx950 Verify correctness against CPU reference Verify CI config builds pass
109 lines
2.5 KiB
Python
109 lines
2.5 KiB
Python
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
GEMM_AXES = [
|
|
"tile_m",
|
|
"tile_n",
|
|
"tile_k",
|
|
"warp_m",
|
|
"warp_n",
|
|
"warp_k",
|
|
"warp_tile_m",
|
|
"warp_tile_n",
|
|
"warp_tile_k",
|
|
"pipeline",
|
|
"epilogue",
|
|
"scheduler",
|
|
"pad_m",
|
|
"pad_n",
|
|
"pad_k",
|
|
"persistent",
|
|
]
|
|
|
|
GEMM_STREAMK_AXES = GEMM_AXES + ["reduction_strategy"]
|
|
|
|
GEMM_ABQUANT_AXES = [
|
|
"tile_m",
|
|
"tile_n",
|
|
"tile_k",
|
|
"warp_m",
|
|
"warp_n",
|
|
"warp_k",
|
|
"warp_tile_m",
|
|
"warp_tile_n",
|
|
"warp_tile_k",
|
|
"pipeline",
|
|
"epilogue",
|
|
"scheduler",
|
|
"pad_m",
|
|
"pad_n",
|
|
"pad_k",
|
|
"a_preshuffle_quant",
|
|
"b_preshuffle_quant",
|
|
"group_size_n",
|
|
]
|
|
|
|
CATEGORICAL_AXES = {
|
|
"pipeline",
|
|
"epilogue",
|
|
"scheduler",
|
|
"reduction_strategy",
|
|
"pad_m",
|
|
"pad_n",
|
|
"pad_k",
|
|
"persistent",
|
|
"a_preshuffle_quant",
|
|
"b_preshuffle_quant",
|
|
}
|
|
|
|
|
|
def normalize_axis_values(feasible_set, axes=None):
|
|
"""Compute normalization metadata for each axis.
|
|
|
|
Returns dict mapping axis name to:
|
|
- For numeric axes: {"type": "numeric", "min": v, "max": v, "range": v}
|
|
- For categorical axes: {"type": "categorical", "values": sorted list, "map": value->index}
|
|
"""
|
|
if axes is None:
|
|
axes = GEMM_AXES
|
|
|
|
meta = {}
|
|
for ax in axes:
|
|
values = [item[ax] for item in feasible_set if ax in item]
|
|
if not values:
|
|
continue
|
|
|
|
if ax in CATEGORICAL_AXES:
|
|
unique = sorted(set(str(v) for v in values))
|
|
meta[ax] = {
|
|
"type": "categorical",
|
|
"values": unique,
|
|
"map": {v: i for i, v in enumerate(unique)},
|
|
"count": len(unique),
|
|
}
|
|
else:
|
|
num_values = [float(v) for v in values]
|
|
mn, mx = min(num_values), max(num_values)
|
|
meta[ax] = {
|
|
"type": "numeric",
|
|
"min": mn,
|
|
"max": mx,
|
|
"range": mx - mn if mx != mn else 1.0,
|
|
}
|
|
return meta
|
|
|
|
|
|
def normalize_point(item, axes, meta):
|
|
"""Normalize a single point to [0, 1] per axis."""
|
|
coords = []
|
|
for ax in axes:
|
|
if ax not in meta or ax not in item:
|
|
coords.append(0.0)
|
|
continue
|
|
m = meta[ax]
|
|
if m["type"] == "numeric":
|
|
coords.append((float(item[ax]) - m["min"]) / m["range"])
|
|
else:
|
|
coords.append(m["map"].get(str(item[ax]), 0) / max(m["count"] - 1, 1))
|
|
return coords
|