Heap simulator assertion failure when scaling up JAX matmul sizes

[seokw-Ryan]

Hi everyone,

I’m experimenting with XLA’s heap simulator (via JAX) to characterize its memory‐allocation behavior as I increase matrix sizes for a simple mat-mul. At a certain point it crashes with:

Estimated size of each array: 400000000000000 bytes, 381469726.56 MB, 372529.03 GB
2025-05-30… E external/xla/xla/service/heap_simulator/heap_simulator.cc:2115] We did not find a place for our sliced allocation. This should not happen because MSA operates on an infinitely sized heap.
… F external/xla/xla/service/heap_simulator/heap_simulator.cc:2371] Check failed: chunks.size() == 1 (1 vs. 0)
Aborted (core dumped)

What I’m trying to do:

• Push the heap simulator to very large matrix sizes
• Observe allocation success, fragmentation behavior, and any internal statistics
• Avoid a hard abort so I can log metrics and continue benchmarking

Questions for the community:

• Has anyone seen this internal CHECK(chunks.size() == 1) failure before?
• Is there a way to catch or bypass that assertion in the heap simulator for non-fatal tracing?

Thanks.