Reverse sort matmul benchmarks by memory use to avoid fragmentation

[jacobhinkle]

Before this PR memory use on H200 grows slowly to reach 126 GiB out of 128 GiB capacity.

After this PR memory use on H200 never goes above about 66 GiB.

Also, we previously had a 20 GiB cutoff for included benchmarks but this changes that to 90% of GPU capacity.

[github-actions]

Review updated until commit be28037

Description

• Reverse sort matmul benchmarks by memory use

• Update OOM cutoff to 90% of GPU capacity

• Introduce maybe_skip_oom_case utility

• Add docstrings and utility functions

---

Changes walkthrough 📝

Relevant files

Enhancement

test_matmul.py Enhance matmul benchmark memory management                              benchmarks/python/test_matmul.py Import functools for comparison utilities Add row_mem function to compute memory usage Add three_way_cmp and mem_cmp functions for sorting Reverse sort matmul problems by memory use Introduce maybe_skip_oom_case to skip OOM cases Update OOM cutoff to 90% of GPU capacity Use maybe_skip_oom_case in test functions +42/-5

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Memory Calculation The memory calculation in row_mem and maybe_skip_oom_case assumes half-precision (FP16) data types. Ensure that this assumption is correct for all use cases and that the memory calculation is accurate. def row_mem(row): """Compute gmem bytes used in a half-precision GEMM This computes only the space required for operands and output, ignoring intermediates like split-K buffers and assuming no bias terms. """ m, n, k, _ = row return ((m + n) * k + m * n) * 2 def three_way_cmp(a, b) -> int: """Perform a three-way comparison like the Python2 cmp() function This returns 0 if a == b, -1 if a < b, and 1 if a > b """ return int(a > b) - int(a < b) def mem_cmp(row1, row2): """Compare two rows based on memory use""" return three_way_cmp(row_mem(row1), row_mem(row2)) # Reverse sort by expected memory use to avoid fragmentation rows.sort(key=functools.cmp_to_key(mem_cmp), reverse=True) return rows def maybe_skip_oom_case(m: int, n: int, k: int): expected_mem = (m * k + n * k + m * n) * 2 # operands plus output expected_mem *= 2 # account for multiple runs/deferred frees _, total = torch.cuda.mem_get_info() max_mem = total * 0.9 if expected_mem > max_mem: pytest.skip( f"Case takes more than {max_mem / (2 ** 30): .2f} GiB. Skipping to avoid OOM" ) Three-Way Comparison The three_way_cmp function is a reimplementation of the Python 2 cmp function. Consider using Python's built-in comparison operators or functools.cmp_to_key directly for simplicity and readability. def three_way_cmp(a, b) -> int: """Perform a three-way comparison like the Python2 cmp() function This returns 0 if a == b, -1 if a < b, and 1 if a > b """ return int(a > b) - int(a < b) Memory Threshold The memory threshold is set to 90% of GPU capacity. Verify that this threshold is appropriate for all target hardware and that it balances performance and resource utilization effectively. max_mem = total * 0.9 if expected_mem > max_mem:

[jacobhinkle]

!build

[Priya2698]

I believe you switched to custom memory checks for matmul since retry_on_oom resulted in errors pertaining to benchmark fixture being used twice?

[Priya2698]

I think this condition can be re-written as return (a > b) - (a < b).

[jacobhinkle]

Clever!

[Priya2698]

Should we also clear allocated memory here?

[jacobhinkle]

Since we reverse sort the benchmarks by required memory, we shouldn't need to clear the allocated memory, since subsequent tests will fit into the already-allocated memory pool.

[Priya2698]

LGTM

[jacobhinkle]

!build