Kernel-level compute-communication overlapping for CP linear attention kernels via nvshmem

[icavan]

Description

Implement kernel-level compute-communication overlapping for context parallelism (CP) linear attention kernels using nvshmem.

Context

For distributed training with context parallelism, overlapping computation with inter-GPU communication at the kernel level can reduce latency. Using nvshmem enables fine-grained, GPU-initiated communication directly from within CUDA kernels, avoiding CPU-side synchronization overhead.

Tasks

• Design the compute-communication overlap strategy for CP linear attention
• Integrate nvshmem for inter-GPU communication within kernels
• Implement overlapping for KDA and/or GDN CP kernels
• Benchmark communication overhead and end-to-end training throughput
• Test across multi-GPU configurations

[Butterfingrz]

Hi! I have a question regarding memory management when using NVSHMEM.

If I want to use nvshmem to achieve communication-computation overlap, does the GPU memory management for this need to be handled outside of the nn.Module? For example, in a scenario where a kernel's output tensor is directly allocated via nvshmem.

[icavan]

Hi! I have a question regarding memory management when using NVSHMEM.

If I want to use nvshmem to achieve communication-computation overlap, does the GPU memory management for this need to be handled outside of the nn.Module? For example, in a scenario where a kernel's output tensor is directly allocated via nvshmem.

Hi @Butterfingrz,

Great question! My feeling here is yes—the GPU memory management for these NVSHMEM buffers probably needs to be handled explicitly, outside the standard dynamic allocation loop of an nn.Module's forward pass.

Here is why I think so:
It mainly comes down to how NVSHMEM manages memory. NVSHMEM strictly requires symmetric memory allocation across all participating GPUs. PyTorch's default Caching Allocator, on the other hand, dynamically assigns asymmetric memory addresses. If we let PyTorch allocate the output tensor on the fly inside the kernel (e.g., via torch::empty()), that memory won't reside on the NVSHMEM symmetric heap. Calling NVSHMEM APIs (like nvshmem_put or nvshmem_get) on such addresses usually leads to invalid memory access or undefined behavior.

How I would suggest handling it:

Pre-allocation and Custom Deleters

Instead of allocating memory dynamically through PyTorch during the forward pass, my suggestion would be to allocate the symmetric buffers using nvshmem_malloc.
A great workaround to integrate this into PyTorch is to wrap this raw NVSHMEM pointer into a standard PyTorch Tensor using the C++ API torch::from_blob, which natively supports passing a custom deleter function.
The key focus here is guaranteeing the validity of the memory's lifecycle. PyTorch's allocator has no knowledge of NVSHMEM memory. By attaching a custom deleter, you tie the NVSHMEM pointer's lifespan directly to the Tensor's reference counting. This ensures the memory remains perfectly valid while computations are running, and allows you to safely free the NVSHMEM memory all at once at the very end of the operator (or when the tensor's lifecycle naturally ends). This completely prevents PyTorch from accidentally triggering a standard cudaFree and causing a crash.

[Butterfingrz]

Thanks for the reply! This is very insightful.