Skip to content

Constraint on split_k on m * n #8404

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
merged 15 commits into from
Oct 11, 2025
+234 −6
Merged

Constraint on split_k on m * n #8404

Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
10e7a7b
Update constraints dict to accept callable for split_k and invoke tha…
afroz-oai Oct 8, 2025
9f20705
Update constraints dict to accept callable for split_k and invoke tha…
afroz-oai Oct 8, 2025
d2efa7c
Make notes that one needs to update constraint checking
afroz-oai Oct 8, 2025
af2feed
Merge branch 'afroz/split-k-callable' into patch-1
afroz-oai Oct 8, 2025
e89313e
Add tests
afroz-oai Oct 8, 2025
fa3ef72
Rework split_k callable to dynamic_split_k instead and add tests
afroz-oai Oct 9, 2025
04cc7d2
Change dynamic_split_k to only return 1 for the cursed shape, else 4
afroz-oai Oct 9, 2025
206ee73
Function that checks shape only on m * n, doesn't involve b
afroz-oai Oct 10, 2025
6967d01
.
afroz-oai Oct 10, 2025
5276531
.
afroz-oai Oct 10, 2025
f5b5ebb
More and better tests
afroz-oai Oct 10, 2025
7fc4959
.
afroz-oai Oct 10, 2025
899749f
,
afroz-oai Oct 10, 2025
30f4e69
Move test_opt_flags_split_k.py into a new test_matmul_details folder
afroz-oai Oct 10, 2025
fba0c8e
Merge branch 'main' into patch-1
afroz-oai Oct 10, 2025
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
25 changes: 22 additions & 3 deletions python/triton_kernels/triton_kernels/matmul_ogs_details/opt_flags.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,10 @@
import torch
from .opt_flags_details import opt_flags_amd, opt_flags_nvidia
from triton_kernels.tensor import bitwidth
from typing import Callable

# Function type: takes four ints (batch_size, m, n, k) and a output dtype, returns an int
CallableSplitK = Callable[[int, int, int, int, torch.dtype], int]


@dataclass
Expand All @@ -33,6 +37,19 @@ def __post_init__(self):
raise ValueError("Not supported")


def get_split_k_from_constraints(
constraints_split_k: int | CallableSplitK,
batch_size: int,
m: int,
n: int,
k: int,
out_dtype: torch.dtype) -> int:
if isinstance(constraints_split_k, int):
return constraints_split_k

return constraints_split_k(batch_size, m, n, k, out_dtype)


def make_default_opt_flags_amd(
out_dtype,
lhs_dtype,
Expand Down Expand Up @@ -91,7 +108,8 @@ def make_default_opt_flags_amd(
is_persistent = constraints.get("is_persistent", False)
# split_k:
if constraints.get("split_k", None) is not None:
split_k = constraints["split_k"]
split_k = get_split_k_from_constraints(
constraints["split_k"], batch_size, m, n, k, out_dtype)
elif is_persistent or enforce_bitwise_invariance:
split_k = 1
else:
Expand Down Expand Up @@ -221,7 +239,8 @@ def make_default_opt_flags_nvidia(
block_n, block_k = block_k, block_n
# split_k
if constraints.get("split_k", None) is not None:
split_k = constraints["split_k"]
split_k = get_split_k_from_constraints(
constraints["split_k"], batch_size, m, n, k, out_dtype)
elif is_persistent or enforce_bitwise_invariance or precision_config.act_scale is not None or precision_config.out_scale is not None:
split_k = 1
else:
Expand Down Expand Up @@ -293,7 +312,7 @@ def make_default_opt_flags_nvidia(
_opt_flags_constraints: dict = dict()
_opt_flags: OptFlags | None = None

def update_opt_flags_constraints(constraints: dict[str, int]):
def update_opt_flags_constraints(constraints: dict[str, int | Callable]):
global _opt_flags_constraints
_opt_flags_constraints.update(constraints)

Expand Down