feat: init stream guard

infini_train/include/core/runtime/device_guard.h

@@ -86,6 +86,12 @@ class DeviceGuardImpl {
 
   virtual Stream *GetStream(Device) const;
 
+  // Set the thread-local current stream for device and return its previous
+  // stream. Does not change the current device or synchronize. Stream ownership
+  // is unchanged. The caller must provide a stream belonging to device and keep
+  // it alive while selected.
+  virtual Stream *ExchangeStream(Device device, Stream *stream) const;
+
   virtual Stream *CreateStream(Device) const;
 
   virtual Stream *CreateStreamWithPriority(Device, int priority) const;

infini_train/include/core/runtime/stream_guard.h

@@ -0,0 +1,36 @@
+#pragma once
+
+#include "infini_train/include/core/runtime/device_guard.h"
+
+namespace infini_train {
+class Device;
+}
+
+namespace infini_train::core {
+
+class Stream;
+
+// Select device and its current stream for this scope. On destruction, restore
+// the target device's previous stream, then restore the original device within
+// the same backend. The stream is borrowed, must belong to device, and must
+// outlive this guard. No synchronization is performed; callers remain
+// responsible for cross-stream dependencies.
+class StreamGuard {
+ public:
+  explicit StreamGuard(Device device, Stream *stream);
+  ~StreamGuard();
+
+  StreamGuard(const StreamGuard &) = delete;
+  StreamGuard &operator=(const StreamGuard &) = delete;
+  StreamGuard(StreamGuard &&) = delete;
+  StreamGuard &operator=(StreamGuard &&) = delete;
+
+ private:
+  // Declared first so device restoration also runs if stream exchange throws
+  // during construction.
+  DeviceGuard device_guard_;
+  DeviceGuardImpl *impl_;
+  Stream *previous_stream_;
+};
+
+}  // namespace infini_train::core

infini_train/src/core/runtime/cpu/cpu_guard_impl.cc

@@ -39,6 +39,13 @@ Stream *CpuGuardImpl::GetStream(Device device) const {
   return nullptr;
 }
 
+Stream *CpuGuardImpl::ExchangeStream(Device device, Stream *stream) const {
+  CHECK(device.type() == Device::DeviceType::kCPU);
+  VLOG(3) << "CpuGuardImpl::ExchangeStream is not supported. "
+             "Return nullptr.";
+  return nullptr;
+}
+
 Stream *CpuGuardImpl::CreateStream(Device device) const {
   CHECK(device.type() == Device::DeviceType::kCPU);
   VLOG(3) << "CpuGuardImpl::CreateStream is not supported. "

infini_train/src/core/runtime/cpu/cpu_guard_impl.h

@@ -27,6 +27,8 @@ class CpuGuardImpl final : public DeviceGuardImpl {
   // Stream management (explicitly unsupported for now)
   Stream *GetStream(Device device) const override;
 
+  Stream *ExchangeStream(Device device, Stream *stream) const override;
+
   Stream *CreateStream(Device device) const override;
 
   Stream *CreateStreamWithPriority(Device device, int priority) const override;

infini_train/src/core/runtime/cuda/cuda_guard_impl.cc

@@ -3,6 +3,7 @@
 #include <array>
 #include <memory>
 #include <mutex>
+#include <vector>
 
 #include "infini_train/include/common/cuda/common_cuda.h"
 #include "infini_train/include/core/runtime/runtime_common.h"
@@ -16,10 +17,19 @@ constexpr int kMaxGpus = 8;
 constexpr size_t kBytesPerMB = 1024ULL * 1024ULL;
 
 static std::array<std::unique_ptr<CudaStream>, kMaxGpus> cuda_streams;
-static std::array<std::unique_ptr<CudaBlasHandle>, kMaxGpus> cuda_blas_handles;
+static thread_local std::array<Stream *, kMaxGpus> current_streams = {};
+// Each thread exclusively reserves one handle per device. The pool retains
+// handles after thread exit for reuse, since CudaBlasHandle deliberately does
+// not destroy its CUDA resources during teardown.
+struct BlasHandlePool {
+  std::mutex mutex;
+  std::array<std::vector<std::shared_ptr<CudaBlasHandle>>, kMaxGpus> handles;
+};
+static BlasHandlePool cuda_blas_handle_pool;
+static thread_local std::array<std::shared_ptr<CudaBlasHandle>, kMaxGpus>
+    cuda_blas_handles;
 
 static std::array<std::once_flag, kMaxGpus> device_stream_flags;
-static std::array<std::once_flag, kMaxGpus> device_handle_flags;
 
 inline void CheckCudaDevice(Device device) {
   CHECK(device.type() == Device::DeviceType::kCUDA) << std::format(
@@ -51,22 +61,37 @@ void CudaGuardImpl::InitSingleStream(Device device) {
   CUDA_CHECK(cudaSetDevice(device.index()));
 
   cuda_streams[device.index()] = std::make_unique<CudaStream>();
+  current_streams[device.index()] = cuda_streams[device.index()].get();
 
   CUDA_CHECK(cudaSetDevice(current_device));
 }
 
 void CudaGuardImpl::InitSingleHandle(Device device) {
   CheckCudaDevice(device);
 
+  auto &handle = cuda_blas_handles.at(device.index());
+  if (handle) {
+    return;
+  }
+  std::lock_guard<std::mutex> lock(cuda_blas_handle_pool.mutex);
+  auto &pooled_handles = cuda_blas_handle_pool.handles.at(device.index());
+  for (const auto &pooled_handle : pooled_handles) {
+    // A pool-only reference means no thread currently reserves this handle.
+    if (pooled_handle.use_count() == 1) {
+      handle = pooled_handle;
+      return;
+    }
+  }
+
   int current_device = -1;
   CUDA_CHECK(cudaGetDevice(&current_device));
   CUDA_CHECK(cudaSetDevice(device.index()));
 
   std::call_once(device_stream_flags.at(device.index()), InitSingleStream,
                  device);
 
-  cuda_blas_handles[device.index()] =
-      std::make_unique<CudaBlasHandle>(cuda_streams[device.index()].get());
+  handle = std::make_shared<CudaBlasHandle>(cuda_streams[device.index()].get());
+  pooled_handles.push_back(handle);
 
   CUDA_CHECK(cudaSetDevice(current_device));
 }
@@ -103,7 +128,23 @@ Stream *CudaGuardImpl::GetStream(Device device) const {
   // times within the same process (e.g. in unit tests).
   std::call_once(device_stream_flags.at(device.index()), InitSingleStream,
                  device);
-  return cuda_streams.at(device.index()).get();
+  if (current_streams.at(device.index()) == nullptr) {
+    current_streams.at(device.index()) = cuda_streams.at(device.index()).get();
+  }
+  return current_streams.at(device.index());
+}
+
+Stream *CudaGuardImpl::ExchangeStream(Device device, Stream *stream) const {
+  CheckCudaDevice(device);
+  std::call_once(device_stream_flags.at(device.index()), InitSingleStream,
+                 device);
+  CHECK_NOTNULL(stream);
+  if (current_streams.at(device.index()) == nullptr) {
+    current_streams.at(device.index()) = cuda_streams.at(device.index()).get();
+  }
+  auto *previous_stream = current_streams.at(device.index());
+  current_streams[device.index()] = stream;
+  return previous_stream;
 }
 
 Stream *CudaGuardImpl::CreateStream(Device device) const {
@@ -229,9 +270,13 @@ void CudaGuardImpl::SynchronizeStream(Stream *stream) const {
 // blas
 BlasHandle *CudaGuardImpl::GetBlasHandle(Device device) const {
   CheckCudaDevice(device);
-  std::call_once(device_handle_flags.at(device.index()), InitSingleHandle,
-                 device);
-  return cuda_blas_handles.at(device.index()).get();
+  if (!cuda_blas_handles.at(device.index())) {
+    InitSingleHandle(device);
+  }
+  auto *handle = cuda_blas_handles.at(device.index()).get();
+  CUBLAS_CHECK(cublasSetStream(handle->cublas_handle(),
+                               GetCudaStream(GetStream(device))));
+  return handle;
 }
 
 // memory

infini_train/src/core/runtime/cuda/cuda_guard_impl.h

@@ -32,6 +32,8 @@ class CudaGuardImpl final : public DeviceGuardImpl {
   // constructor/destructor of class Stream
   Stream *GetStream(Device device) const override;
 
+  Stream *ExchangeStream(Device device, Stream *stream) const override;
+
   Stream *CreateStream(Device device) const override;
 
   Stream *CreateStreamWithPriority(Device device, int priority) const override;

infini_train/src/core/runtime/device_guard.cc

@@ -29,6 +29,11 @@ Stream *DeviceGuardImpl::GetStream(Device) const {
   return nullptr;  // unreachable
 }
 
+Stream *DeviceGuardImpl::ExchangeStream(Device, Stream *) const {
+  LOG(FATAL) << "DeviceGuardImpl::ExchangeStream is not implemented.";
+  return nullptr;  // unreachable
+}
+
 Stream *DeviceGuardImpl::CreateStream(Device) const {
   LOG(FATAL) << "DeviceGuardImpl::CreateStream is not implemented.";
   return nullptr;  // unreachable

infini_train/src/core/runtime/stream_guard.cc

@@ -0,0 +1,18 @@
+#include "infini_train/include/core/runtime/stream_guard.h"
+
+#include "infini_train/include/device.h"
+
+namespace infini_train::core {
+
+StreamGuard::StreamGuard(Device device, Stream *stream)
+    : device_guard_(device),
+      impl_(GetDeviceGuardImpl(device.type())),
+      previous_stream_(impl_->ExchangeStream(device, stream)) {}
+
+StreamGuard::~StreamGuard() {
+  impl_->ExchangeStream(device_guard_.current_device(), previous_stream_);
+  // device_guard_ restores the original device after the stream has been
+  // restored.
+}
+
+}  // namespace infini_train::core

tests/backend/CMakeLists.txt

@@ -7,3 +7,13 @@ infini_train_add_test(test_privateuse1_backend_validation
   SOURCES test_privateuse1_backend_validation.cc
   LABELS cpu
 )
+
+if(USE_CUDA)
+  infini_train_add_test(test_stream_guard_cuda
+    SOURCES test_stream_guard.cc
+    LABELS cuda accelerator hardware
+    BINARY_REGISTRATION
+    RUN_SERIAL
+    TEST_TIMEOUT 60
+  )
+endif()