Skip to content

Correct the memory hierarchy: the earlier table measured the thread count, not the part #32

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
AndreSlavescu merged 1 commit into from
Jul 30, 2026
Merged

Correct the memory hierarchy: the earlier table measured the thread count, not the part #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
99 changes: 51 additions & 48 deletions benchmarks/agx_memory_hierarchy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,21 +13,20 @@

Three things make the numbers trustworthy rather than suggestive:

saturation the sweep has to reach the known 120.6 GB/s at large sizes, or the thread count is too
low to saturate and every number is a lower bound on something else. Printed as a check
rather than assumed.
occupancy every working set is measured at several threadgroup counts and the best is kept. A fixed
count is not safe: the first version of this used 64 for everything, which saturates DRAM
at 128 MB but starves an 8 MB working set by a factor of twelve, and the resulting table
described the thread count rather than the hierarchy.
amortisation each dispatch reads gigabytes and runs for tens of milliseconds, so launch overhead is
far below the noise. Timing small kernels through a host round trip is what produced
three fabricated results earlier in this project.
interleaving sizes are measured in rotating order across rounds, because a sweep that walks from
small to large measures thermal drift as much as it measures the hierarchy.

One limit of this probe is worth knowing before reading its output. Inside the resident regime bandwidth
*rises* with working set -- 1545, 2006, 2138 and 2386 GB/s at 256 KB, 512 KB, 1 MB and 2 MB -- which
cannot be a property of a cache. A smaller working set means fewer inner iterations per pass, so the
pass loop's bookkeeping takes a larger share, and the effect shrinks as the set grows. The probe
therefore establishes that a resident set runs at 2386 GB/s or better and says nothing reliable about
how that varies below 2 MB. `metile.target.agx` records it as one number for that reason.
One limit remains after the occupancy sweep. The smallest working sets cannot be given both enough
threads to saturate and enough work per thread to amortise the pass loop, because the two demands
conflict once the set is only a few times the thread count. Numbers below about 1 MB are floors rather
than levels, and `metile.target.agx` records the resident regime as a single figure for that reason.

usage:
python benchmarks/agx_memory_hierarchy.py
Expand All @@ -50,11 +49,15 @@
from metile.target import agx

THREADGROUP = 256
# Enough threads to saturate DRAM, and few enough that a 256 KB working set still gives every thread
# something to do: the inner loop strides by the total thread count, so a working set smaller than that
# leaves threads idle and measures parallelism instead of bandwidth.
THREADGROUPS = 64
VECTOR_BYTES = 16 # float4
VECTOR_BYTES = 16

# Threadgroup counts tried at every working set, with the best taken. A single count cannot serve the
# whole sweep and picking one produced a wrong answer that sat in the target model until the occupancy
# probe contradicted it: at 64 groups an 8 MB working set reads 196 GB/s and at 512 it reads 2403, so the
# "level" recorded at 8 MB was the thread count, not the memory system. Too few threads starves the path;
# too many leaves each thread one inner iteration per pass, where loop bookkeeping competes with the
# loads. Only the maximum over the sweep is a property of the part.
THREADGROUP_COUNTS = (32, 64, 128, 256, 512, 1024) # float4

SIZES_KB = (256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144)

Expand All @@ -63,7 +66,6 @@ def _arguments():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--traffic", type=float, default=4.0, help="GB read per dispatch")
parser.add_argument("--rounds", type=int, default=5)
parser.add_argument("--threadgroups", type=int, default=THREADGROUPS)
return parser.parse_args()


Expand Down Expand Up @@ -93,75 +95,76 @@ def main():
from metile.runtime.metal_device import MetalDevice

device = MetalDevice.get()
threads = arguments.threadgroups * THREADGROUP
grid, block = (threads, 1, 1), (THREADGROUP, 1, 1)
pipeline = device.compile_msl(READ_KERNEL, "probe")

