mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-09 08:37:03 +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
12 lines
618 B
Python
12 lines
618 B
Python
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
from sampling.sampler import sample_feasible_set as sample_feasible_set
|
|
from sampling.seed import make_seed as make_seed
|
|
from sampling.budget import allocate_budget as allocate_budget
|
|
from sampling.budget import load_op_weights as load_op_weights
|
|
from sampling.manifest import write_manifest as write_manifest
|
|
from sampling.feasible_set import GEMM_AXES as GEMM_AXES
|
|
from sampling.feasible_set import GEMM_ABQUANT_AXES as GEMM_ABQUANT_AXES
|
|
from sampling.feasible_set import normalize_axis_values as normalize_axis_values
|