diff --git a/openair2/LAYER2/NR_MAC_UE/config_ue.c b/openair2/LAYER2/NR_MAC_UE/config_ue.c
index d0675b30fe..2e349f8919 100644
--- a/openair2/LAYER2/NR_MAC_UE/config_ue.c
+++ b/openair2/LAYER2/NR_MAC_UE/config_ue.c
@@ -1930,29 +1930,54 @@ static void configure_BWPs(NR_UE_MAC_INST_t *mac, NR_ServingCellConfig_t *scd)
 
 static void handle_mac_uecap_info(NR_UE_MAC_INST_t *mac, NR_UE_NR_Capability_t *ue_Capability)
 {
-  if (ue_Capability->featureSets) {
-    if (ue_Capability->featureSets->featureSetsDownlinkPerCC) {
-      struct NR_FeatureSets__featureSetsDownlinkPerCC *fs_dlcc_list= ue_Capability->featureSets->featureSetsDownlinkPerCC;
-      for (int i = 0; i < fs_dlcc_list->list.count; i++) {
-        NR_FeatureSetDownlinkPerCC_t *fs_dl_cc = fs_dlcc_list->list.array[i];
-        if (mac->current_DL_BWP->scs != fs_dl_cc->supportedSubcarrierSpacingDL)
-          continue;
-        int uecap_bw_index;
-        if (fs_dl_cc->supportedBandwidthDL.present == NR_SupportedBandwidth_PR_fr1) {
-          uecap_bw_index = fs_dl_cc->supportedBandwidthDL.choice.fr1;
-          // 90 MHz option is indicated by a separate pointer in case indicated supported BW is 100MHz
-          // so we need to increase the index by 1 unit to point to 100 MHz if not 90MHz
-          if (uecap_bw_index == NR_SupportedBandwidth__fr1_mhz100 && !fs_dl_cc->channelBW_90mhz)
-            uecap_bw_index++;
-        }
-        else
-          uecap_bw_index = fs_dl_cc->supportedBandwidthDL.choice.fr2;
-        int dl_bw_mhz = mac->phy_config.config_req.carrier_config.dl_bandwidth;
-        if (dl_bw_mhz != get_supported_bw_mhz(mac->frequency_range, uecap_bw_index))
-          continue;
-        if (fs_dl_cc->maxNumberMIMO_LayersPDSCH)
-          mac->uecap_maxMIMO_PDSCH_layers = 2 << *fs_dl_cc->maxNumberMIMO_LayersPDSCH;
+  if (!ue_Capability->featureSets)
+    return;
+  if (ue_Capability->featureSets->featureSetsDownlinkPerCC) {
+    struct NR_FeatureSets__featureSetsDownlinkPerCC *fs_dlcc_list = ue_Capability->featureSets->featureSetsDownlinkPerCC;
+    for (int i = 0; i < fs_dlcc_list->list.count; i++) {
+      NR_FeatureSetDownlinkPerCC_t *fs_dl_cc = fs_dlcc_list->list.array[i];
+      if (mac->current_DL_BWP->scs != fs_dl_cc->supportedSubcarrierSpacingDL)
+        continue;
+      int uecap_bw_index;
+      if (fs_dl_cc->supportedBandwidthDL.present == NR_SupportedBandwidth_PR_fr1) {
+        uecap_bw_index = fs_dl_cc->supportedBandwidthDL.choice.fr1;
+        // 90 MHz option is indicated by a separate pointer in case indicated supported BW is 100MHz
+        // so we need to increase the index by 1 unit to point to 100 MHz if not 90MHz
+        if (uecap_bw_index == NR_SupportedBandwidth__fr1_mhz100 && !fs_dl_cc->channelBW_90mhz)
+          uecap_bw_index++;
       }
+      else
+        uecap_bw_index = fs_dl_cc->supportedBandwidthDL.choice.fr2;
+      int dl_bw_mhz = mac->phy_config.config_req.carrier_config.dl_bandwidth;
+      if (dl_bw_mhz != get_supported_bw_mhz(mac->frequency_range, uecap_bw_index))
+        continue;
+      if (fs_dl_cc->maxNumberMIMO_LayersPDSCH)
+        mac->uecap_maxMIMO_PDSCH_layers = 2 << *fs_dl_cc->maxNumberMIMO_LayersPDSCH;
+    }
+  }
+  if (ue_Capability->featureSets->featureSetsUplinkPerCC) {
+    struct NR_FeatureSets__featureSetsUplinkPerCC *fs_ulcc_list = ue_Capability->featureSets->featureSetsUplinkPerCC;
+    for (int i = 0; i < fs_ulcc_list->list.count; i++) {
+      NR_FeatureSetUplinkPerCC_t *fs_ul_cc = fs_ulcc_list->list.array[i];
+      if (mac->current_UL_BWP->scs != fs_ul_cc->supportedSubcarrierSpacingUL)
+        continue;
+      int uecap_bw_index;
+      if (fs_ul_cc->supportedBandwidthUL.present == NR_SupportedBandwidth_PR_fr1) {
+        uecap_bw_index = fs_ul_cc->supportedBandwidthUL.choice.fr1;
+        // 90 MHz option is indicated by a separate pointer in case indicated supported BW is 100MHz
+        // so we need to increase the index by 1 unit to point to 100 MHz if not 90MHz
+        if (uecap_bw_index == NR_SupportedBandwidth__fr1_mhz100 && !fs_ul_cc->channelBW_90mhz)
+          uecap_bw_index++;
+      }
+      else
+        uecap_bw_index = fs_ul_cc->supportedBandwidthUL.choice.fr2;
+      int ul_bw_mhz = mac->phy_config.config_req.carrier_config.uplink_bandwidth;
+      if (ul_bw_mhz != get_supported_bw_mhz(mac->frequency_range, uecap_bw_index))
+        continue;
+      if (fs_ul_cc->maxNumberMIMO_LayersNonCB_PUSCH)
+        mac->uecap_maxMIMO_PUSCH_layers_nocb = 1 << *fs_ul_cc->maxNumberMIMO_LayersNonCB_PUSCH;
+      if (fs_ul_cc->mimo_CB_PUSCH && fs_ul_cc->mimo_CB_PUSCH->maxNumberMIMO_LayersCB_PUSCH)
+        mac->uecap_maxMIMO_PUSCH_layers_cb = 1 << *fs_ul_cc->mimo_CB_PUSCH->maxNumberMIMO_LayersCB_PUSCH;
     }
   }
 }
diff --git a/openair2/LAYER2/NR_MAC_UE/mac_defs.h b/openair2/LAYER2/NR_MAC_UE/mac_defs.h
index 3be1fb8414..087e7ee6cd 100644
--- a/openair2/LAYER2/NR_MAC_UE/mac_defs.h
+++ b/openair2/LAYER2/NR_MAC_UE/mac_defs.h
@@ -475,7 +475,8 @@ typedef struct NR_UE_MAC_INST_s {
   bool harq_ACK_SpatialBundlingPUSCH;
 
   uint32_t uecap_maxMIMO_PDSCH_layers;
-  uint32_t uecap_maxMIMO_PUSCH_layers;
+  uint32_t uecap_maxMIMO_PUSCH_layers_cb;
+  uint32_t uecap_maxMIMO_PUSCH_layers_nocb;
 
   NR_UL_TIME_ALIGNMENT_t ul_time_alignment;
   NR_TDD_UL_DL_ConfigCommon_t *tdd_UL_DL_ConfigurationCommon;
diff --git a/openair2/LAYER2/NR_MAC_UE/main_ue_nr.c b/openair2/LAYER2/NR_MAC_UE/main_ue_nr.c
index 06b4310f60..973ecb9fa3 100644
--- a/openair2/LAYER2/NR_MAC_UE/main_ue_nr.c
+++ b/openair2/LAYER2/NR_MAC_UE/main_ue_nr.c
@@ -68,7 +68,8 @@ void nr_ue_init_mac(module_id_t module_idP)
   mac->harq_ACK_SpatialBundlingPUCCH = false;
   mac->harq_ACK_SpatialBundlingPUSCH = false;
   mac->uecap_maxMIMO_PDSCH_layers = 0;
-  mac->uecap_maxMIMO_PUSCH_layers = 0;
+  mac->uecap_maxMIMO_PUSCH_layers_cb = 0;
+  mac->uecap_maxMIMO_PUSCH_layers_nocb = 0;
 
   memset(&mac->ssb_measurements, 0, sizeof(mac->ssb_measurements));
   memset(&mac->ul_time_alignment, 0, sizeof(mac->ul_time_alignment));
diff --git a/openair2/LAYER2/NR_MAC_UE/nr_ue_scheduler.c b/openair2/LAYER2/NR_MAC_UE/nr_ue_scheduler.c
index 9a89ff6474..48bb4925c8 100644
--- a/openair2/LAYER2/NR_MAC_UE/nr_ue_scheduler.c
+++ b/openair2/LAYER2/NR_MAC_UE/nr_ue_scheduler.c
@@ -692,16 +692,28 @@ int nr_config_pusch_pdu(NR_UE_MAC_INST_t *mac,
       N_PRB_oh = 0;
 
     if (sc_info->rateMatching_PUSCH) {
-      long *maxMIMO_Layers = sc_info->maxMIMO_Layers_PUSCH;
-      if (!maxMIMO_Layers)
-        maxMIMO_Layers = pusch_Config ? pusch_Config->maxRank : NULL;
-      AssertFatal (maxMIMO_Layers != NULL,"Option with max MIMO layers not configured is not supported\n");
-      pusch_config_pdu->tbslbrm = nr_compute_tbslbrm(pusch_config_pdu->mcs_table, sc_info->ul_bw_tbslbrm, *maxMIMO_Layers);
+      long maxMIMO_Layers = 0;
+      if (sc_info->maxMIMO_Layers_PUSCH)
+        maxMIMO_Layers = *sc_info->maxMIMO_Layers_PUSCH;
+      else if (pusch_Config && pusch_Config->maxRank)
+        maxMIMO_Layers = *pusch_Config->maxRank;
+      else {
+        if (pusch_Config && pusch_Config->txConfig) {
+          if (*pusch_Config->txConfig == NR_PUSCH_Config__txConfig_codebook)
+            maxMIMO_Layers = mac->uecap_maxMIMO_PUSCH_layers_cb;
+          else
+            maxMIMO_Layers = mac->uecap_maxMIMO_PUSCH_layers_nocb;
+        } else
+          maxMIMO_Layers = 1; // single antenna port
+      }
+      AssertFatal (maxMIMO_Layers > 0, "Invalid number of max MIMO layers for PUSCH\n");
+      pusch_config_pdu->tbslbrm = nr_compute_tbslbrm(pusch_config_pdu->mcs_table, sc_info->ul_bw_tbslbrm, maxMIMO_Layers);
     } else
       pusch_config_pdu->tbslbrm = 0;
 
     /* PTRS */
-    if (pusch_Config && pusch_Config->dmrs_UplinkForPUSCH_MappingTypeB && pusch_Config->dmrs_UplinkForPUSCH_MappingTypeB->choice.setup->phaseTrackingRS) {
+    if (pusch_Config && pusch_Config->dmrs_UplinkForPUSCH_MappingTypeB &&
+        pusch_Config->dmrs_UplinkForPUSCH_MappingTypeB->choice.setup->phaseTrackingRS) {
       if (pusch_config_pdu->transform_precoding == NR_PUSCH_Config__transformPrecoder_disabled) {
         nfapi_nr_ue_ptrs_ports_t ptrs_ports_list;
         pusch_config_pdu->pusch_ptrs.ptrs_ports_list = &ptrs_ports_list;
@@ -718,7 +730,8 @@ int nr_config_pusch_pdu(NR_UE_MAC_INST_t *mac,
         if(valid_ptrs_setup == true) {
           pusch_config_pdu->pdu_bit_map |= PUSCH_PDU_BITMAP_PUSCH_PTRS;
         }
-        LOG_D(NR_MAC, "UL PTRS values: PTRS time den: %d, PTRS freq den: %d\n", pusch_config_pdu->pusch_ptrs.ptrs_time_density, pusch_config_pdu->pusch_ptrs.ptrs_freq_density);
+        LOG_D(NR_MAC, "UL PTRS values: PTRS time den: %d, PTRS freq den: %d\n",
+              pusch_config_pdu->pusch_ptrs.ptrs_time_density, pusch_config_pdu->pusch_ptrs.ptrs_freq_density);
       }
     }
   }