@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
patch 92 lines 3.61 KB
Raw
1 From 18c0599c2f07ec965132efa25961dc8179c2dda3 Mon Sep 17 00:00:00 2001
2 From: Yongji Wu <wuyongji317@gmail.com>
3 Date: Tue, 20 May 2025 13:41:12 -0700
4 Subject: [PATCH] fix reinit issues due to states not cleaned up
5
6 fix double free
7 ---
8 src/host/init/init.cu | 10 ++++++++++
9 .../internal/host/nvshmemi_mem_transport.hpp | 15 +++++++++++++++
10 src/modules/bootstrap/uid/bootstrap_uid.cpp | 5 +++++
11 3 files changed, 30 insertions(+)
12
13 diff --git a/src/host/init/init.cu b/src/host/init/init.cu
14 index b1c5dbf..1fecb4b 100644
15 --- a/src/host/init/init.cu
16 +++ b/src/host/init/init.cu
17 @@ -43,6 +43,8 @@
18 #include "internal/host/nvshmemi_types.h"
19 #include "internal/host/shared_memory.h"
20 #include "internal/host/nvshmemi_symmetric_heap.hpp"
21 +// eep-dev
22 +#include "internal/host/nvshmemi_mem_transport.hpp"
23
24 extern __constant__ nvshmemi_device_host_state_t nvshmemi_device_state_d;
25 static std::map<void *, int> registered_device_states;
26 @@ -1293,6 +1295,14 @@ void nvshmemid_hostlib_finalize(void *device_ctx, void *transport_device_ctx) {
27 /* Multi-init Multi-fini*/
28 nvshmemi_state = NULL;
29 nvshmemi_device_state.nvshmemi_is_nvshmem_initialized = 0;
30 +
31 + // eep-dev
32 + nvshmemi_mem_p2p_transport::destroy_instance();
33 + nvshmemi_mem_remote_transport::destroy_instance();
34 + free(nvshmemi_default_session);
35 + nvshmemi_default_session = nullptr;
36 + nvshmemi_device_state.nvshmemi_is_nvshmem_bootstrapped = false;
37 +
38 nvshmemi_is_device_state_ready = false;
39 } else
40 nvshmemi_boot_handle.barrier(&nvshmemi_boot_handle);
41 diff --git a/src/include/internal/host/nvshmemi_mem_transport.hpp b/src/include/internal/host/nvshmemi_mem_transport.hpp
42 index 2495844..e4f408a 100644
43 --- a/src/include/internal/host/nvshmemi_mem_transport.hpp
44 +++ b/src/include/internal/host/nvshmemi_mem_transport.hpp
45 @@ -36,6 +36,13 @@ class nvshmemi_mem_p2p_transport final {
46 return p2p_objref_;
47 }
48 }
49 + // eep-dev
50 + static void destroy_instance(void) {
51 + if (p2p_objref_ != nullptr) {
52 + delete p2p_objref_;
53 + p2p_objref_ = nullptr;
54 + }
55 + }
56
57 void print_mem_handle(int pe_id, int transport_idx, nvshmemi_symmetric_heap &obj);
58
59 @@ -87,6 +94,14 @@ class nvshmemi_mem_remote_transport final {
60 }
61 }
62
63 + // eep-dev
64 + static void destroy_instance(void) {
65 + if (remote_objref_ != nullptr) {
66 + delete remote_objref_;
67 + remote_objref_ = nullptr;
68 + }
69 + }
70 +
71 int gather_mem_handles(nvshmemi_symmetric_heap &obj, uint64_t heap_offset, size_t size);
72 /* On-demand registration and release of memory */
73 int register_mem_handle(nvshmem_mem_handle_t *local_handles, int transport_idx,
74 diff --git a/src/modules/bootstrap/uid/bootstrap_uid.cpp b/src/modules/bootstrap/uid/bootstrap_uid.cpp
75 index a1fa748..788fa96 100644
76 --- a/src/modules/bootstrap/uid/bootstrap_uid.cpp
77 +++ b/src/modules/bootstrap/uid/bootstrap_uid.cpp
78 @@ -630,6 +630,11 @@ int nvshmemi_bootstrap_plugin_pre_init(bootstrap_handle_t* handle, const int abi
79 // Discover the network for bootstrap, if not done previously.
80 // This code needs to be stateful to be able to be called multiple times by the caller
81 BOOTSTRAP_CHECK(bootstrap_net_init());
82 + // eep-dev
83 + if (handle->pre_init_ops != nullptr) {
84 + BOOTSTRAP_PTR_FREE(handle->pre_init_ops);
85 + handle->pre_init_ops = nullptr;
86 + }
87 if (handle->pre_init_ops == nullptr) {
88 BOOTSTRAP_CALLOC(&handle->pre_init_ops, 1);
89 handle->pre_init_ops->get_unique_id = bootstrap_get_unique_id;
90 --
91 2.43.0
92