Hi @atalman @malfet do you mind having a look at this section? Today we gate the NVSHMEM feature with build-time flag USE_NVSHMEM=1. We may change it later.

Why do these two APIs have the nvshmem_ prefix? Do we plan to unify nvshmem_reduce_scatter_out with reduce_scatter_out?

Correct. I expect we can remove the prefix by the release.

Just for my learning, we need to use nvshmemx_collective_launch because nvshmem_all_reduce_kernel uses some nvshmem APIs. Is this correct?

From documentation:

The nvshmemx_collective_launch function must be used to launch CUDA kernels on the GPU when the CUDA kernels use NVSHMEM synchronization or collective APIs (e.g., nvshmem_wait, nvshmem_barrier, nvshmem_barrier_all, or any other collective operation). CUDA kernels that do not use synchronizing NVSHMEM APIs (or that do not use NVSHMEM APIs at all), are not required to be launched by this API. This call is collective across the PEs in the NVSHMEM job.

what's your plan for documenting this to users?
a) how long-lived is it? (iiuc, it's long-lived since we have to support 'Cuda' backend for AMD
b) what are the user-facing tradeoffs? Why would i choose CUDASymmetricMemory vs NVSHMEMSymmetricMemory in a given scenario?

a) The env itself could be longer term due to possibility of "NCCL" backend. The "CUDASymmetricMemory" impl would stay as long as needed for the reason you mentioned.
b) intra vs inter capability is the main difference from user point of view.

(b) is still not clear to me. Does nvshem backend support intra+inter or is it purely intra?

if nvshmem supports the superset of cuda backend except for AMD support, then I would argue we delete the env, and autodetect AMD to switch it to the cuda backend. otherwise, default to the nvshmem backend?

The env is also for NCCLSymmetricMemory vs NVSHMEMSymmetricMemory.

i was gonna suggest having a dedicated folder for the nvshmem backend files so its clear to readers that they are built into a separate compilation unit.

however maybe its an annoyingly long path? csrc/distributed/c10d/symm_mem/common and /nvshmem

Yeah, totally agree with a dedicated folder. The c10d folder is a mix of everything today, will refactor together with other libraries later.

maybe we can have a symmetric memory dedicated folder?

you should really do the code-motion part in a previous PR so its easier to tell what is real changes

3 functions moved to utils file; not checking the details, assuming its just a direct move

planning to document these ENVs? any idea if they are going to be long-lived?

We can probably clean this env.

(b) is still not clear to me. Does nvshem backend support intra+inter or is it purely intra?

if nvshmem supports the superset of cuda backend except for AMD support, then I would argue we delete the env, and autodetect AMD to switch it to the cuda backend. otherwise, default to the nvshmem backend?

just curious, since this is declared as 'extern' in the header, what does that mean when there is a definition in both torch and nvshmem backend? won't the torch definition always exist and take precedence even when we are using nvshmem backend? (did you confirm in tests that the nvshmem backend prefix is being used)

You are right that usually compiler does not allow defining the same variable in multiple source files.
Here we are enjoying a special rule for extern const where static linkage applies -- meaning the variable is not visible to other translation units.

@wconstab thanks for your suggestion. I made a little class StoreExchange to encapsulate those store_ methods and the prefix. Then each SymmMem implementation can locally instantiate a StoreExchange instance to avoid namespace & prefix contention.

nit: would it be clearer to do (void) reduce_op in the meta kernel? its not a big deal though

I think dispatcher would complain about the signature mismatch.

no, you keep the signature.

in theory this is the preferred way to tell the compiler/linter that you're intentionally not using a provided value and iiuc it should silence linter errors, but .. your way is fine too

void foo(int x){
   (void)x;
   return;
}

Ah I see what you mean now. Thanks for the suggestion!
In this specific case, the linter wants me to change std::string to const std::string& to avoid a copy in argument passing. But const std::string& is unfortunately not supported by torch dispatcher signature definition (if I am not mistaken).

is this NYI or NotImplemented? (Will we implement it for Cuda backend? should this be documented? (why is this different than c10d existing rank mapping helpers)

Yes, it stands for unimplemented.
Good question whether other backends need it. Perhaps I should move them into NVSHMEMSymmetricMemory for now.

i didn't review the IpcChannel code yet, maybe someone else can take a close look there

It is a move. Sorry

Still want to ask the nvshmem_ prefix. Do we imagine that in the future different backend will export different broadcast or should we just have one broadcast and error out if the underlying backend doesn't support it.

The latter I guess.

one minor nit: the CUDA backend puts the ops (e.g., one_shot_all_reduce) in CUDASymmetricMemoryOps.cu. Do we want to follow that convention?

It is a matter of renaming this file. I think we can do it later (after we sedate the matter of whether NVSHMEM should be a SymmetricMemory class by itself or an extension to CUDASymmetricMemory)