[nccl-ep] Add ncclEpUpdateHandle to rebind topk_idx without reallocating

[kwen2501]

Summary

• Add ncclEpUpdateHandle API that rebinds topk_idx on an existing handle without reallocating GPU buffers. This avoids unnecessary cudaMalloc/cudaFree traffic (both device-synchronizing) in MoE hot loops where only the routing changes between iterations.
• Refactor ncclEpCreateHandle to delegate the computation phase to ncclEpUpdateHandle, eliminating code duplication.
• Fix dense_prob_buffer sizing to use max_tokens_per_rank instead of num_tokens so the buffer is reusable across different token counts.

RFC: #2084

API

ncclResult_t ncclEpUpdateHandle(
    ncclEpHandle_t handle,
    const ncclNDTensor_t* topk_idx,
    ncclNDTensor_t* const* local_tensors,
    unsigned int num_local_tensors,
    cudaStream_t stream
);

Usage

ncclEpCreateHandle(&handle, group, topk_idx_0, ...);

// Each iteration: rebind routing without realloc
ncclEpUpdateHandle(handle, topk_idx_1, ...);
ncclEpDispatch(handle, ...);
ncclEpCombine(handle, ...);

ncclEpHandleDestroy(handle);

Changes

• nccl_ep.h: Add ncclEpUpdateHandle declaration
• nccl_ep.cc: Add ncclEpUpdateHandle implementation (resets buffers + re-runs convert/allgather/preprocess). Refactor ncclEpCreateHandle to be allocation-only, delegating computation to ncclEpUpdateHandle.
• nccl_wrapper.py: Add ctypes binding for ncclEpUpdateHandle

Co-authored-by: Claude noreply@anthropic.com

Made with Cursor