common: tighten peer-handle docstrings

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Author: phu0ngng

transformer_engine/common/comm_handle.cpp

@@ -6,9 +6,8 @@
 
 #include "transformer_engine/comm_handle.h"
 
-#include "transformer_engine/nccl_comm.h"
-
 #include "common.h"
+#include "transformer_engine/nccl_comm.h"
 #include "util/logging.h"
 
 using transformer_engine::convertNVTETensor;

transformer_engine/common/common.h

@@ -180,16 +180,11 @@ struct Tensor {
    */
   bool row_scaled_nvfp4 = false;
 
-  /*! \brief Optional peer-handle annotation for zero-copy comm fast paths.
-   *
-   *  Borrowed reference to a comm-backend-specific resource (NCCL window,
-   *  NVSHMEM pointer, CUDA IPC handle, etc.) that lets a consumer initiate
-   *  one-sided remote-memory ops against this tensor's storage. The kind tag
-   *  selects which backend owns ``peer_handle_data``; the tensor never owns
-   *  the resource — the caller keeps it valid for the tensor's lifetime.
-   *  Code paths that don't issue cross-rank ops ignore these fields. */
+  /*! \brief Optional borrowed peer handle for one-sided RMA against this tensor.
+   *         ``peer_handle_kind`` selects the backend owning ``peer_handle_data``;
+   *         the caller keeps the resource valid for the tensor's lifetime. */
   NVTEPeerHandleKind peer_handle_kind = NVTE_PEER_HANDLE_NONE;
-  void* peer_handle_data = nullptr;
+  void *peer_handle_data = nullptr;
   uint64_t peer_handle_offset = 0;
 
   /*! Map from NVTETensorParam to parameter sizes */

transformer_engine/common/include/transformer_engine/comm_handle.h

@@ -5,17 +5,10 @@
  ************************************************************************/
 
 /*! \file comm_handle.h
- *  \brief Generic peer-handle annotation on NVTETensor.
+ *  \brief Generic peer-handle annotation on NVTETensor for one-sided RMA.
  *
- *  A peer handle is an opaque, comm-backend-specific reference that lets a
- *  consumer initiate one-sided remote-memory operations against a peer's
- *  buffer (e.g. NCCL symmetric-memory window, NVSHMEM pointer, CUDA-IPC
- *  handle). The annotation is borrowed; the resource outlives the call and
- *  the tensor never owns it.
- *
- *  Backends register their setter/getter under a dedicated header (e.g.
- *  ``nccl_comm.h`` for NCCL windows). This header exposes only the kind tag
- *  and detach, both of which are payload-agnostic.
+ *  The annotation is borrowed; the tensor never owns the underlying resource.
+ *  Per-backend setters/getters live in dedicated headers (e.g. ``nccl_comm.h``).
  */
 
 #ifndef TRANSFORMER_ENGINE_COMM_HANDLE_H_
@@ -27,17 +20,16 @@
 extern "C" {
 #endif
 
-/*! \brief Kind tag identifying which comm backend owns a tensor's peer handle. */
+/*! \brief Comm backend that owns a tensor's peer handle. */
 typedef enum {
   NVTE_PEER_HANDLE_NONE = 0,
   NVTE_PEER_HANDLE_NCCL_WINDOW = 1,
-  /* Reserved for future backends: NVSHMEM_PTR, CUDA_IPC, UCX_RKEY, ... */
 } NVTEPeerHandleKind;
 
-/*! \brief Return the peer-handle kind currently attached to ``t``. */
+/*! \brief Peer-handle kind attached to ``t``. */
 NVTEPeerHandleKind nvte_tensor_peer_handle_kind(const NVTETensor t);
 
-/*! \brief Clear any peer handle attached to ``t``; no-op when none is set. */
+/*! \brief Clear any peer handle attached to ``t``. */
 void nvte_tensor_detach_peer_handle(NVTETensor t);
 
 #ifdef __cplusplus

transformer_engine/common/include/transformer_engine/nccl_comm.h

@@ -5,12 +5,10 @@
  ************************************************************************/
 
 /*! \file nccl_comm.h
- *  \brief NCCL-backed peer-handle setter/getter for NVTETensor.
+ *  \brief Attach a registered NCCL symmetric-memory window to an NVTETensor.
  *
- *  Attaches a registered NCCL symmetric-memory window + byte offset onto a
- *  tensor so consumers (e.g. the EP backend) can issue one-sided put/get over
- *  the window instead of staging through the raw data pointer. The window is
- *  caller-owned; ``attach`` does not register or rendezvous it.
+ *  The window is caller-owned and must outlive the tensor; ``attach`` does
+ *  not register or rendezvous it.
  */
 
 #ifndef TRANSFORMER_ENGINE_NCCL_COMM_H_
@@ -23,23 +21,16 @@
 extern "C" {
 #endif
 
-/*! \brief Attach an NCCL window + byte offset to ``t``.
- *
- *  Sets the tensor's peer-handle kind to ``NVTE_PEER_HANDLE_NCCL_WINDOW``. The
- *  window must stay registered for the tensor's lifetime. Pass ``window=NULL``
- *  to detach (equivalent to ``nvte_tensor_detach_peer_handle``).
+/*! \brief Attach an NCCL window + byte offset to ``t``. Pass ``window=NULL`` to detach.
  *
  *  \param[in,out] t      Tensor to annotate.
  *  \param[in]     window Opaque ncclWindow_t (caller-owned), or NULL to clear.
  *  \param[in]     offset Byte offset into the window where this tensor starts.
  */
 void nvte_tensor_attach_nccl_window(NVTETensor t, void* window, uint64_t offset);
 
-/*! \brief Read the NCCL window + offset attached to ``t``.
- *
- *  ``*window`` is set to ``NULL`` and ``*offset`` to 0 when no NCCL window is
- *  attached (including when the tensor carries a different peer-handle kind).
- *  Either out-pointer may be ``NULL`` to skip that field.
+/*! \brief Read the NCCL window + offset attached to ``t``; yields (NULL, 0) when unset.
+ *         Either out-pointer may be NULL to skip that field.
  */
 void nvte_tensor_nccl_window(const NVTETensor t, void** window, uint64_t* offset);