diff --git a/src/usock/ums/kmod/ums/cm/ums_core.c b/src/usock/ums/kmod/ums/cm/ums_core.c
index 82b087f3a3b9947a1bbe556546551291c4fb1750..5b014e3d6d5c7b8b85c7ac886be56f8aee66c07f 100644
--- a/src/usock/ums/kmod/ums/cm/ums_core.c
+++ b/src/usock/ums/kmod/ums/cm/ums_core.c
@@ -211,7 +211,7 @@ static int ums_lgr_register_conn(struct ums_connection *conn, bool first)
 	 */
 	sock_hold(&ums->sk); /* sock_put in ums_lgr_unregister_conn() */
 	while (conn->conn_id == 0) {
-		conn->conn_id = (u32)atomic_inc_return(&nexttoken);
+		conn->conn_id = (u32)atomic_inc_return(&nexttoken) & UMS_CONN_ID_MASK;
 		if (ums_lgr_find_conn(conn->conn_id, conn->lgr))
 			conn->conn_id = 0;
 	}
diff --git a/src/usock/ums/kmod/ums/cm/ums_core.h b/src/usock/ums/kmod/ums/cm/ums_core.h
index 296c6aa54a5dedabecc28f5b37b0512dbf2f1e0e..50644bea184df1ebb65b385d536153555a8b43eb 100644
--- a/src/usock/ums/kmod/ums/cm/ums_core.h
+++ b/src/usock/ums/kmod/ums/cm/ums_core.h
@@ -21,6 +21,7 @@
 #include "ums_ubcore.h"
 
 #define GID_SPRINTF_BUFF_SIZE 40
+#define UMS_CONN_ID_MASK 0x3FFFFF
 
 struct ums_rtoken_key_info {
 	__be32 nw_rkey;
diff --git a/src/usock/ums/kmod/ums/io/ums_tx.c b/src/usock/ums/kmod/ums/io/ums_tx.c
index 88e7a002c830f170a67586332e2b1c204f402071..a7f9b890d455c2f98b4d84418088677c5a152548 100644
--- a/src/usock/ums/kmod/ums/io/ums_tx.c
+++ b/src/usock/ums/kmod/ums/io/ums_tx.c
@@ -424,7 +424,7 @@ static inline u32 ums_get_imm(struct ums_connection *conn, int skip_flag)
 	/* in DMS, we store conn->token, credit and write_blocked in the IMM */
 	imm.write_blocked = conn->local_tx_ctrl.prod_flags.write_blocked;
 	imm.skip_flag = ((u32)skip_flag) & 0x1;
-	imm.token = conn->local_tx_ctrl.token & 0xFFFFFF;
+	imm.token = conn->local_tx_ctrl.token & UMS_CONN_ID_MASK;
 	if (ums_wr_rx_credits_need_announce_frequent(link))
 		saved_credits = (u8)ums_wr_rx_get_credits(link);
 	imm.credits = saved_credits;