Skip to content

feat: init stream guard - #247

Open
Chamberlain0w0 wants to merge 2 commits into
masterfrom
feat/stream_guard
Open

Chamberlain0w0 wants to merge 2 commits into
masterfrom
feat/stream_guard

Conversation

@Chamberlain0w0

@Chamberlain0w0 Chamberlain0w0 commented Oct 12, 2026 •

Copy link
Copy Markdown
Collaborator

概述

由于 CP 实现涉及切换计算流操作,所以新增公共运行时组件 StreamGuard,支持在作用域内切换 device 和 stream,并在退出时自动恢复,供多流计算等场景复用。

主要修改

  • 新增 DeviceGuardImpl::ExchangeStream:切换指定设备在线程内的当前 stream,返回原 stream,不切换当前设备、不执行同步。
  • 新增独立的 stream_guard.h/.cc:进入作用域时保存并切换 device 和 stream;退出时先恢复目标设备的原 stream,再恢复原 device。禁止拷贝和移动。
  • CUDA 后端使用线程局部状态维护各设备的当前 stream,并在获取 cuBLAS handle 时绑定当前 stream。
  • 补齐 CPU 和默认后端接口,遵循现有不支持 stream 操作时的处理方式。

StreamGuard 不持有 stream 所有权,也不执行同步;调用方负责 stream 的设备归属、生命周期和跨 stream 依赖。

Comment thread infini_train/src/core/runtime/cuda/cuda_guard_impl.cc
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants