transformers
0bb187fe - [CB] Make CB more device agnostic and add support for XPU (#49156)

Commit
39 days ago
[CB] Make CB more device agnostic and add support for XPU (#49156) * This PR generalizes continuous batching graph support from CUDA-only graphs to accelerator graphs, with support for both CUDA and XPU graph capture. The main changes are: - Add `use_accelerator_graph` to `ContinuousBatchingConfig` while keeping `use_cuda_graph` and `cuda_graph_booleans` as backward-compatible aliases. - Rename the continuous batching graph decision path from CUDA graphs to accelerator graphs. - Add device-module helpers for CUDA/XPU graph capture, graph pools, memory pool contexts, streams, events, and cache cleanup. - Use `torch.xpu.XPUGraph` / `torch.xpu.graph` when running on XPU, and keep `torch.cuda.CUDAGraph` / `torch.cuda.graph` for CUDA. - Detect accelerator graph support through graph-capture APIs, not just device availability. - Update async batching, IO buffers, offloading, and manager cleanup to use the active accelerator module instead of hard-coded `torch.cuda`. - Update continuous batching docs and test terminology from CUDA graphs to accelerator graphs where applicable. - Keep Inductor graph-mode detection based on the existing `triton.cudagraphs` option, since current PyTorch/XPU builds do not expose a separate `triton.xpugraphs` option. Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * format Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * fmt Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * update Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * update Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * update Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * extend tp test Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * fmt Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * fix test crash Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * update Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * update Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * fmt Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * update Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * fix test Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> * Core * Note * Fix merge main --------- Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> Signed-off-by: Wang, Yi <yi.a.wang@intel.com> Co-authored-by: Wang, Yi A <yi.a.wang@intel.com>
Author
Parents
Loading