MI300X Cache / HBM Bandwidth Microbenchmark (Triton)

This repository contains a Triton-based microbenchmark designed to measure effective memory bandwidth on AMD MI300X across different levels of the memory hierarchy, including L1, L2, last-level cache (often referred to as Infinity Cache), and HBM.

The benchmark works by sweeping the per-program working set size and observing how the achieved bandwidth changes as the working set exceeds different cache capacities.

Important note on terminology
MI300X features per-CU L1 caches, per-XCD L2 cache, and a large shared on-package last-level cache in the MI300 family (often called Infinity Cache in AMD documentation). This benchmark does not directly identify cache boundaries by name; instead, it detects bandwidth regimes corresponding to where data is effectively served from.

---

1. What this benchmark measures

1.1 Kernel behavior

The core kernel measure_bandwidth_kernel is launched as a 1D grid of Triton programs. Each program (pid) operates on an independent input stream:

• Input tensor x: shape (grids, stream_elems), float32
• Output tensor y: shape (grids, BLOCK_SIZE), float32
• Each program loads all stream_elems elements from x[pid, :]
• Loads are done in chunks of BLOCK_SIZE
• The full stream is loaded load_iterations times
• Accumulation (val_sum += val) prevents dead-code elimination
• A final store writes BLOCK_SIZE elements per program

Conceptually, per program:

for i in range(load_iterations):
    read stream_elems * 4 bytes from x[pid, :]
write BLOCK_SIZE * 4 bytes to y[pid, :]

The volatile=True flag on tl.load discourages the compiler from hoisting or reusing loads across iterations, helping ensure each iteration issues real memory traffic.

---

1.2 Timing and bandwidth calculation

Timing is measured using CUDA events:

• Several warmup iterations are executed
• The kernel is then run n_repeat times
• The average execution time avg_ms is recorded

Approximate byte traffic per run:

• Load bytes: stream_elems * 4 * load_iters * grids
• Store bytes: BLOCK_SIZE * 4 * grids
• Total bytes: load + store

Observed bandwidth:

Bandwidth (GB/s) = total_bytes / (avg_ms / 1000) / 1e9

---

1.3 Linear regression to isolate memory cost

For each stream_elems, the benchmark runs multiple values of load_iters and fits:

avg_ms ≈ intercept_ms + load_ms_per_iter * load_iters

This separates:

• Intercept: launch overhead, loop overhead, stores
• Slope: pure streaming load cost per iteration

Derived bandwidth:

bandwidth (GB/s) = grids * stream_elems * 4 / load_ms_per_iter / 1e6

This regression is implemented in compute_bandwidth() using sklearn.linear_model.LinearRegression.

---

2. Why this reveals different cache levels

The benchmark repeatedly accesses the same memory addresses. After warmup:

• If the working set fits in L1, loads hit L1 → very high bandwidth
• If it exceeds L1 but fits in L2, bandwidth drops to L2 level
• If it exceeds L2 but fits in the large shared cache, another plateau appears
• Once all caches are exceeded, traffic comes from HBM, giving the lowest plateau

Because accesses are repeated many times, the measured bandwidth reflects the highest cache level that can retain the working set under current pressure.

---

2.1 Two working sets matter

There are two distinct working set sizes:

1. Per-program working set

   stream_elems * 4 bytes
   

2. Aggregated working set across concurrently active programs

   active_programs * stream_elems * 4 bytes
   

Even if the per-program working set fits in cache, the aggregated working set may overflow it. This is critical for understanding the role of grid_multiplier.

---

3. Key parameters and their effects

3.1 --stream-elements (primary cache sweep knob)

Controls the per-program working set size.

• Small values → L1 / L2 dominated
• Larger values → L2 / LLC / HBM

In benchmark mode, this parameter is swept automatically. You can modify the stream_elems_sweep variable in the run_benchmark function to adjust the range and change granularity.

---

3.2 --load-iterations

Repeats the full stream load multiple times.

• Improves signal-to-noise ratio
• Enables linear regression to isolate memory cost

---

3.3 --block-size (BLOCK_SIZE)

