Skip to content

[PerfXLab] Add and optimize gemm op #2220

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

Open
bin913 wants to merge 3 commits into
base: master
Choose a base branch
from
Open

[PerfXLab] Add and optimize gemm op #2220

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3b29e6e
optimize gemm
bin913 Mar 29, 2026
a39dec1
add alpha and beta in test
bin913 Apr 2, 2026
32b96fa
fix stream_k
bin913 Apr 2, 2026
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
29 changes: 29 additions & 0 deletions benchmark/test_blas_perf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -350,3 +350,32 @@ def addr_input_fn(m, n, cur_dtype, device):
dtypes=FLOAT_DTYPES,
)
bench.run()


class GemmBenchmark(BlasBenchmark):
"""
benchmark for gemm
"""

Copy link
Copy Markdown
Collaborator

@you-and-you you-and-you Apr 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is GemmBenchmark set to pass? Does the base class BlasBenchmark, with its default generated input shapes and calling method, achieve 100% compatibility with the standard gemm interface requirements? If it is compatible, it is recommended to add a comment line to clarify this.

pass


@pytest.mark.gemm
def test_gemm_benchmark():
def gemm_input_fn(b, m, n, k, cur_dtype, device, b_column_major):
inp1 = torch.randn([m, k], dtype=cur_dtype, device=device)
if b_column_major:
inp2 = torch.randn([n, k], dtype=cur_dtype, device=device)
yield inp1, inp2.t()
else:
inp2 = torch.randn([k, n], dtype=cur_dtype, device=device)
yield inp1, inp2

bench = GemmBenchmark(
input_fn=gemm_input_fn,
op_name="gemm",
torch_op=torch.Tensor.mm,
dtypes=FLOAT_DTYPES,
)
bench.set_gems(flag_gems.ops.gemm)
bench.run()
2 changes: 2 additions & 0 deletions src/flag_gems/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,8 @@ def torch_ge(v):
("minimum", minimum),
("mm", mm),
("mm.out", mm_out),
("gemm", gemm),
("gemm_out", gemm_out),
("mse_loss", mse_loss),
("mul.Tensor", mul),
("mul_.Tensor", mul_),
Expand Down
3 changes: 3 additions & 0 deletions src/flag_gems/ops/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@
from flag_gems.ops.gather import gather, gather_backward
from flag_gems.ops.ge import ge, ge_scalar
from flag_gems.ops.gelu import gelu, gelu_, gelu_backward
from flag_gems.ops.gemm import gemm, gemm_out
from flag_gems.ops.get_scheduler_metadata import get_scheduler_metadata
from flag_gems.ops.glu import glu, glu_backward
from flag_gems.ops.groupnorm import group_norm, group_norm_backward
Expand Down Expand Up @@ -375,6 +376,8 @@
"gelu_",
"gelu_backward",
"get_scheduler_metadata",
"gemm",
"gemm_out",
"glu",
"glu_backward",
"group_norm",
Expand Down
Loading