ROCm / ROCm/aiter

[Issue]: The `bias` is not used consistently in `gemm_a8w8_bpreshuffle`

Open
#670 5 comments 0 reactions 1 assignee View on GitHub

@junhaha666 is already working on this.

Since Jul 18, 2025.

Dominant language
Python
Stars
565
Forks
585
Avg merge
3d 4h
Merged PRs (30d)
366

Description

Problem Description

The gemm_a8w8_bpreshuffle changed in this commit caf1e1914c28efac71697d2046bd9ebef7b7ec6a. it changed the API of gemm_a8w8_bpreshuffle and introduced bias term.
The bias is not used consistently.
bias is only used with asm kernel gemm_a8w8_ASM but not ck kernels, this makes the API interface usage confusing. Since gemm_a8w8_bpreshuffle is a tuning API interface, it is hard to know if when bias is passed into this API, will the bias be applied.

def gemm_a8w8_bpreshuffle(
    XQ: Tensor,
    WQ: Tensor,
    x_scale: Tensor,
    w_scale: Tensor,
    bias: Optional[Tensor] = None,
    dtype=torch.float16,
    check=False,
):
    assert dtype in [
        torch.bfloat16,
        torch.float16,
    ], f"Output {dtype=} is currently not supported in gemm_a8w8"
    m = XQ.shape[0]
    n = WQ.shape[0]
    k = XQ.shape[-1]

    ck_config = get_CKGEMM_config(m, n, k, "a8w8_bpreshuffle_tuned_gemm.csv")
    if (
        ck_config is None
        and dtype == dtypes.bf16
        and bias is not None
        and WQ.dtype != dtypes.i8
    ):
        res = gemm_a8w8_ASM(XQ, WQ, x_scale, w_scale, bias, dtype=dtype, check=check)
        if res is not None:
            return res
    Y = torch.empty(m, n, dtype=dtype, device=XQ.device)
    return gemm_a8w8_bpreshuffle_ck(XQ, WQ, x_scale, w_scale, Y)

https://github.com/ROCm/aiter/blob/2b083e30d524f0bbd4921e2922103322ac3efd04/aiter/ops/gemm_op_a8w8.py#L214

Operating System

Ubuntu 22.04.4 LTS (Jammy Jellyfish)

CPU

AMD EPYC 9654 96-Core Processor

GPU

AMD Instinct MI300X

ROCm Version

ROCm 6.4.1

ROCm Component

No response

Steps to Reproduce

No response

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.