Controls vector width per Triton program.

• Larger blocks increase per-program parallelism and may increase memory coalescing efficiency.
• Too large can reduce occupancy or increase register pressure, potentially reducing throughput.
• Must be one of: 64, 128, 256, 512, 1024, 2048.

---

3.4 --num-warps (fix it to 4)

Controls the number of warps per program.

• More warps improve latency hiding (especially for HBM)
• Too many warps increase register pressure and reduce occupancy

---

3.5 --num-stages (fix it to 1)

Software pipelining depth.

• Modest values can help hide memory latency
• Too large increases register usage

---

3.6 --grid-multiplier (CRITICAL PARAMETER)

This parameter controls how many Triton programs are launched:

grids = num_cu * grid_multiplier

Why this matters

• Each program has its own independent stream

• Increasing grid_multiplier increases:

  • Concurrency
  • Aggregate working set size
  • Cache pressure

• Decreasing grid_multiplier reduces:

  • Cache pressure
  • Memory-level parallelism
  • Overall GPU utilization

Observed behavior (important)

Through experimentation, the following behavior was observed:

• Too small grid_multiplier (e.g., 1 or 2):

  • GPU is underutilized
  • Not enough concurrent programs to hide memory latency
  • Measured bandwidth is artificially low

• Moderate grid_multiplier (≈ 4):

  • Sufficient concurrency to saturate memory pipelines
  • Cache levels saturate at realistic working set sizes
  • Bandwidth measurements become stable and representative

• Very large grid_multiplier:

  • Aggressive cache pressure
  • Cache knees shift to smaller per-program working sets
  • Useful for stress-testing cache capacity, but harder to interpret

Recommendation: A grid_multiplier of 4 is the smallest value that consistently produces reliable and realistic bandwidth measurements on MI300X. Smaller values tend to under-drive the GPU and underestimate effective bandwidth.

This parameter directly determines how fast high-level caches (L1/L2/LLC) become saturated, because it controls the aggregated working set size seen by the cache hierarchy.

---

3.7 --num-cu

Logical number of compute units used to size the grid. Default is 304 for MI300X.

---

3.8 --num-repeat

Number of timing repetitions.

• Increase for small working sets and fast kernels
• Improves timing stability

---

4. Expected results and interpretation

When sweeping stream_elements, you should observe:

1. Highest bandwidth plateau (L1-dominated)
2. Lower plateau (L2-dominated)
3. Another plateau (shared LLC / Infinity Cache)
4. Lowest plateau (HBM bandwidth ceiling)

Additional observations:

• Cache knee points shift with grid_multiplier
• Underutilization shows up as uniformly low bandwidth
• Regression r2 close to 1.0 indicates clean measurements

---

5. How to run

5.1 Requirements

• MI300X system

• Triton + PyTorch environment targeting the GPU

• Python packages:

  • torch
  • triton
  • numpy
  • pandas
  • scikit-learn

---

5.2 Run a single configuration

python mi300x_bw.py run \
  --stream-elements 1024 \
  --load-iterations 10 \
  --block-size 1024 \
  --num-warps 4 \
  --num-stages 1 \
  --num-cu 304 \
  --grid-multiplier 4 \
  --num-repeat 200

---

5.3 Run the full benchmark sweep

You can modify the stream_elems_sweep variable in the run_benchmark function to adjust the data size range.

python mi300x_bw.py benchmark \
  --test-name mi300x_bw_test \
  --grid-multiplier 4 \
  --num-warps 4 \
  --block-size 1024

Outputs are written to ./results/:

• *_details_*.csv: raw timing data
• *_summary_*.csv: regression-based bandwidth estimates

---

6. Best practices

• Use grid_multiplier >= 4
• Keep the system free of other GPU workloads
• Watch regression r2 values
• Compare multiple grid_multiplier values to understand cache pressure effects
• Fix or monitor GPU clocks if possible

---

7. Limitations

• Cache residency is probabilistic, not guaranteed
• Hardware prefetching and replacement policies still apply
• Measures effective bandwidth, not theoretical peak
• Results depend on concurrency, occupancy, and scheduling