Skip to content

[Triton/Gluon] Move gluon gemm_a8w8_blockscale kernel into _gluon_kernels #4917

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vgokhale merged 1 commit into from
Aug 25, 2026
Merged

[Triton/Gluon] Move gluon gemm_a8w8_blockscale kernel into _gluon_kernels #4917

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion aiter/ops/triton/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@
which depend on the old module names, can still import it the old "way" of importing.
"""
# Paths that only exist for backward compatibility and are on their way out.
_DEPRECATED_COMPAT_PATHS = ("gluon.gemm_a8w8",)
_DEPRECATED_COMPAT_PATHS = ("gluon.gemm_a8w8", "gluon.gemm_a8w8_blockscale")


def _warn_if_deprecated(name, new_path):
Expand Down Expand Up @@ -87,6 +87,7 @@ def _warn_if_deprecated(name, new_path):
"gemm_a8w8_per_token_scale": "gemm.basic.gemm_a8w8_per_token_scale",
"gemm_a8w8": "gemm.basic.gemm_a8w8",
"gluon.gemm_a8w8": "gemm.basic.gemm_a8w8",
"gluon.gemm_a8w8_blockscale": "gemm.basic.gemm_a8w8_blockscale",
"gemm_a8wfp4": "gemm.basic.gemm_a8wfp4",
"gemm_afp4wfp4_pre_quant_atomic": "gemm.basic.gemm_afp4wfp4_pre_quant_atomic",
"gemm_afp4wfp4": "gemm.basic.gemm_afp4wfp4",
Expand Down
286 changes: 1 addition & 285 deletions .../ops/triton/gluon/gemm_a8w8_blockscale.py → ...gfx950/gemm/basic/gemm_a8w8_blockscale.py
View file
Open in desktop
Comment thread
Boss2002n marked this conversation as resolved.
Comment thread
Boss2002n marked this conversation as resolved.
Original file line number Diff line number Diff line change
@@ -1,24 +1,15 @@
# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

import functools
"""Gluon FP8 block-scale GEMM kernels for gfx950."""

import torch
import triton
from triton import language as tl
from triton.experimental import gluon
from triton.experimental.gluon import language as gl
from triton.runtime.jit import constexpr_function

from aiter.ops.triton.utils._triton import arch_info
from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid, remap_xcd
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
from aiter.ops.triton.utils.logger import AiterTritonLogger

_LOGGER = AiterTritonLogger()


# Supported (BLOCK_M, BLOCK_N) tiles; BLOCK_K=128 and NUM_WARPS=4 are baked in.
_SUPPORTED_TILES = ((64, 128), (128, 128), (128, 256))


Expand Down Expand Up @@ -893,278 +884,3 @@ def _gemm_a8w8_blockscale_kernel(
NEED_N_MASK=NEED_N_MASK,
NUM_WARPS=NUM_WARPS,
)


@gluon.jit
def _gemm_a8w8_blockscale_reduce_kernel(
c_in_ptr,
c_out_ptr,
M,
N,
stride_c_in_k,
stride_c_in_m,
stride_c_in_n,
stride_c_out_m,
stride_c_out_n,
BLOCK_SIZE_M: gl.constexpr, # Note: Can be distinct from GEMM block size
BLOCK_SIZE_N: gl.constexpr,
ACTUAL_KSPLIT: gl.constexpr,
MAX_KSPLIT: gl.constexpr,
):

pid_m = gl.program_id(axis=0)
pid_n = gl.program_id(axis=1)

blocked_read: gl.constexpr = gl.BlockedLayout( # (MAX_KSPLIT, BLOCK_M, BLOCK_N)
size_per_thread=[1, 1, 4],
threads_per_warp=[1, 8, 8],
warps_per_cta=[1, 4, 1],
order=[2, 1, 0],
)

# blocked_write: gl.constexpr = gl.BlockedLayout(
# size_per_thread=[1, 4], # (BLOCK_M, BLOCK_N)
# threads_per_warp=[8, 8],
# warps_per_cta=[4, 1],
# order=[1, 0],
# )

offs_m = pid_m * BLOCK_SIZE_M + gl.arange(
0,
BLOCK_SIZE_M, # keep dim 1
gl.SliceLayout(0, gl.SliceLayout(2, blocked_read)),
)
offs_n = pid_n * BLOCK_SIZE_N + gl.arange(
0,
BLOCK_SIZE_N, # keep dim 2
gl.SliceLayout(0, gl.SliceLayout(1, blocked_read)),
)
offs_k = gl.arange(
0, MAX_KSPLIT, gl.SliceLayout(1, gl.SliceLayout(2, blocked_read)) # keep dim 0
)
c_in_offs = (
(offs_k[:, None, None] * stride_c_in_k)
+ (offs_m[None, :, None] * stride_c_in_m)
+ (offs_n[None, None, :] * stride_c_in_n)
)
if ACTUAL_KSPLIT == MAX_KSPLIT:
c_in_mask = (offs_m[None, :, None] < M) & (offs_n[None, None, :] < N)
c = gl.amd.cdna4.buffer_load(c_in_ptr, c_in_offs, mask=c_in_mask, cache=".ca")
else:
c_in_mask = (
(offs_m[None, :, None] < M)
& (offs_n[None, None, :] < N)
& (offs_k[:, None, None] < ACTUAL_KSPLIT)
)
c = gl.amd.cdna4.buffer_load(
c_in_ptr, c_in_offs, mask=c_in_mask, cache=".ca"
) # , other=0.0)
c = tl.sum(c, 0)

c = c.to(c_out_ptr.type.element_ty)

offs_cm = pid_m * BLOCK_SIZE_M + gl.arange(
0, BLOCK_SIZE_M, gl.SliceLayout(1, gl.SliceLayout(0, blocked_read))
)
offs_cn = pid_n * BLOCK_SIZE_N + gl.arange(
0, BLOCK_SIZE_N, gl.SliceLayout(0, gl.SliceLayout(0, blocked_read))
)
c_out_offs = (offs_cm[:, None] * stride_c_out_m) + (
offs_cn[None, :] * stride_c_out_n
)
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)

gl.amd.cdna4.buffer_store(
stored_value=c, ptr=c_out_ptr, offsets=c_out_offs, mask=c_mask
)


@functools.lru_cache(maxsize=1024)
def _get_config_cached(
M: int,
N: int,
K: int,
):
if not arch_info.is_gluon_avail():
raise ValueError(
"Gluon implementation is not supported on this device (requires CDNA4)."
)

dev = arch_info.get_arch()

# Try specialized config first.
config_dict = load_config_json(
f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE-N={N}-K={K}.json",
required=False,
)
# Fall back to the general config (must exist).
if config_dict is None:
config_dict = load_config_json(
f"{AITER_TRITON_CONFIGS_PATH}/gemm/gluon/{dev}-GEMM-A8W8_BLOCKSCALE.json"
)

# Config keys should be named M_LEQ_<bound> or "any"
bounds = []
for setting in config_dict:
potential_block_m = setting.replace("M_LEQ_", "")
if potential_block_m.isnumeric():
bounds.append(int(potential_block_m))

# Walk buckets in ascending-M order; pick the smallest one whose tile
# the kernel currently supports. Unsupported buckets are skipped (those
# configs become live again once the kernel grows the corresponding
# padded-LDS layouts), so we may fall through to "any".
config = config_dict["any"]
for bound in sorted(bounds):
if M > bound or f"M_LEQ_{bound}" not in config_dict:
continue
candidate = config_dict[f"M_LEQ_{bound}"]
if (candidate["BLOCK_SIZE_M"], candidate["BLOCK_SIZE_N"]) in _SUPPORTED_TILES:
config = candidate
break

return config


def _get_config(
M: int,
N: int,
K: int,
):
# Fresh copy per call, outside the lru boundary — the caller writes
# derived fields (SPLITK_BLOCK_SIZE here, GROUP_K/GROUP_N at the call
# site) into the returned dict.
config = _get_config_cached(M, N, K).copy()

block_size_k = config["BLOCK_SIZE_K"]
num_k_blocks = triton.cdiv(K, block_size_k)
num_k_blocks_per_split = triton.cdiv(num_k_blocks, config["NUM_KSPLIT"])
config["SPLITK_BLOCK_SIZE"] = num_k_blocks_per_split * block_size_k

return config


def gemm_a8w8_blockscale(
x: torch.Tensor,
w: torch.Tensor,
x_scale: torch.Tensor,
w_scale: torch.Tensor,
dtype: float | None = torch.bfloat16,
y: torch.Tensor | None = None,
config: dict | None = None,
):
"""
Computes the 8 bit matmul Y = X x WT using the block-scale quantization approach.

