UCT/IB/MLX5: add WQE token and completion suppress by jeynmann · Pull Request #11684 · openucx/ucx

jeynmann · 2026-07-21T02:33:33Z

What?

Save tokens on the tx side and add support for completion suppress.

Why?

Need to keep tokens on the tx side so the plugin can purge outstanding operations using tokens received from the remote side.
Need to suppress the completions ailover and tx progress do not act on the wqe and pending operations together.

How?

Add completion suppress support in uct ep invalidte.
Add an flag to indicate the completion ownership moved to plugin.

evgeny-leksikov

let's move UCP changes to next PR for simplisity.
Please keep here only UCT API + IMPL + new UCT tests with new behavior:
ep_invalidate(NO_COMPLETION)
post new ops // to be consistent we need to post WQEs here
short_progress_loop // make sure no completions
purge_outstanding(EP) // all completions (for ops before and after invaslidate) appears here

evgeny-leksikov · 2026-08-10T08:11:19Z

    void                                *arg;
+
+    /** Completion status for purging all outstanding operations. */
+    ucs_status_t                        status;


is this in or out parameter? what is purpose of the status? I would expect that status of any outstanding operation is unknown.

evgeny-leksikov · 2026-08-10T08:16:34Z

+    if (!(ep->super.ext_flags & UCT_RC_EP_EXT_FLAG_FAILOVER_ARMED)) {
+        outstanding = txwq->bb_max - uct_rc_txqp_available(&ep->super.txqp);
+        txwq->ft_ci = txwq->prev_sw_pi - outstanding;
+        ep->super.ext_flags |= UCT_RC_EP_EXT_FLAG_FAILOVER_ARMED;


how does it work? AFAIR per offline discussion, ep_invalidate should suppress completions only by new flag in uct_ep_invalidate_params_t

evgeny-leksikov · 2026-08-10T08:17:31Z

+        outstanding = txwq->bb_max - uct_rc_txqp_available(&ep->super.txqp);
+        txwq->ft_ci = txwq->prev_sw_pi - outstanding;
+        ep->super.ext_flags |= UCT_RC_EP_EXT_FLAG_FAILOVER_ARMED;
+        ucs_assert(txwq->ft_ci == txwq->hw_ci);


how about to swap the logic:
assign txwq->ft_ci = txwq->hw_ci
and assert the calculation only for debugging

evgeny-leksikov · 2026-08-10T08:24:38Z

+    base_iface = ucs_derived_of(&iface->super.super.super, uct_base_iface_t);
+    is_flush_cancel = ep->super.flags & UCT_RC_EP_FLAG_FLUSH_CANCEL;
+    if ((base_iface->err_handler == NULL) || is_flush_cancel) {
+        uct_rc_txqp_purge_outstanding(iface, &ep->super.txqp, ep_status, pi, 0);
+    }


can we simplify? smth like

if (!(ep->super.flags & UCT_RC_EP_FLAG_SUPPRESS_COMPLETIONS)) { uct_rc_txqp_purge_outstanding(iface, &ep->super.txqp, ep_status, pi, 0); }

evgeny-leksikov · 2026-08-10T08:27:17Z

+    uint32_t                    next_token; /* Token assigned to the next packet */
+    uint32_t                    *tokens;    /* Token indexed by WQEBB */
+    uint16_t                    token_mask; /* Mask for token array */


why do we need the array? can we get each wqe size traversing WQ during purge?

evgeny-leksikov · 2026-08-10T08:27:27Z


+enum {
+    /* Failover owns this endpoint's outstanding WQEs and send operations. */
+    UCT_RC_EP_EXT_FLAG_FAILOVER_ARMED = UCS_BIT(0)


need to rename

evgeny-leksikov · 2026-08-11T07:45:53Z

+ */
+typedef enum {
+    /** Move the endpoint QP to the error state. */
+    UCT_EP_INVALIDATE_FLAG_MODIFY_QP_TO_ERR     = UCS_BIT(0),


why needed?

we dont use IB terms in UCX API

evgeny-leksikov · 2026-08-11T07:50:36Z

+     * Suppress completion of outstanding operations by the normal endpoint
+     * error path. The caller transfers outstanding-operation ownership to an
+     * external provider and is responsible for completing them. This state
+     * remains active until endpoint cancellation or destruction. This flag
+     * does not modify QP state.
+     */
+    UCT_EP_INVALIDATE_FLAG_SUPPRESS_COMPLETIONS = UCS_BIT(1)


IMO "suppress" is not correct wording here since completions of posted operations won't be handled by normal iface_progress and that's user's responsibility to "purge" them.

evgeny-leksikov · 2026-08-11T07:51:53Z

- * @ref uct_ep_outstanding_purge_params_t::cb is invoked once for each
- * undelivered outstanding operation, in the original endpoint posting order.
+ * The transport delegates outstanding-operation classification to its
+ * external provider. The provider invokes @ref


what is external provider?

evgeny-leksikov · 2026-08-11T07:52:52Z

+ * in the original endpoint posting order. No transport-local status purge is
+ * provided by this API.


why need to mention status here?

evgeny-leksikov · 2026-08-11T07:53:39Z

+        !(params->field_mask & UCT_EP_INVALIDATE_PARAM_FIELD_FLAGS)) {
+        return UCS_ERR_INVALID_PARAM;


does not it break compatibility?

evgeny-leksikov · 2026-08-11T07:54:48Z

           uct_rc_mlx5_base_ep_is_connected(tl_ep, params);
 }

+static void uct_rc_gga_iface_handle_failure(uct_ib_iface_t *ib_iface, void *arg,


why do we need special err_hadler for GGA?

evgeny-leksikov · 2026-08-11T07:56:15Z

+        message_length  = length + sizeof(*am);
        fm_ce_se        |= uct_rc_iface_tx_moderation(&iface->super, txqp, MLX5_WQE_CTRL_CQ_UPDATE);


evgeny-leksikov · 2026-08-11T07:58:01Z

    ucs_trace_poll("rc_mlx5 iface %p tx_cqe: ep %p qpn 0x%x hw_ci %d", iface,
                   ep, qp_num, hw_ci);

+    ucs_assert(!(ep->flags & UCT_RC_MLX5_EP_FLAG_SUPPRESS_COMPLETIONS));


evgeny-leksikov · 2026-08-11T08:03:35Z

    txwq->prev_sw_pi = UINT16_MAX;
-#if UCS_ENABLE_ASSERT
+    txwq->ft_ci      = UINT16_MAX;
+    txwq->next_token = 0;


is it still needed?

IMO, yes.

For ft_ci:

ft_ci is used to save which wqe to start purge.

ft_ci is set when invalidate being called.

ft_ci will be updated by plugin, so we can retry/reentrant from last pos.

ft_ci might be diff from hw_ci.

For next token:

Used to track token for each message.

Used to correct token when message is partially delivered.

ft_ci is clear, my question is about next_token:

Used to correct psn when message is partially delivered

what does happen is this case?
I do not see why we cannot just repost full message:

In case of put/get, previously arrived data will be over written.

What happens in case of send? I would expect it should not generate CQE (with err) on sender side and it mean that data will be dropped and we need to resend full message any way.

do I miss something else?

As discussed offline, next token is the first token for the next message. We can't distinguish whether it's first or middle from HW.

IMO we need to repost entire undelivered QWE, even if failure happened in the middle. Otherwise we lose data [0; first_failed_packet) for SEND op code, partially delivered QWE won't consume CQE and (partial) data will be dropped on receiver side.

evgeny-leksikov · 2026-08-11T08:05:16Z

+        if (qp->type == UCT_IB_MLX5_OBJ_TYPE_DEVX) {
+            char in[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_in)]   = {};
+            char out[UCT_IB_MLX5DV_ST_SZ_BYTES(query_qp_out)] = {};
+            void *qpc;
+            int ret;
+
+            UCT_IB_MLX5DV_SET(query_qp_in, in, opcode,
+                              UCT_IB_MLX5_CMD_OP_QUERY_QP);
+            UCT_IB_MLX5DV_SET(query_qp_in, in, qpn, qp->qp_num);
+            ret = mlx5dv_devx_obj_query(qp->devx.obj, in, sizeof(in), out,
+                                        sizeof(out));
+            if (ret != 0) {
+                return UCS_ERR_IO_ERROR;
+            }
+
+            qpc  = UCT_IB_MLX5DV_ADDR_OF(query_qp_out, out, qpc);
+            *psn = UCT_IB_MLX5DV_GET(qpc, qpc, next_send_psn);
+            return UCS_OK;
+        }


why do we have it in tests?

To make sure token calculation matches with HW.

gleon99 · 2026-08-12T08:36:49Z

@roiedanino .

evgeny-leksikov · 2026-08-12T12:55:29Z

+    /**
+     * Defer completion of outstanding operations to the caller. The normal
+     * endpoint error path will not complete these operations; the caller must
+     * purge them with @ref uct_ep_outstanding_purge. A successful purge returns
+     * ownership of subsequent operations to UCT.
+     */
+    UCT_EP_INVALIDATE_FLAG_DEFER_COMPLETIONS = UCS_BIT(0)


Suggested change

/**

* Defer completion of outstanding operations to the caller. The normal

* endpoint error path will not complete these operations; the caller must

* purge them with @ref uct_ep_outstanding_purge. A successful purge returns

* ownership of subsequent operations to UCT.

*/

UCT_EP_INVALIDATE_FLAG_DEFER_COMPLETIONS = UCS_BIT(0)

/**

* Do not complete outstanding operations. UCT will not invoke completion

* callbacks for issued operations when the endpoint is invalidated;

* the caller must complete them with @ref uct_ep_outstanding_purge.

*/

UCT_EP_INVALIDATE_FLAG_NO_COMPLETIONS = UCS_BIT(0)

evgeny-leksikov · 2026-08-12T13:01:31Z

 * This routine invalidates the endpoint and moves it to the error state.
- * All the incomplete and subsequent operations on the endpoint will be
- * completed with error.
+ * Incomplete operations are completed with error unless
+ * @ref UCT_EP_INVALIDATE_FLAG_DEFER_COMPLETIONS transfers their ownership to
+ * the caller.


pls revert, we dont need to mention all possible flags in the function description, only basic behavior

evgeny-leksikov · 2026-08-12T13:05:50Z

     *  uct_ep_outstanding_purge_field_t. @ref
     *  UCT_EP_OUTSTANDING_FIELD_RX_TOKEN and @ref
-     *  UCT_EP_OUTSTANDING_FIELD_CB are required. */
+     *  UCT_EP_OUTSTANDING_FIELD_CB are required for token-based purging. */


do we have "something-else"-based purging? are rx_tocken and cb mandatory? check doc for field_mask in other params-like structures, it's only for compatibility, keep fields self-documented and add cross references when it's really needed

evgeny-leksikov · 2026-08-12T13:17:18Z

+
+#if HAVE_MLX5_MMO
+    if ((qp_type == IBV_QPT_RC) && (opcode != MLX5_OPCODE_NOP) &&
+        (opcode != MLX5_OPCODE_MMO)) {


why is MMO exception?

evgeny-leksikov · 2026-08-12T13:26:11Z

+        return status;
+    }
+
+    if ((params != NULL) &&


I think check params == NULL is redundant, if we call it somewhere this way, let's fix it.

evgeny-leksikov · 2026-08-12T13:26:26Z

+        !(ep->super.flags & UCT_RC_MLX5_EP_FLAG_DEFER_COMPLETIONS)) {
+        ep->super.flags |= UCT_RC_MLX5_EP_FLAG_DEFER_COMPLETIONS;


Is it correct to call ep_invalidate twice?

Added assertion instead.

evgeny-leksikov · 2026-08-12T13:28:19Z

+    status = uct_ib_mlx5_ext_ep_outstanding_purge(tl_ep, params);
+    if (status == UCS_OK) {
+        ep->flags &= ~UCT_RC_MLX5_EP_FLAG_DEFER_COMPLETIONS;
+    }
+
+    return status;


code style:

Suggested change

status = uct_ib_mlx5_ext_ep_outstanding_purge(tl_ep, params);

if (status == UCS_OK) {

ep->flags &= ~UCT_RC_MLX5_EP_FLAG_DEFER_COMPLETIONS;

}

return status;

status = uct_ib_mlx5_ext_ep_outstanding_purge(tl_ep, params);

if (status != UCS_OK) {

return status;

}

ep->flags &= ~UCT_RC_MLX5_EP_FLAG_DEFER_COMPLETIONS;

return UCS_OK;

evgeny-leksikov · 2026-08-12T13:38:12Z

+void uct_rc_mlx5_ep_update_tx_res(uct_ep_h tl_ep)
+{
+    UCT_RC_MLX5_BASE_EP_DECL(tl_ep, iface, ep);
+    uct_ib_mlx5_txwq_t *txwq = &ep->tx.wq;
+    uint16_t available;
+
+    ucs_assert(ep->flags & UCT_RC_MLX5_EP_FLAG_DEFER_COMPLETIONS);
+    available = txwq->bb_max - (txwq->prev_sw_pi - txwq->hw_ci);
+    ucs_assert(available >= uct_rc_txqp_available(&ep->super.txqp));
+    if (available > uct_rc_txqp_available(&ep->super.txqp)) {
+        uct_rc_mlx5_iface_update_tx_res(&iface->super, ep, txwq->hw_ci);
+    }
+}


is this new function? where is it called from?

Renamed to uct_rc_mlx5_ep_update_tx_res

Refactored uct_rc_mlx5_iface_update_tx_res.

Normal path will update both iface and qp tx credits.

No completion path will update iface tx credits only.

Uct purge will update QP tx credits once fnished.

jeynmann force-pushed the failover_ext_uct_wqe_meta branch 4 times, most recently from 738e18e to 35acf5a Compare July 27, 2026 10:12

jeynmann changed the title ~~UCT/IB/MLX5: add wqe token and failover arm~~ [DNM] UCT/IB/MLX5: add wqe token and failover arm Jul 29, 2026

jeynmann force-pushed the failover_ext_uct_wqe_meta branch 3 times, most recently from e311e37 to d3448c3 Compare August 6, 2026 10:30

JeynmannZ added 6 commits August 7, 2026 15:18

UCT/RC/MLX5: Delegate failure purge to UCP

e386c6b

UCP/CORE: Delegate failure purge to UCP

8d98636

GTEST/UCT: Delegate failure purge to UCP

9f08a84

UCT: fix CI

5e6b26c

UCT/IB/MLX5: fix restart before reconfig.

5ad8404

UCP: set ep fail before purge

3328c1a

jeynmann force-pushed the failover_ext_uct_wqe_meta branch from 796b669 to 2d1fe63 Compare August 10, 2026 02:13

JeynmannZ added 4 commits August 10, 2026 06:36

UCT: address comments

0e19522

UCP: address comments

7804552

UCP: address comments

7885eec

UCT/IB: add meta for uct

3b2ec01

jeynmann force-pushed the failover_ext_uct_wqe_meta branch from 2d1fe63 to 3b2ec01 Compare August 10, 2026 06:08

UCT/IB/MLX5: refactor and fix CI

9aee3e5

evgeny-leksikov requested changes Aug 10, 2026

View reviewed changes

JeynmannZ added 3 commits August 10, 2026 20:08

UCT: address comments

cc5224d

UCP: remove ucp changes from this PR

31e7197

UCT/IB/MLX5: address comments

17ba322

jeynmann changed the title ~~[DNM] UCT/IB/MLX5: add wqe token and failover arm~~ UCT/IB/MLX5: add WQE token and completion suppress Aug 11, 2026

UCT/IB/MLX5: refactor token record

a4c34c9

jeynmann force-pushed the failover_ext_uct_wqe_meta branch 2 times, most recently from 934ec35 to 70dd829 Compare August 11, 2026 06:40

evgeny-leksikov reviewed Aug 11, 2026

View reviewed changes

JeynmannZ added 4 commits August 12, 2026 05:54

UCT/IB/MLX5: cleanup

7ba1e58

UCT/IB/MLX5: address comments

a4a68ee

UCT: fix format and comments

d2a7d17

UCT: cleanup protection for flush cancel

f9f10bf

jeynmann force-pushed the failover_ext_uct_wqe_meta branch from cc80d4e to f9f10bf Compare August 12, 2026 03:26

gleon99 requested review from evgeny-leksikov and roiedanino August 12, 2026 08:35

evgeny-leksikov reviewed Aug 12, 2026

View reviewed changes

JeynmannZ added 2 commits August 12, 2026 18:17

UCT: address comments

e6f3742

UCT: address comments

b99a63c

jeynmann force-pushed the failover_ext_uct_wqe_meta branch from ed603a8 to b99a63c Compare August 13, 2026 04:52

UCT: fix CI

7870232

		* in the original endpoint posting order. No transport-local status purge is
		* provided by this API.

		!(params->field_mask & UCT_EP_INVALIDATE_PARAM_FIELD_FLAGS)) {
		return UCS_ERR_INVALID_PARAM;

		message_length = length + sizeof(*am);
		fm_ce_se \|= uct_rc_iface_tx_moderation(&iface->super, txqp, MLX5_WQE_CTRL_CQ_UPDATE);

		!(ep->super.flags & UCT_RC_MLX5_EP_FLAG_DEFER_COMPLETIONS)) {
		ep->super.flags \|= UCT_RC_MLX5_EP_FLAG_DEFER_COMPLETIONS;

Uh oh!

Conversation

jeynmann commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What?

Why?

How?

Uh oh!

evgeny-leksikov left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jeynmann Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jeynmann Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gleon99 commented Aug 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

jeynmann commented Jul 21, 2026 •

edited

Loading

jeynmann Aug 11, 2026 •

edited

Loading

jeynmann Aug 11, 2026 •

edited

Loading