print(f"device: {device.name}")
print(f"{threads} threads, {arguments.traffic:g} GB read per dispatch")
print(f"{arguments.traffic:g} GB read per dispatch, best of {THREADGROUP_COUNTS} threadgroups")
print(f"recorded streaming ceiling: {agx.STREAMING_READ_GBPS} GB/s\n")

target_bytes = arguments.traffic * 1e9
cases = []
largest = max(SIZES_KB) * 1024
data = metile.Buffer(data=np.ones(largest // 4, dtype=np.float32))
out = metile.Buffer(data=np.zeros(threads * 4, dtype=np.float32))
out = metile.Buffer(data=np.zeros(max(THREADGROUP_COUNTS) * THREADGROUP * 4, dtype=np.float32))

cases = []
for size_kb in SIZES_KB:
size = size_kb * 1024
vectors = size // VECTOR_BYTES
if vectors < threads:
continue
passes = max(1, int(target_bytes // size))
for count in THREADGROUP_COUNTS:
if vectors < count * THREADGROUP:
continue
cases.append((size, count, vectors, passes))

def measure(case):
_, count, vectors, passes = case
threads = count * THREADGROUP
buffers = [
data.metal_buffer,
out.metal_buffer,
metile.Buffer(data=np.array([vectors], dtype=np.uint32)).metal_buffer,
metile.Buffer(data=np.array([passes], dtype=np.uint32)).metal_buffer,
metile.Buffer(data=np.array([threads], dtype=np.uint32)).metal_buffer,
]
cases.append((size, vectors, passes, buffers))

def measure(buffers):
started = time.perf_counter_ns()
device.dispatch_kernel(pipeline, buffers, grid, block)
device.dispatch_kernel(pipeline, buffers, (threads, 1, 1), (THREADGROUP, 1, 1))
device.sync()
return (time.perf_counter_ns() - started) / 1e9
seconds = (time.perf_counter_ns() - started) / 1e9
return vectors * VECTOR_BYTES * passes / seconds / 1e9

for _, _, _, buffers in cases:
measure(buffers)
for case in cases:
measure(case)

samples = {size: [] for size, _, _, _ in cases}
samples = {case: [] for case in cases}
for index in range(arguments.rounds):
ordered = cases[index % len(cases) :] + cases[: index % len(cases)]
for size, _, _, buffers in ordered:
samples[size].append(measure(buffers))
for case in ordered:
samples[case].append(measure(case))

rates = {case: statistics.median(values) for case, values in samples.items()}

print(f"{'working set':>13}{'passes':>8}{'ms':>9}{'GB/s':>9}{'vs DRAM':>9}")
print(f"{'working set':>13}{'best groups':>13}{'GB/s':>9}{'vs DRAM':>9}{'worst groups':>14}")
results = []
for size, vectors, passes, _ in cases:
seconds = statistics.median(samples[size])
gbps = (vectors * VECTOR_BYTES * passes) / seconds / 1e9
results.append((size, gbps))
for size_kb in SIZES_KB:
size = size_kb * 1024
here = [(case, rate) for case, rate in rates.items() if case[0] == size]
if not here:
continue
best_case, best = max(here, key=lambda pair: pair[1])
worst = min(rate for _, rate in here)
results.append((size, best))
label = f"{size // 1024} KB" if size < 1024 * 1024 else f"{size // (1024 * 1024)} MB"
print(
f"{label:>13}{passes:>8}{seconds * 1e3:>9.1f}{gbps:>9.1f}"
f"{gbps / agx.STREAMING_READ_GBPS:>8.2f}x"
f"{label:>13}{best_case[1]:>13}{best:>9.0f}"
f"{best / agx.STREAMING_READ_GBPS:>8.2f}x{best / worst:>13.1f}x"
)

dram = min(gbps for _, gbps in results)
peak = max(gbps for _, gbps in results)
print(f"\nslowest {dram:.1f} GB/s, fastest {peak:.1f} GB/s, ratio {peak / dram:.2f}x")
saturated = dram >= 0.9 * agx.STREAMING_READ_GBPS
print(f"\nslowest {dram:.0f} GB/s, fastest {peak:.0f} GB/s, ratio {peak / dram:.2f}x")
print(
f"saturation check: largest working sets reach {dram:.1f} GB/s against a recorded "
f"{agx.STREAMING_READ_GBPS} GB/s -- {'ok' if saturated else 'TOO LOW, raise --threadgroups'}"
"last column is how much the thread count alone moves that size, which is why it is swept."
)
if not saturated:
print(
"Until that passes, every number here is a lower bound and the knees may be artefacts."
)
return 1

# Report the knees rather than leaving them to be eyeballed: a level boundary is where bandwidth
# drops materially between adjacent sizes.
Expand Down
158 changes: 158 additions & 0 deletions benchmarks/agx_occupancy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
"""How much parallelism a kernel needs before memory bandwidth saturates.

Decode kernels achieve 80 to 128 GB/s while the memory level their working set sits at delivers 128 to
555. That gap is not tiling -- decode reads each weight once, so there is no reuse to capture -- which
leaves the question of what it is. The first suspect is that the kernel is not running enough
threadgroups to saturate the path in the first place.

That is a property of the part, not of any kernel, so it can be measured directly: hold the working set
fixed and sweep the number of threadgroups. The curve's knee is the parallelism a kernel needs, and any
kernel launching fewer than that is leaving bandwidth behind for a reason a tile-size change could fix.

The number matters for tile selection specifically. A GEMM tiled by output width launches
`output_features / block_n` threadgroups, so a wider tile means fewer of them: at N=8960 a block_n of 64
gives 140 threadgroups and 256 gives 35. If saturation needs more than 35, the widest tiles are
self-limiting and the compiler can rule them out before measuring anything.

Same controls as the other probes here: each dispatch reads gigabytes so launch overhead is negligible,
thread counts are measured in rotating order so drift does not land on one end of the sweep, and elapsed
time has to scale with the pass count or the loop is being collapsed.
"""

import argparse
import statistics
import sys
import time
from pathlib import Path

import numpy as np

_root = str(Path(__file__).resolve().parent.parent)
sys.path.insert(0, _root)

import metile

THREADGROUP = 256
VECTOR_BYTES = 16
THREADGROUP_COUNTS = (1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 192, 256, 384, 512)

# Two footprints: one in the resident regime and one streaming, because the parallelism needed to
# saturate a cache and to saturate DRAM are not obviously the same number.
WORKING_SETS = (1024 * 1024, 8 * 1024 * 1024, 64 * 1024 * 1024)

KERNEL = """#include <metal_stdlib>
using namespace metal;

kernel void probe(device const float4* data [[buffer(0)]],
device float4* out [[buffer(1)]],
constant uint& vectors [[buffer(2)]],
constant uint& passes [[buffer(3)]],
constant uint& stride [[buffer(4)]],
uint gid [[thread_position_in_grid]]) {
float4 total = float4(0.0f);
for (uint pass = 0; pass < passes; ++pass) {
for (uint index = gid; index < vectors; index += stride) {
total += data[index];
}
}
out[gid] = total;
}
"""


def _arguments():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--traffic", type=float, default=4.0, help="GB read per dispatch")
parser.add_argument("--rounds", type=int, default=3)
return parser.parse_args()


def main():
arguments = _arguments()
from metile.runtime.metal_device import MetalDevice

device = MetalDevice.get()
pipeline = device.compile_msl(KERNEL, "probe")
largest = max(WORKING_SETS)
data = metile.Buffer(data=np.ones(largest // 4, dtype=np.float32))
out = metile.Buffer(data=np.zeros(max(THREADGROUP_COUNTS) * THREADGROUP * 4, dtype=np.float32))

print(f"device: {device.name}")
print(f"{THREADGROUP} threads per group, {arguments.traffic:g} GB read per dispatch\n")

def measure(threadgroups, size):
threads = threadgroups * THREADGROUP
vectors = size // VECTOR_BYTES
passes = max(1, int(arguments.traffic * 1e9 // size))
buffers = [
data.metal_buffer,
out.metal_buffer,
metile.Buffer(data=np.array([vectors], dtype=np.uint32)).metal_buffer,
metile.Buffer(data=np.array([passes], dtype=np.uint32)).metal_buffer,
metile.Buffer(data=np.array([threads], dtype=np.uint32)).metal_buffer,
]
started = time.perf_counter_ns()
device.dispatch_kernel(pipeline, buffers, (threads, 1, 1), (THREADGROUP, 1, 1))
device.sync()
seconds = (time.perf_counter_ns() - started) / 1e9
return vectors * VECTOR_BYTES * passes / seconds / 1e9

cases = [
(count, size)
for size in WORKING_SETS
for count in THREADGROUP_COUNTS
if size // VECTOR_BYTES >= count * THREADGROUP
]
for count, size in cases:
measure(count, size)

samples = {case: [] for case in cases}
for index in range(arguments.rounds):
ordered = cases[index % len(cases) :] + cases[: index % len(cases)]
for case in ordered:
samples[case].append(measure(*case))

rates = {case: statistics.median(values) for case, values in samples.items()}

header = f"{'groups':>7}{'threads':>9}" + "".join(
f"{size // (1024 * 1024) if size >= 1024 * 1024 else size // 1024:>10}"
+ ("MB" if size >= 1024 * 1024 else "KB")
for size in WORKING_SETS
)
print(header)
print("-" * len(header))
for count in THREADGROUP_COUNTS:
cells = []
for size in WORKING_SETS:
rate = rates.get((count, size))
cells.append(f"{rate:>10.0f} " if rate else f"{'-':>12}")
print(f"{count:>7}{count * THREADGROUP:>9}" + "".join(cells))

print("\nthreadgroups needed to reach 90% of the best this sweep saw, per working set:")
knees = {}
for size in WORKING_SETS:
curve = [
(count, rates[count, size]) for count in THREADGROUP_COUNTS if (count, size) in rates
]
peak = max(rate for _, rate in curve)
knee = next(count for count, rate in curve if rate >= 0.9 * peak)
knees[size] = knee
label = f"{size // (1024 * 1024)} MB" if size >= 1024 * 1024 else f"{size // 1024} KB"
print(f" {label:>7}: {knee:>4} groups ({knee * THREADGROUP} threads) for {peak:.0f} GB/s")

# What this means for tile selection, which is the reason to measure it.
print("\nfor a GEMM tiled by output width, groups launched = output_features / block_n:")
print(f"{'N':>7}" + "".join(f"{f'bn={bn}':>10}" for bn in (32, 64, 128, 256)))
worst_knee = max(knees.values())
for output_features in (4864, 8192, 8960, 17408):
cells = []
for block_n in (32, 64, 128, 256):
groups = output_features // block_n
cells.append(f"{groups:>7}{' ok' if groups >= worst_knee else ' LOW'}")
print(f"{output_features:>7}" + "".join(cells))
print(f"\nLOW marks fewer than {worst_knee} groups, the most demanding knee above.")
return 0


if __name__ == "__main__":
raise SystemExit(main())
2 changes: 2 additions & 0 deletions metile/target/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from metile.target.agx import (
BANDWIDTH_BY_WORKING_SET_GBPS,
GROUPS_FOR_RESIDENT_BANDWIDTH,
ILP_CEILING,
MATRIX_PEAK_TFLOPS,
REGISTER_BUDGET,
Expand All @@ -31,6 +32,7 @@

__all__ = [
"BANDWIDTH_BY_WORKING_SET_GBPS",
"GROUPS_FOR_RESIDENT_BANDWIDTH",
"ILP_CEILING",
"MATRIX_PEAK_TFLOPS",
"REGISTER_BUDGET",
Expand Down
Loading
Loading