Lol I interned inside pytorch a few years ago (you and I even met/talked about tangential things :)) and worked on tracking such allocations (although I didn't hook it up to profiler). Spoiler alert: you can't track such provenance because everything gets muddled in the dispatcher.

EDIT: not completely accurate to say you can't do it. I prototyped a little allocator that would stamp every allocation (the pointer itself, in the unused bits, a trick I learned from zach) with the thread id and a timestamp (just an incrementing counter) and then percolate that up to the surface. Obv that didn't land lol.