trunk/c34b6eb9467e0924cbbe79208de91a8bb51d92e2: [CUDA] Pass device and stream to Python MemPool callbacks (#200427)
- PyTorch 近 90 天出现 2131 次
- PyTorch 近 90 天第 2102 次发版
- 上一次：同一天稍早 · trunk/2208903bd2a09f30786d54b3d2426c98013d0159: [Inductor] Add gfx950 FlyDSL FlexAttention forward (#194309)
发生了什么
Pass the allocation device and raw stream address directly to Python-backed MemPool allocation and free callbacks. Remove the implicit CUDAStreamGuard so callbacks can decide how to establish and use their CUDA context. Update the public callback signatures, documentation, examples, and tests. The explicit-stream test also verifies that invoking a callback does not change its current stream. Test Plan: python test/t…
摘要按规则整理自下方来源原文