ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 01:28:27 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
6667a9021713f794a2c9aee4696c19f6cf376235
composable_kernel/tile_engine/ops/gemm/gemm_validation_utils.py
Thrupti Raj Lakshmana Gowda b0f200713a [rocm-libraries] ROCm/rocm-libraries#8519 (commit 9637390) 
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
2026-07-07 18:22:48 +00:00

1629 lines
47 KiB
Python
Raw Blame History

# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
import logging
from typing import Tuple, List
GEMM_PIPELINES = ["mem", "compv3", "compv4"]
GEMM_PRESHUFFLE_PIPELINES = ["preshufflev2"]
GEMM_ROWCOLQUANT_PIPELINES = ["compv3"]
GEMM_MX_PIPELINES = ["comp_async"]
GEMM_BQUANT_PIPELINES = ["compv3"]
GEMM_ABQUANT_PIPELINES = ["compv3"]
LAYOUT_MAP = {
"r": "ck_tile::tensor_layout::gemm::RowMajor",
"c": "ck_tile::tensor_layout::gemm::ColumnMajor",
}
ELEMENT_SIZE_MAP = {
"fp16": 2,
"bf16": 2,
"int8": 1,
"fp8": 1,
"bf8": 1,
"int4": 0.5,
"fp4": 0.5,
"int32": 4,
"fp32": 4,
"fp64": 8,
}
def get_warp_size_for_gpu(gpu_target: str) -> int:
"""Get the warp size for a given GPU target.
CDNA architectures (gfx9xx) use WAVE64 (64 threads per wavefront).
RDNA architectures (gfx10xx, gfx11xx, gfx12xx) use WAVE32 (32 threads per wavefront).
"""
if gpu_target.startswith("gfx9"):
return 64 # CDNA - WAVE64
return 32 # RDNA and others - WAVE32
WARP_SUPPORTED_COMBINATIONS = {
"gfx90a": [
[1, 4, 1],
[2, 2, 1],
[4, 1, 1],
],
"gfx942": [
[1, 4, 1],
[2, 2, 1],
[4, 1, 1],
],
"gfx950": [
[1, 4, 1],
[2, 2, 1],
[4, 1, 1],
],
"gfx1201": [
[2, 4, 1],
[1, 8, 1],
[8, 1, 1],
[4, 2, 1],
],
}
GEMM_PRESHUFFLE_WARP_TILE_SUPPORTED_COMBINATIONS = {
"gfx90a": {
"fp16_fp16_fp16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"bf16_bf16_bf16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"fp8_fp8_fp16": [[32, 32, 16], [32, 32, 32]],
"bf8_bf8_fp16": [[32, 32, 16], [32, 32, 32]],
},
"gfx942": {
"fp16_fp16_fp16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"bf16_bf16_bf16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"fp8_fp8_fp16": [[32, 32, 16], [32, 32, 32], [16, 16, 32], [16, 16, 64]],
"bf8_bf8_fp16": [[32, 32, 16], [32, 32, 32], [16, 16, 64], [16, 16, 32]],
"int8_int8_int32": [[16, 16, 32], [32, 32, 16]],
},
"gfx950": {
"fp16_fp16_fp16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"bf16_bf16_bf16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"fp8_fp8_fp16": [
[32, 32, 16],
[32, 32, 32],
[16, 16, 32],
[16, 16, 64],
[16, 16, 128],
[32, 32, 64],
],
"bf8_bf8_fp16": [
[32, 32, 16],
[32, 32, 32],
[16, 16, 64],
[16, 16, 32],
[16, 16, 128],
[32, 32, 64],
],
},
}
GEMM_WARP_TILE_SUPPORTED_COMBINATIONS = {
"gfx90a": {
"fp16_fp16_fp16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"bf16_bf16_bf16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"fp8_fp8_fp16": [[32, 32, 16], [32, 32, 32]],
"bf8_bf8_fp16": [[32, 32, 16], [32, 32, 32]],
},
"gfx942": {
"fp16_fp16_fp16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"bf16_bf16_bf16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"fp8_fp8_fp16": [[32, 32, 16], [32, 32, 32], [16, 16, 32], [16, 16, 64]],
"bf8_bf8_fp16": [[32, 32, 16], [32, 32, 32], [16, 16, 64], [16, 16, 32]],
"int8_int8_int32": [[16, 16, 32], [32, 32, 16]],
},
"gfx950": {
"fp16_fp16_fp16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"bf16_bf16_bf16": [
[32, 32, 8],
[16, 16, 16],
[32, 32, 16],
[16, 16, 32],
[64, 4, 16],
],
"fp8_fp8_fp16": [
[32, 32, 16],
[32, 32, 32],
[16, 16, 32],
[16, 16, 64],
[16, 16, 128],
[32, 32, 64],
],
"bf8_bf8_fp16": [
[32, 32, 16],
[32, 32, 32],
[16, 16, 64],
[16, 16, 32],
[16, 16, 128],
[32, 32, 64],
],
},
"gfx1201": { # Check how to handle for GEMM and Multi D
"fp16_fp16_fp16": [
[16, 16, 16],
],
},
}
GEMM_MX_WARP_TILE_SUPPORTED_COMBINATIONS = {
"gfx950": {
"fp4_fp4_fp16": [[16, 16, 128]],
"fp8_fp8_fp16": [[16, 16, 128]],
},
}
TRAIT_UNSUPPORTED_COMBINATIONS = {
("compv3", "cshuffle", "interwave"),
("compv3", "default", "interwave"),
("compv4", "cshuffle", "interwave"),
("compv4", "default", "interwave"),
("comp_async", "default", "intrawave"),
("comp_async", "default", "interwave"),
("comp_async", "cshuffle", "interwave"),
}
AQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
("mem", "default", "interwave"),
("mem", "cshuffle", "interwave"),
("compv3", "default", "interwave"),
("compv3", "cshuffle", "interwave"),
}
BQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
("compv3", "default", "interwave"),
("compv3", "cshuffle", "interwave"),
}
ABQUANT_TRAIT_UNSUPPORTED_COMBINATIONS = {
("compv3", "default", "interwave"),
("compv3", "cshuffle", "interwave"),
}
def element_size(data_type: str) -> float:
"""Calculate the size (in bytes) of a single element for given data type."""
data_type = data_type.lower()
if data_type not in ELEMENT_SIZE_MAP:
raise ValueError(f"Unsupported data type: {data_type}")
return ELEMENT_SIZE_MAP[data_type]
def is_trait_combination_valid(
pipeline: str,
epilogue: str,
scheduler: str,
persistent_or_preshuffle_quant=None,
kernel_name_prefix: str = "",
layout: str = "",
) -> bool:
"""Check if a trait combination is valid."""
if kernel_name_prefix == "gemm_aquant":
if (pipeline, epilogue, scheduler) in AQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
return False
# mem pipeline does not support preshuffle
if pipeline == "mem" and persistent_or_preshuffle_quant is True:
return False
return True
elif kernel_name_prefix == "gemm_bquant":
if (pipeline, epilogue, scheduler) in BQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
return False
# bquant only supports compv3 + intrawave
if pipeline != "compv3" or scheduler != "intrawave":
return False
# BPreshuffleQuant requires ColumnMajor BLayout (second char of layout is 'c')
if persistent_or_preshuffle_quant is True and len(layout) >= 2 and layout[1] != "c":
return False
return True
elif (
kernel_name_prefix == "gemm_rowcolquant"
or kernel_name_prefix == "grouped_gemm_rowcolquant"
or kernel_name_prefix == "grouped_gemm_tensorquant"
):
# rowcolquant and tensorquant only supports compv3 + intrawave + cshuffle
if pipeline != "compv3" or scheduler != "intrawave" or epilogue != "cshuffle":
return False
return True
elif kernel_name_prefix == "gemm_abquant":
if (pipeline, epilogue, scheduler) in ABQUANT_TRAIT_UNSUPPORTED_COMBINATIONS:
return False
# abquant only supports compv3 + intrawave
if pipeline != "compv3" or scheduler != "intrawave":
return False
# BPreshuffleQuant requires ColumnMajor BLayout (second char of layout is 'c')
if persistent_or_preshuffle_quant is True and len(layout) >= 2 and layout[1] != "c":
return False
return True
else:
return (pipeline, epilogue, scheduler) not in TRAIT_UNSUPPORTED_COMBINATIONS
def validate_warp_configuration(
warp_m: int,
warp_n: int,
warp_k: int,
gpu_name: str,
) -> bool:
"""Validate warp configuration."""
current_combination = [warp_m, warp_n, warp_k]
allowed_combinations = WARP_SUPPORTED_COMBINATIONS.get(gpu_name, {})
if not allowed_combinations:
# If GPU not recognized, try to be permissive but log warning
logging.warning(f"No warp_[m/n/k] combinations found for GPU: {gpu_name}")
return True
# Check if current combination is in the allowed list
if current_combination not in allowed_combinations:
return False
return True
def validate_dimension_alignment(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
) -> Tuple[bool, List[str]]:
"""Check if tile dimensions are properly aligned with warp dimensions."""
alignment_issues = []
if tile_m % (warp_m * warp_tile_m) != 0:
alignment_issues.append(
f"tile_m({tile_m}) % [{warp_m}x{warp_tile_m}] = {tile_m % (warp_m * warp_tile_m)}"
)
if tile_n % (warp_n * warp_tile_n) != 0:
alignment_issues.append(
f"tile_n({tile_n}) % [{warp_n}x{warp_tile_n}] = {tile_n % (warp_n * warp_tile_n)}"
)
if tile_k % (warp_k * warp_tile_k) != 0:
alignment_issues.append(
f"tile_k({tile_k}) % [{warp_k}x{warp_tile_k}] = {tile_k % (warp_k * warp_tile_k)}"
)
return len(alignment_issues) == 0, alignment_issues
LDS_SIZE_MAP = {
"gfx90a": 2**16, # 64KB
"gfx942": 2**16, # 64KB
"gfx950": 160 * 1024, # 160KB
"gfx1201": 2**16, # 64KB
}
DEFAULT_LDS_SIZE = 2**16 # 64KB
def validate_lds_capacity(
tile_m: int,
tile_n: int,
tile_k: int,
a_datatype: str,
b_datatype: str,
pipeline: str,
gpu_target: str = "",
) -> Tuple[bool, str]:
"""Validate LDS capacity requirements."""
matrix_a_size = (tile_m * tile_k) * element_size(a_datatype)
matrix_b_size = (tile_n * tile_k) * element_size(b_datatype)
total_tile_in_lds = matrix_a_size + matrix_b_size
base_gpu_target = gpu_target.split(":")[0] if gpu_target else gpu_target
hw_lds_size = LDS_SIZE_MAP.get(base_gpu_target, DEFAULT_LDS_SIZE)
double_buffer = pipeline in ["preshufflev2", "compv4"]
max_tile_size = hw_lds_size // 2 if double_buffer else hw_lds_size
if total_tile_in_lds > max_tile_size:
error_msg = (
f"LDS capacity exceeded: Total required {total_tile_in_lds:,}B ({total_tile_in_lds / 1024:.1f}KB) > "
f"maximum allowed {max_tile_size:,}B ({max_tile_size / 1024}KB) "
f"[{base_gpu_target}, {'double' if double_buffer else 'single'} buffer]. Breakdown:\n"
f"- Matrix A ({a_datatype}): {tile_m}x{tile_k} = {matrix_a_size:,}B\n"
f"- Matrix B ({b_datatype}): {tile_n}x{tile_k} = {matrix_b_size:,}B"
)
return False, error_msg
return True, ""
def validate_gemm_warp_tile_combination(
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
gpu_name: str,
) -> Tuple[bool, str]:
"""Validate warp tile combination against GPU-specific supported combinations."""
# Construct the key for looking up supported combinations
warp_tile_key = f"{a_datatype}_{b_datatype}_{c_datatype}"
current_combination = [warp_tile_m, warp_tile_n, warp_tile_k]
# Check if we have GPU-specific combinations
gpu_warp_tile_combinations = GEMM_WARP_TILE_SUPPORTED_COMBINATIONS.get(gpu_name, {})
if not gpu_warp_tile_combinations:
# If GPU not recognized, try to be permissive but log warning
logging.warning(f"No warp tile combinations found for GPU: {gpu_name}")
return True, ""
# Check if we have combinations for this data type combination
allowed_combinations = gpu_warp_tile_combinations.get(warp_tile_key, [])
if not allowed_combinations:
# For data type combinations not in the list, be permissive
logging.debug(
f"No warp tile combinations found for data types: {warp_tile_key}"
)
return True, ""
# Check if current combination is in the allowed list
if current_combination not in allowed_combinations:
error_msg = (
f"Invalid warp tile combination: {current_combination} not in allowed list. "
f"Valid combinations for '{warp_tile_key}' on {gpu_name}: {allowed_combinations}"
)
return False, error_msg
return True, ""
def validate_gemm_preshuffle_warp_tile_combination(
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
gpu_name: str,
) -> Tuple[bool, str]:
"""Validate warp tile combination against GPU-specific supported combinations."""
# Construct the key for looking up supported combinations
warp_tile_key = f"{a_datatype}_{b_datatype}_{c_datatype}"
current_combination = [warp_tile_m, warp_tile_n, warp_tile_k]
# Check if we have GPU-specific combinations
gpu_warp_tile_combinations = GEMM_PRESHUFFLE_WARP_TILE_SUPPORTED_COMBINATIONS.get(
gpu_name, {}
)
if not gpu_warp_tile_combinations:
# If GPU not recognized, try to be permissive but log warning
logging.warning(f"No warp tile combinations found for GPU: {gpu_name}")
return True, ""
# Check if we have combinations for this data type combination
allowed_combinations = gpu_warp_tile_combinations.get(warp_tile_key, [])
if not allowed_combinations:
# For data type combinations not in the list, be permissive
logging.debug(
f"No warp tile combinations found for data types: {warp_tile_key}"
)
return True, ""
# Check if current combination is in the allowed list
if current_combination not in allowed_combinations:
error_msg = (
f"Invalid warp tile combination: {current_combination} not in allowed list. "
f"Valid combinations for '{warp_tile_key}' on {gpu_name}: {allowed_combinations}"
)
return False, error_msg
return True, ""
def validate_gemm_mx_warp_tile_combination(
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
gpu_name: str,
) -> Tuple[bool, str]:
"""Validate MX GEMM warp tile combinations against supported scaled MFMA shapes."""
warp_tile_key = f"{a_datatype}_{b_datatype}_{c_datatype}"
current_combination = [warp_tile_m, warp_tile_n, warp_tile_k]
gpu_warp_tile_combinations = GEMM_MX_WARP_TILE_SUPPORTED_COMBINATIONS.get(
gpu_name, {}
)
if not gpu_warp_tile_combinations:
logging.warning(f"No MX GEMM warp tile combinations found for GPU: {gpu_name}")
return False, f"MX GEMM is not enabled for GPU: {gpu_name}"
allowed_combinations = gpu_warp_tile_combinations.get(warp_tile_key, [])
if not allowed_combinations:
return (
False,
f"No MX GEMM warp tile combinations found for data types: {warp_tile_key}",
)
if current_combination not in allowed_combinations:
error_msg = (
f"Invalid MX GEMM warp tile combination: {current_combination} not in "
f"allowed list. Valid combinations for '{warp_tile_key}' on {gpu_name}: "
f"{allowed_combinations}"
)
return False, error_msg
return True, ""
def is_tile_config_valid(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
kernel_name_prefix: str = "",
group_size_k: int = 128,
) -> bool:
"""
Comprehensive tile configuration validation.
Returns True if configuration is valid, False otherwise.
"""
# Basic sanity checks
if tile_m <= 0 or tile_n <= 0 or tile_k <= 0:
return False
if warp_m <= 0 or warp_n <= 0 or warp_k <= 0:
return False
if warp_tile_m <= 0 or warp_tile_n <= 0 or warp_tile_k <= 0:
return False
# Check that warp tiles fit within block tiles
if warp_m * warp_tile_m > tile_m:
return False
if warp_n * warp_tile_n > tile_n:
return False
if warp_k * warp_tile_k > tile_k:
return False
# Validate warp configuration
if not validate_warp_configuration(warp_m, warp_n, warp_k, gpu_target):
logging.debug(
f"Invalid warp configuration: warp_m({warp_m}), warp_n({warp_n}), warp_k({warp_k})"
)
return False
# Validate dimension alignment
is_aligned, alignment_issues = validate_dimension_alignment(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
)
if not is_aligned:
logging.debug(
f"Dimension alignment failed: {', '.join(alignment_issues)}. "
f"Tile dimensions {tile_m}x{tile_n}x{tile_k} must be divisible by "
f"[warp]: {warp_m}x{warp_n}x{warp_k} x [warp_tile]: {warp_tile_m}x{warp_tile_n}x{warp_tile_k}"
)
return False
# Validate LDS capacity
lds_valid, lds_error = validate_lds_capacity(
tile_m, tile_n, tile_k, a_datatype, b_datatype, pipeline, gpu_target
)
if not lds_valid:
logging.debug(f"LDS validation failed: {lds_error}")
return False
if pipeline in GEMM_PIPELINES:
gemm_valid, gemm_valid_error = validate_gemm(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
)
if not gemm_valid:
logging.debug(f"GEMM validation failed: {gemm_valid_error}")
return False
# Validate warp tile combination
warp_tile_valid, warp_tile_error = validate_gemm_warp_tile_combination(
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
gpu_target,
)
if not warp_tile_valid:
logging.debug(f"Warp tile validation failed: {warp_tile_error}")
return False
elif pipeline in GEMM_MX_PIPELINES:
mx_valid, mx_error = validate_gemm_mx(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
)
if not mx_valid:
logging.debug(f"MX GEMM validation failed: {mx_error}")
return False
warp_tile_valid, warp_tile_error = validate_gemm_mx_warp_tile_combination(
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
gpu_target,
)
if not warp_tile_valid:
logging.debug(f"MX warp tile validation failed: {warp_tile_error}")
return False
elif pipeline in GEMM_PRESHUFFLE_PIPELINES:
preshuffle_valid, preshuffle_valid_error = validate_gemm_preshuffle(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
)
if not preshuffle_valid:
logging.debug(
f"GEMM Preshuffle validation failed: {preshuffle_valid_error}"
)
return False
# Validate warp tile combination
warp_tile_valid, warp_tile_error = (
validate_gemm_preshuffle_warp_tile_combination(
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
gpu_target,
)
)
if not warp_tile_valid:
logging.debug(f"Warp tile validation failed: {warp_tile_error}")
return False
# Additional operator-specific validation (runs after pipeline validation)
if (
kernel_name_prefix == "gemm_rowcolquant"
or kernel_name_prefix == "grouped_gemm_rowcolquant"
or kernel_name_prefix == "grouped_gemm_tensorquant"
):
rowcol_tensor_quant_valid, rowcol_tensor_quant_valid_error = (
validate_gemm_rowcol_tensor_quant(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
)
)
if not rowcol_tensor_quant_valid:
logging.debug(
f"GEMM RowColQuant/TensorQuant validation failed: {rowcol_tensor_quant_valid_error}"
)
return False
if kernel_name_prefix == "gemm_aquant":
aquant_valid, aquant_valid_error = validate_gemm_aquant(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
group_size_k,
)
if not aquant_valid:
logging.debug(f"GEMM AQuant validation failed: {aquant_valid_error}")
return False
elif kernel_name_prefix == "gemm_bquant":
bquant_valid, bquant_valid_error = validate_gemm_bquant(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
group_size_k,
)
if not bquant_valid:
logging.debug(f"GEMM BQuant validation failed: {bquant_valid_error}")
return False
elif kernel_name_prefix == "gemm_abquant":
abquant_valid, abquant_valid_error = validate_gemm_abquant(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
group_size_k,
)
if not abquant_valid:
logging.debug(f"GEMM ABQuant validation failed: {abquant_valid_error}")
return False
return True
# [TODO] Handle this while moving code to commons Add more datatype to this function if needed
def get_dtype_string(datatype: str) -> str:
"""Get C++ type string for datatype"""
dtype_map = {
"fp16": "ck_tile::fp16_t",
"fp4": "ck_tile::pk_fp4_t",
"fp8": "ck_tile::fp8_t",
"bf8": "ck_tile::bf8_t",
"bf16": "ck_tile::bf16_t",
"fp32": "float",
"fp64": "double",
}
return dtype_map.get(datatype, "float")
def get_abc_layouts(layout_code: str) -> Tuple[str, str, str]:
"""
Return (ALayout, BLayout, CLayout) from a 3-letter code like 'rcr', 'ccr', 'crr', 'rrr'.
"""
code = str(layout_code).strip().lower()
a_layout = LAYOUT_MAP[code[0]]
b_layout = LAYOUT_MAP[code[1]]
c_layout = LAYOUT_MAP[code[2]]
return a_layout, b_layout, c_layout
def get_abcd_layouts(layout_code: str) -> Tuple[str, str, str, List[str]]:
"""
Return (ALayout, BLayout, CLayout) from a 3-letter code like 'rcrr', 'ccrr', 'crrr', 'rrrr'.
"""
code = str(layout_code).strip().lower()
a_layout = LAYOUT_MAP[code[0]]
b_layout = LAYOUT_MAP[code[1]]
c_layout = LAYOUT_MAP[code[2]]
d0_layout = LAYOUT_MAP[code[3]]
d1_layout = LAYOUT_MAP[code[3]]
return a_layout, b_layout, c_layout, [d0_layout, d1_layout]
def validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target: str = "gfx90a",
) -> Tuple[bool, str]:
# Validate whole workgroup cover configuration
warp_size = get_warp_size_for_gpu(gpu_target)
NumWarps = warp_m * warp_n * warp_k
BlockSize = NumWarps * warp_size
XPerTile = 0
YPerTile = 0
vector_load_size = 0
# A matrix validation
if layout[0] == "r":
vector_load_size = get_global_vector_load_size(
BlockSize, tile_k, a_datatype, tile_m, tile_k
)
XPerTile = tile_k
YPerTile = tile_m
elif layout[0] == "c":
vector_load_size = get_global_vector_load_size(
BlockSize, tile_k, a_datatype, tile_m, tile_m
)
# Validate distribution
XPerTile = tile_k
YPerTile = tile_m
wg_cover_core_valid, wg_cover_core_error = wg_cover_core_validation(
XPerTile, YPerTile, BlockSize, vector_load_size, warp_size
)
if not wg_cover_core_valid:
logging.debug(
f"whole workgroup cover failed for Matrix A distribution: {wg_cover_core_error}"
)
return False, wg_cover_core_error
XPerTile = tile_m
YPerTile = tile_k
wg_cover_core_valid, wg_cover_core_error = wg_cover_core_validation(
XPerTile, YPerTile, BlockSize, vector_load_size, warp_size
)
if not wg_cover_core_valid:
logging.debug(
f"whole workgroup cover failed for Matrix A: {wg_cover_core_error}"
)
return False, wg_cover_core_error
# B matrix validation
if layout[1] == "r":
vector_load_size = get_global_vector_load_size(
BlockSize, tile_k, b_datatype, tile_n, tile_n
)
# Validate distribution
XPerTile = tile_k
YPerTile = tile_n
wg_cover_core_valid, wg_cover_core_error = wg_cover_core_validation(
XPerTile, YPerTile, BlockSize, vector_load_size, warp_size
)
if not wg_cover_core_valid:
logging.debug(
f"whole workgroup cover failed for Matrix B distribution: {wg_cover_core_error}"
)
return False, wg_cover_core_error
XPerTile = tile_n
YPerTile = tile_k
elif layout[1] == "c":
XPerTile = tile_k
YPerTile = tile_n
vector_load_size = get_global_vector_load_size(
BlockSize, tile_k, b_datatype, tile_n, tile_k
)
wg_cover_core_valid, wg_cover_core_error = wg_cover_core_validation(
XPerTile, YPerTile, BlockSize, vector_load_size, warp_size
)
if not wg_cover_core_valid:
logging.debug(
f"whole workgroup cover failed for Matrix B: {wg_cover_core_error}"
)
return False, wg_cover_core_error
return True, ""
def wg_cover_core_validation(
XPerTile: int,
YPerTile: int,
BlockSize: int,
vector_load_size: int,
warp_size: int,
) -> Tuple[bool, str]:
if XPerTile % vector_load_size != 0:
return False, "XPerTile is not divisible by vector_load_size"
num_warps = BlockSize / warp_size
LargestVec = (XPerTile * YPerTile) / (num_warps * warp_size)
X1 = LargestVec if vector_load_size > LargestVec else vector_load_size
X0 = XPerTile / X1
Y1 = warp_size // X0
if X0 * Y1 != warp_size:
return False, "X0 * Y1 != warp_size"
return True, ""
def validate_cshuffle_epilogue_distribution(
tile_m: int,
tile_n: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_size: int,
c_datatype: str,
) -> Tuple[bool, str]:
"""
Validate that the CShuffleEpilogue tile distribution pattern is valid.
This mirrors the static_assert in static_encoding_pattern.hpp:
static_assert(X0 * Y1 == warp_size, "X0 * Y1 must cover whole wavefront!");
The CShuffleEpilogue creates a tile_distribution_encoding_pattern_2d<BlockSize, YPerTile, XPerTile, VecSize, thread_raked>
where:
- BlockSize = warp_m * warp_n * warp_k * warp_size
- YPerTile = MPerIterationShuffle (derived from tile_m / (warp_m * warp_tile_m / some_factor))
- XPerTile = NPerIterationShuffle (derived from tile_n)
- VecSize = vector size based on element size (typically 8 for fp16)
The key constraint is that X0 must evenly divide warp_size, where:
- X0 = min(warp_size, XPerTile / X1)
- X1 = min(VecSize, LargestVec)
- LargestVec = (XPerTile * YPerTile) / (num_warps * warp_size)
"""
NumWarps = warp_m * warp_n * warp_k
BlockSize = NumWarps * warp_size
elem_size = ELEMENT_SIZE_MAP.get(c_datatype, 2)
VecSize = 16 // elem_size
XPerTile = tile_n
YPerTile = tile_m // warp_m
if XPerTile <= 0 or YPerTile <= 0:
return (
False,
f"Invalid tile dimensions: XPerTile={XPerTile}, YPerTile={YPerTile}",
)
num_warps = BlockSize // warp_size
if num_warps * warp_size == 0:
return False, "Invalid BlockSize or warp_size"
LargestVec = (XPerTile * YPerTile) // (num_warps * warp_size)
if LargestVec <= 0:
LargestVec = 1
X1 = min(VecSize, LargestVec) if LargestVec > 0 else VecSize
if X1 <= 0:
X1 = 1
X0 = min(warp_size, XPerTile // X1) if X1 > 0 else warp_size
Y1 = warp_size // X0 if X0 > 0 else 0
if X0 * Y1 != warp_size:
return (
False,
f"CShuffleEpilogue distribution invalid: X0({X0}) * Y1({Y1}) = {X0 * Y1} != warp_size({warp_size}). "
f"XPerTile={XPerTile}, YPerTile={YPerTile}, VecSize={VecSize}, BlockSize={BlockSize}",
)
return True, ""
def get_global_vector_load_size(
BlockSize: int,
KPerBlock: int,
DataType: str,
MNPerBlock: int,
XPerTile: int,
) -> int:
elements_per_thread = MNPerBlock * KPerBlock / BlockSize
PackedSize = 1
if (
PackedSize == 2
and XPerTile % (PackedSize * 32 / element_size(DataType)) == 0
and elements_per_thread % (PackedSize * 32 / element_size(DataType)) == 0
):
return PackedSize * 32 / element_size(DataType)
elif (
XPerTile % (PackedSize * 16 / element_size(DataType)) == 0
and elements_per_thread % (PackedSize * 16 / element_size(DataType)) == 0
):
return int(PackedSize * 16 / element_size(DataType))
elif (
XPerTile % (PackedSize * 8 / element_size(DataType)) == 0
and elements_per_thread % (PackedSize * 8 / element_size(DataType)) == 0
):
return int(PackedSize * 8 / element_size(DataType))
elif (
element_size(DataType) >= PackedSize * 4
and XPerTile % (PackedSize * 4 / element_size(DataType)) == 0
and elements_per_thread % (PackedSize * 4 / element_size(DataType)) == 0
):
return int(PackedSize * 4 / element_size(DataType))
elif (
element_size(DataType) >= PackedSize * 2
and XPerTile % (PackedSize * 2 / element_size(DataType)) == 0
and elements_per_thread % (PackedSize * 2 / element_size(DataType)) == 0
):
return int(PackedSize * 2 / element_size(DataType))
else:
return PackedSize
def validate_gemm(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
trait_name: str = None,
) -> Tuple[bool, str]:
# GEMM Validation
warp_size = get_warp_size_for_gpu(gpu_target)
# Validate whole workgroup cover configuration
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target,
)
)
if not whole_workgroup_cover_valid:
logging.debug(
f"Whole workgroup cover configuration validation failed: {whole_workgroup_cover_error}"
)
return False, whole_workgroup_cover_error
# Validate CShuffleEpilogue distribution pattern (for cshuffle epilogue)
# This validation ensures the tile distribution pattern is valid for the output tile
cshuffle_valid, cshuffle_error = validate_cshuffle_epilogue_distribution(
tile_m,
tile_n,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_size,
c_datatype,
)
if not cshuffle_valid:
logging.debug(f"CShuffleEpilogue validation failed: {cshuffle_error}")
return False, cshuffle_error
return True, ""
def validate_gemm_mx(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
trait_name: str = None,
) -> Tuple[bool, str]:
# MX GEMM uses the scaled MFMA path from ck_tile example 42.
if layout.lower() != "rcr":
return False, "MX GEMM currently supports only rcr layout"
if a_datatype not in ["fp4", "fp8"] or b_datatype not in ["fp4", "fp8"]:
return False, "MX GEMM currently supports fp4 and fp8 inputs only"
if c_datatype != "fp16":
return False, "MX GEMM currently expects fp16 output"
if tile_k % 32 != 0 or warp_tile_k % 32 != 0:
return False, "MX GEMM tile K and warp tile K must be multiples of 32"
warp_size = get_warp_size_for_gpu(gpu_target)
if warp_size != 64:
return False, "MX GEMM scaled MFMA path currently requires wave64"
k_lane_m = warp_size // warp_tile_m if warp_tile_m != 0 else 0
k_lane_n = warp_size // warp_tile_n if warp_tile_n != 0 else 0
if k_lane_m == 0 or k_lane_n == 0:
return False, "Invalid MX GEMM warp tile M/N for wave size"
if warp_tile_k // 32 // k_lane_m == 0 or warp_tile_k // 32 // k_lane_n == 0:
return False, "MX GEMM scale distribution requires non-zero K per lane"
return validate_gemm(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
warp_tile_m,
warp_tile_n,
warp_tile_k,
a_datatype,
b_datatype,
c_datatype,
pipeline,
layout,
gpu_target,
trait_name,
)
def validate_gemm_preshuffle(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
trait_name: str = None,
) -> Tuple[bool, str]:
# Preshuffle Validations
warp_size = get_warp_size_for_gpu(gpu_target)
# Validate vector load alignment
m_iter_per_warp = tile_m / (warp_m * warp_tile_m)
vector_valid, vector_error = validate_vector_load_alignment(
warp_tile_m,
warp_tile_k,
a_datatype,
m_iter_per_warp,
wave_size=warp_size,
vector_load_size=16,
)
if not vector_valid:
logging.debug(f"Vector load alignment failed: {vector_error}")
return False, "vector load alignment error"
# Validate M0, M1, M2 configuration for matrix A row-major layout
m0_m1_m2_valid, m0_m1_m2_error = validate_m0_m1_m2_configuration(
tile_m,
tile_k,
warp_m,
warp_n,
warp_k,
a_datatype,
vector_load_size=16,
warp_size=warp_size,
)
if not m0_m1_m2_valid:
logging.debug(f"M0/M1/M2 configuration validation failed: {m0_m1_m2_error}")
return False, m0_m1_m2_error
return True, ""
def validate_vector_load_alignment(
wg_m: int,
wg_k: int,
a_datatype: str,
m_iter_per_warp: int,
wave_size: int,
vector_load_size: int,
) -> Tuple[bool, str]:
try:
# Calculate the memory access pattern size
a_element_size = element_size(a_datatype)
access_size = (wg_m * wg_k * a_element_size * m_iter_per_warp) / wave_size
# Check if it's aligned to vector load size
if access_size % vector_load_size != 0:
error_msg = (
f"Vector load alignment violation: "
f"({wg_m} * {wg_k} * {a_element_size} * {m_iter_per_warp} / {wave_size}) "
f"% {vector_load_size} = {access_size % vector_load_size} != 0. "
f"Access size: {access_size} bytes"
)
return False, error_msg
return True, ""
except Exception as e:
return False, f"Error in vector load validation: {str(e)}"
def validate_m0_m1_m2_configuration(
tile_m: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
a_datatype: str,
vector_load_size: int = 16,
warp_size: int = 64,
) -> Tuple[bool, str]:
"""
Validate M0, M1, M2 configuration for matrix A row-major layout.
This ensures proper memory access pattern alignment.
"""
try:
# Validation for A as row-major
MPerBlock = tile_m
# Calculate K1 using element size
K1 = vector_load_size / element_size(a_datatype)
# Check if K1 is valid (must be integer)
if K1 != int(K1):
return (
False,
f"K1 = {K1} is not an integer. vector_load_size({vector_load_size}) must be divisible by element_size({a_datatype})",
)
K1 = int(K1)
# Calculate K0
if tile_k % K1 != 0:
return False, f"tile_k({tile_k}) must be divisible by K1({K1})"
K0 = tile_k // K1
# Calculate M2
if warp_size % K0 != 0:
return False, f"warp_size({warp_size}) must be divisible by K0({K0})"
M2 = warp_size // K0
# Calculate number of warps and block size
NumWarps = warp_m * warp_n * warp_k
BlockSize = NumWarps * warp_size
# Calculate M0 (assuming get_warp_size() returns warp_size)
M0 = BlockSize // warp_size # This should equal NumWarps
# Calculate M1
if (M2 * M0) == 0:
return False, f"M2({M2}) * M0({M0}) cannot be zero"
if MPerBlock % (M2 * M0) != 0:
return (
False,
f"MPerBlock({MPerBlock}) must be divisible by M2({M2}) * M0({M0}) = {M2 * M0}",
)
M1 = MPerBlock // (M2 * M0)
# Validate the assertion: M0 * M1 * M2 == MPerBlock
calculated_m_per_block = M0 * M1 * M2
if calculated_m_per_block != MPerBlock:
error_msg = (
f"Incorrect M0, M1, M2 configuration! "
f"M0({M0}) * M1({M1}) * M2({M2}) = {calculated_m_per_block} != MPerBlock({MPerBlock}). "
f"Configuration: K0={K0}, K1={K1}, NumWarps={NumWarps}, BlockSize={BlockSize}"
)
return False, error_msg
return True, ""
except ZeroDivisionError as e:
return False, f"Division by zero in M0/M1/M2 calculation: {str(e)}"
except Exception as e:
return False, f"Error in M0/M1/M2 validation: {str(e)}"
def _validate_fp8_mfma_warp_tile_k(
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
gpu_target: str,
op_label: str = "",
) -> Tuple[bool, str]:
"""Validate MFMA warp-tile constraints shared by all fp8/bf8 quantized GEMM ops.
Checks:
- warp_tile_m == warp_tile_n (square MFMA requirement)
- warp_tile_k matches the ISA-mandated K-block for the given warp_tile_m and gpu_target
"""
suffix = f" ({op_label})" if op_label else ""
if warp_tile_m != warp_tile_n:
return False, (
f"warp_tile_m({warp_tile_m}) must equal warp_tile_n({warp_tile_n})"
f" — MFMA requires a square warp tile{suffix}"
)
if a_datatype in ["fp8", "bf8"]:
# MFMA instruction shapes for fp8/bf8 (from CDNA ISA reference manual):
# gfx90a/gfx942: MFMA_F32_16x16x128_F8 (warp_tile_m=16) → warp_tile_k=64
# MFMA_F32_32x32x64_F8 (warp_tile_m=32) → warp_tile_k=32
# gfx950 doubles the K-block:
# MFMA_F32_16x16x256_F8 (warp_tile_m=16) → warp_tile_k=128
# MFMA_F32_32x32x128_F8 (warp_tile_m=32) → warp_tile_k=64
if gpu_target == "gfx950":
expected_k = 64 if warp_tile_m == 32 else 128
else:
expected_k = 32 if warp_tile_m == 32 else 64
if warp_tile_k != expected_k:
return False, (
f"For {a_datatype} on {gpu_target}, warp_tile_m={warp_tile_m} "
f"requires warp_tile_k={expected_k}, got warp_tile_k={warp_tile_k}{suffix}"
)
return True, ""
def validate_gemm_rowcol_tensor_quant(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
) -> Tuple[bool, str]:
"""Validate RowColQuant / TensorQuant GEMM-specific constraints."""
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target,
)
)
if not whole_workgroup_cover_valid:
return False, whole_workgroup_cover_error
return _validate_fp8_mfma_warp_tile_k(
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
op_label="RowColQuant / TensorQuant"
)
def validate_gemm_aquant(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
group_size_k: int = 128,
) -> Tuple[bool, str]:
"""Validate AQuant GEMM-specific constraints."""
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target,
)
)
if not whole_workgroup_cover_valid:
logging.debug(
f"Whole workgroup cover configuration validation failed: {whole_workgroup_cover_error}"
)
return False, whole_workgroup_cover_error
if tile_k % group_size_k != 0 or tile_k < group_size_k:
return False, (
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
f"and tile_k >= group_size_k"
)
if group_size_k % warp_tile_k != 0:
return False, (
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
)
return _validate_fp8_mfma_warp_tile_k(
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
op_label="AQuant"
)
def validate_gemm_bquant(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
group_size_k: int = 128,
) -> Tuple[bool, str]:
"""Validate BQuant GEMM-specific constraints."""
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target,
)
)
if not whole_workgroup_cover_valid:
return False, whole_workgroup_cover_error
if tile_k % group_size_k != 0 or tile_k < group_size_k:
return False, (
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
f"and tile_k >= group_size_k"
)
if group_size_k % warp_tile_k != 0:
return False, (
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
)
return _validate_fp8_mfma_warp_tile_k(
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
op_label="BQuant"
)
def validate_gemm_abquant(
tile_m: int,
tile_n: int,
tile_k: int,
warp_m: int,
warp_n: int,
warp_k: int,
warp_tile_m: int,
warp_tile_n: int,
warp_tile_k: int,
a_datatype: str,
b_datatype: str,
c_datatype: str,
pipeline: str,
layout: str,
gpu_target: str,
group_size_k: int = 128,
) -> Tuple[bool, str]:
"""Validate ABQuant GEMM-specific constraints."""
whole_workgroup_cover_valid, whole_workgroup_cover_error = (
validate_whole_wg_cover_configuration(
tile_m,
tile_n,
tile_k,
warp_m,
warp_n,
warp_k,
layout,
a_datatype,
b_datatype,
gpu_target,
)
)
if not whole_workgroup_cover_valid:
return False, whole_workgroup_cover_error
if tile_k % group_size_k != 0 or tile_k < group_size_k:
return False, (
f"tile_k({tile_k}) must be a multiple of group_size_k({group_size_k}) "
f"and tile_k >= group_size_k"
)
if group_size_k % warp_tile_k != 0:
return False, (
f"group_size_k({group_size_k}) must be divisible by warp_tile_k({warp_tile_k})"
)
return _validate_fp8_mfma_warp_tile_k(
warp_tile_m, warp_tile_n, warp_tile_k, a_datatype, gpu_target,
op_label="ABQuant"
)
Reference in New Issue View Git Blame Copy Permalink