Skip to content

[Bugfix][NVSHMEM] Support v2 state in legacy IBGDA kernels #765

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
XFDG wants to merge 1 commit into from
Closed

[Bugfix][NVSHMEM] Support v2 state in legacy IBGDA kernels #765

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions csrc/kernels/legacy/ibgda_device.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
#pragma once


#include <type_traits>

#include <nvshmem.h>
#include <device_host_transport/nvshmem_common_ibgda.h>
#include <non_abi/device/threadgroup/nvshmemi_common_device_defines.cuh>
Expand Down Expand Up @@ -78,11 +80,21 @@ __device__ static __forceinline__ nvshmemi_ibgda_device_state_t* ibgda_get_state
return &nvshmemi_ibgda_device_state_d;
}

template <typename StateType>
__device__ static __forceinline__ nvshmemi_ibgda_device_qp_t* ibgda_get_rc_impl(StateType* state, int pe, int id) {
const auto num_rc_per_pe = state->num_rc_per_pe;

if constexpr (std::is_same_v<StateType, nvshmemi_ibgda_device_state_v1>) {
return &state->globalmem
.rcs[pe * num_rc_per_pe * state->num_devices_initialized + id % (num_rc_per_pe * state->num_devices_initialized)];
} else {
return &state->globalmem.rcs[pe + nvshmemi_device_state_d.npes * id];
}
}

__device__ static __forceinline__ nvshmemi_ibgda_device_qp_t* ibgda_get_rc(int pe, int id) {
auto state = ibgda_get_state();
const auto num_rc_per_pe = ibgda_get_state()->num_rc_per_pe;
return &state->globalmem
.rcs[pe * num_rc_per_pe * state->num_devices_initialized + id % (num_rc_per_pe * state->num_devices_initialized)];
return ibgda_get_rc_impl(state, pe, id);
}

__device__ static __forceinline__ void ibgda_lock_acquire(int* lock) {
Expand Down
7 changes: 7 additions & 0 deletions deep_ep/include/deep_ep/common/compiled.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,13 @@
#define __CUDACC__
#endif

// Define __CUDACC_RDC__ so NVSHMEM device symbols use the correct extern declarations.
#ifndef DISABLE_NVSHMEM
#ifndef __CUDACC_RDC__
#define __CUDACC_RDC__ // NOLINT(*-reserved-identifier)
#endif
#endif

// Remove Torch restrictions
#ifdef __CUDA_NO_HALF_CONVERSIONS__
#undef __CUDA_NO_HALF_CONVERSIONS__
Expand Down