Key parameters:
- X: Matrix X with shape (M, K).
- W: Matrix W with shape (N, K).
- X_scale: Scale tensor for X with shape (M, *scale_k).
- W_scale: Scale tensor for W with shape (**scale_n, *scale_k).

Returns:
- Y: The output matrix with shape (M, N).

*scale_k = (K + scale_block_size_k - 1) // scale_block_size_k
**scale_n = (N + scale_block_size_n - 1) // scale_block_size_n
"""
_LOGGER.info(
f"GEMM_A8W8_BLOCKSCALE: x={tuple(x.shape)} w={tuple(w.shape)} x_scale={tuple(x_scale.shape)} w_scale={tuple(w_scale.shape)}"
)

M, K = x.shape
N, K = w.shape

# Check constraints.
assert x.shape[1] == w.shape[1], "Incompatible dimensions!!!"

# Transpose w and w_scale
w = w.T
w_scale = w_scale.T

if y is None:
y = torch.empty((M, N), dtype=dtype, device=x.device)

if config is None:
config = _get_config(M, N, K)

# Scale block sizes
# TODO: need a better way to pass scale block sizes around
config["GROUP_K"] = triton.next_power_of_2(triton.cdiv(K, w_scale.shape[0]))
config["GROUP_N"] = triton.next_power_of_2(triton.cdiv(N, w_scale.shape[1]))

if config["NUM_KSPLIT"] == 1:
assert (
config["GROUP_K"] == config["BLOCK_SIZE_K"]
), f"GROUP_K: {config['GROUP_K']} must equal BLOCK_SIZE_K: {config['BLOCK_SIZE_K']} when not using KSPLIT"

if config["NUM_KSPLIT"] > 1:
y_pp = torch.empty(
(config["NUM_KSPLIT"], M, N), dtype=torch.float32, device=y.device
)
else:
y_pp = None

num_stages = config.get("num_stages", 2)
num_stages = max(num_stages, 2)

# grid = (config["NUM_KSPLIT"], triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]),)
grid = lambda META: (
(
META["NUM_KSPLIT"]
* triton.cdiv(M, META["BLOCK_SIZE_M"])
* triton.cdiv(N, META["BLOCK_SIZE_N"])
),
)
_gemm_a8w8_blockscale_kernel[grid](
x,
w,
y if config["NUM_KSPLIT"] == 1 else y_pp,
x_scale,
w_scale,
M,
N,
K,
x.stride(0),
x.stride(1),
w.stride(0),
w.stride(1),
0 if config["NUM_KSPLIT"] == 1 else y_pp.stride(0),
y.stride(0) if config["NUM_KSPLIT"] == 1 else y_pp.stride(1),
y.stride(1) if config["NUM_KSPLIT"] == 1 else y_pp.stride(2),
x_scale.stride(0),
x_scale.stride(1),
w_scale.stride(0),
w_scale.stride(1),
NUM_WARPS=config["num_warps"],
NUM_STAGES=num_stages,
**config,
)

if config["NUM_KSPLIT"] > 1:
REDUCE_BLOCK_SIZE_M = 32
REDUCE_BLOCK_SIZE_N = 32
ACTUAL_KSPLIT = triton.cdiv(K, config["SPLITK_BLOCK_SIZE"])

grid_reduce = (
triton.cdiv(M, REDUCE_BLOCK_SIZE_M),
triton.cdiv(N, REDUCE_BLOCK_SIZE_N),
)

_gemm_a8w8_blockscale_reduce_kernel[grid_reduce](
y_pp,
y,
M,
N,
y_pp.stride(0),
y_pp.stride(1),
y_pp.stride(2),
y.stride(0),
y.stride(1),
REDUCE_BLOCK_SIZE_M,
REDUCE_BLOCK_SIZE_N,
ACTUAL_KSPLIT,
triton.next_power_of_2(config["NUM_KSPLIT"]),
)

return y
4 changes: 1 addition & 3 deletions aiter/ops/triton/configs/CLAUDE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,9 +123,7 @@ Consequences to keep in mind:
instead).

Direct-path loaders bypass the resolver's directory probe. Grep for
`f"{AITER_TRITON_CONFIGS_PATH}/..."` before moving anything —
`gluon/gemm_a8w8_blockscale.py` still builds legacy `gemm/gluon/` paths by
hand (via `load_config_json`) and must be edited when its configs move.
`f"{AITER_TRITON_CONFIGS_PATH}/..."` before moving anything.
`gluon/gemm_afp4wfp4.py` goes through `get_gemm_config(backend="gluon")` and
needs no changes.

Expand Down
48 changes: 0 additions & 48 deletions aiter/ops/triton/configs/gemm/gluon/gfx950-GEMM-A8W8_BLOCKSCALE-N=2112-K=7168.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,52 +1,4 @@
{
"M_LEQ_16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 2,
"num_stages": 1,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"cache_modifier": ".cg",
"NUM_KSPLIT": 14
},
"M_LEQ_32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 4,
"num_warps": 2,
"num_stages": 2,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"cache_modifier": null,
"NUM_KSPLIT": 7
},
"M_LEQ_64": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 4,
"num_warps": 2,
"num_stages": 1,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"cache_modifier": null,
"NUM_KSPLIT": 7
},
"M_LEQ_128": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 2,
"num_stages": 2,
"waves_per_eu": 1,
"matrix_instr_nonkdim": 16,
"cache_modifier": ".ca",
"NUM_KSPLIT": 7
},
"M_LEQ_2048": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
Expand Down
Loading