[Issue]: The `bias` is not used consistently in `gemm_a8w8_bpreshuffle`
@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)
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
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.