&req, sizeof(req), true);
 }
 
+static struct sk_buff *
+mt7615_mcu_alloc_sta_req(struct mt7615_vif *mvif, struct mt7615_sta *msta)
+{
+       struct sta_req_hdr hdr = {
+               .bss_idx = mvif->idx,
+               .wlan_idx = msta->wcid.idx,
+               .muar_idx = mvif->omac_idx,
+               .is_tlv_append = 1,
+       };
+       struct sk_buff *skb;
+
+       skb = mt7615_mcu_msg_alloc(NULL, MT7615_STA_UPDATE_MAX_SIZE);
+       if (!skb)
+               return ERR_PTR(-ENOMEM);
+
+       skb_put_data(skb, &hdr, sizeof(hdr));
+
+       return skb;
+}
+
+static struct wtbl_req_hdr *
+mt7615_mcu_alloc_wtbl_req(struct mt7615_sta *msta, int cmd,
+                         void *sta_wtbl, struct sk_buff **skb)
+{
+       struct sta_rec_wtbl *sta_hdr = sta_wtbl;
+       struct wtbl_req_hdr hdr = {
+               .wlan_idx = msta->wcid.idx,
+               .operation = cmd,
+       };
+       struct sk_buff *nskb = *skb;
+
+       if (!nskb) {
+               nskb = mt7615_mcu_msg_alloc(NULL, MT7615_WTBL_UPDATE_BA_SIZE);
+               if (!nskb)
+                       return ERR_PTR(-ENOMEM);
+
+               *skb = nskb;
+       }
+
+       if (sta_hdr)
+               sta_hdr->len = cpu_to_le16(sizeof(hdr));
+
+       return skb_put_data(nskb, &hdr, sizeof(hdr));
+}
+
+static struct tlv *
+mt7615_mcu_add_nested_tlv(struct sk_buff *skb, int tag, int len,
+                         void *sta_ntlv, void *sta_wtbl)
+{
+       struct sta_ntlv_hdr *ntlv_hdr = sta_ntlv;
+       struct sta_rec_wtbl *sta_hdr = sta_wtbl;
+       struct tlv *ptlv, tlv = {
+               .tag = cpu_to_le16(tag),
+               .len = cpu_to_le16(len),
+       };
+       u16 ntlv;
+
+       ptlv = skb_put(skb, len);
+       memcpy(ptlv, &tlv, sizeof(tlv));
+
+       ntlv = le16_to_cpu(ntlv_hdr->tlv_num);
+       ntlv_hdr->tlv_num = cpu_to_le16(ntlv + 1);
+
+       if (sta_hdr) {
+               u16 size = le16_to_cpu(sta_hdr->len);
+
+               sta_hdr->len = cpu_to_le16(size + len);
+       }
+
+       return ptlv;
+}
+
+static struct tlv *
+mt7615_mcu_add_tlv(struct sk_buff *skb, int tag, int len)
+{
+       return mt7615_mcu_add_nested_tlv(skb, tag, len, skb->data, NULL);
+}
+
+static void
+mt7615_mcu_sta_ba_tlv(struct sk_buff *skb,
+                     struct ieee80211_ampdu_params *params,
+                     bool enable, bool tx)
+{
+       struct sta_rec_ba *ba;
+       struct tlv *tlv;
+
+       tlv = mt7615_mcu_add_tlv(skb, STA_REC_BA, sizeof(*ba));
+
+       ba = (struct sta_rec_ba *)tlv;
+       ba->ba_type = tx ? MT_BA_TYPE_ORIGINATOR : MT_BA_TYPE_RECIPIENT,
+       ba->winsize = cpu_to_le16(params->buf_size);
+       ba->ssn = cpu_to_le16(params->ssn);
+       ba->ba_en = enable << params->tid;
+       ba->amsdu = params->amsdu;
+       ba->tid = params->tid;
+}
+
+static void
+mt7615_mcu_wtbl_ba_tlv(struct sk_buff *skb,
+                      struct ieee80211_ampdu_params *params,
+                      bool enable, bool tx, void *sta_wtbl,
+                      void *wtbl_tlv)
+{
+       struct wtbl_ba *ba;
+       struct tlv *tlv;
+
+       tlv = mt7615_mcu_add_nested_tlv(skb, WTBL_BA, sizeof(*ba),
+                                       wtbl_tlv, sta_wtbl);
+
+       ba = (struct wtbl_ba *)tlv;
+       ba->tid = params->tid;
+
+       if (tx) {
+               ba->ba_type = MT_BA_TYPE_ORIGINATOR;
+               ba->sn = enable ? cpu_to_le16(params->ssn) : 0;
+               ba->ba_winsize = cpu_to_le16(params->buf_size);
+               ba->ba_en = enable;
+       } else {
+               memcpy(ba->peer_addr, params->sta->addr, ETH_ALEN);
+               ba->ba_type = MT_BA_TYPE_RECIPIENT;
+               ba->rst_ba_tid = params->tid;
+               ba->rst_ba_sel = RST_BA_MAC_TID_MATCH;
+               ba->rst_ba_sb = 1;
+       }
+
+       if (enable && tx) {
+               u8 ba_range[] = { 4, 8, 12, 24, 36, 48, 54, 64 };
+               int i;
+
+               for (i = 7; i > 0; i--) {
+                       if (params->buf_size >= ba_range[i])
+                               break;
+               }
+               ba->ba_winsize_idx = i;
+       }
+}
+
+static int
+mt7615_mcu_wtbl_tx_ba(struct mt7615_dev *dev,
+                     struct ieee80211_ampdu_params *params,
+                     bool enable)
+{
+       struct mt7615_sta *msta = (struct mt7615_sta *)params->sta->drv_priv;
+       struct mt7615_vif *mvif = msta->vif;
+       struct wtbl_req_hdr *wtbl_hdr;
+       struct sk_buff *skb = NULL;
+       int err;
+
+       wtbl_hdr = mt7615_mcu_alloc_wtbl_req(msta, WTBL_SET, NULL, &skb);
+       if (IS_ERR(wtbl_hdr))
+               return PTR_ERR(wtbl_hdr);
+
+       mt7615_mcu_wtbl_ba_tlv(skb, params, enable, true, NULL, wtbl_hdr);
+
+       err = __mt76_mcu_skb_send_msg(&dev->mt76, skb,
+                                     MCU_EXT_CMD_WTBL_UPDATE, true);
+       if (err < 0)
+               return err;
+
+       skb = mt7615_mcu_alloc_sta_req(mvif, msta);
+       if (IS_ERR(skb))
+               return PTR_ERR(skb);
+
+       mt7615_mcu_sta_ba_tlv(skb, params, enable, true);
+
+       return __mt76_mcu_skb_send_msg(&dev->mt76, skb,
+                                      MCU_EXT_CMD_STA_REC_UPDATE, true);
+}
+
+static int
+mt7615_mcu_wtbl_rx_ba(struct mt7615_dev *dev,
+                     struct ieee80211_ampdu_params *params,
+                     bool enable)
+{
+       struct mt7615_sta *msta = (struct mt7615_sta *)params->sta->drv_priv;
+       struct mt7615_vif *mvif = msta->vif;
+       struct wtbl_req_hdr *wtbl_hdr;
+       struct sk_buff *skb;
+       int err;
+
+       skb = mt7615_mcu_alloc_sta_req(mvif, msta);
+       if (IS_ERR(skb))
+               return PTR_ERR(skb);
+
+       mt7615_mcu_sta_ba_tlv(skb, params, enable, false);
+
+       err = __mt76_mcu_skb_send_msg(&dev->mt76, skb,
+                                     MCU_EXT_CMD_STA_REC_UPDATE, true);
+       if (err < 0 || !enable)
+               return err;
+
+       skb = NULL;
+       wtbl_hdr = mt7615_mcu_alloc_wtbl_req(msta, WTBL_SET, NULL, &skb);
+       if (IS_ERR(wtbl_hdr))
+               return PTR_ERR(wtbl_hdr);
+
+       mt7615_mcu_wtbl_ba_tlv(skb, params, enable, false, NULL, wtbl_hdr);
+
+       return __mt76_mcu_skb_send_msg(&dev->mt76, skb,
+                                      MCU_EXT_CMD_WTBL_UPDATE, true);
+}
+
+static const struct mt7615_mcu_ops wtbl_update_ops = {
+       .add_tx_ba = mt7615_mcu_wtbl_tx_ba,
+       .add_rx_ba = mt7615_mcu_wtbl_rx_ba,
+};
+
+static int
+mt7615_mcu_sta_ba(struct mt7615_dev *dev,
+                 struct ieee80211_ampdu_params *params,
+                 bool enable, bool tx)
+{
+       struct mt7615_sta *msta = (struct mt7615_sta *)params->sta->drv_priv;
+       struct mt7615_vif *mvif = msta->vif;
+       struct wtbl_req_hdr *wtbl_hdr;
+       struct tlv *sta_wtbl;
+       struct sk_buff *skb;
+
+       skb = mt7615_mcu_alloc_sta_req(mvif, msta);
+       if (IS_ERR(skb))
+               return PTR_ERR(skb);
+
+       mt7615_mcu_sta_ba_tlv(skb, params, enable, tx);
+       sta_wtbl = mt7615_mcu_add_tlv(skb, STA_REC_WTBL,
+                                     sizeof(struct sta_rec_wtbl));
+
+       wtbl_hdr = mt7615_mcu_alloc_wtbl_req(msta, WTBL_SET, sta_wtbl, &skb);
+       mt7615_mcu_wtbl_ba_tlv(skb, params, enable, tx, sta_wtbl, wtbl_hdr);
+
+       return __mt76_mcu_skb_send_msg(&dev->mt76, skb,
+                                      MCU_EXT_CMD_STA_REC_UPDATE, true);
+}
+
+static int
+mt7615_mcu_sta_tx_ba(struct mt7615_dev *dev,
+                    struct ieee80211_ampdu_params *params,
+                    bool enable)
+{
+       return mt7615_mcu_sta_ba(dev, params, enable, true);
+}
+
+static int
+mt7615_mcu_sta_rx_ba(struct mt7615_dev *dev,
+                    struct ieee80211_ampdu_params *params,
+                    bool enable)
+{
+       return mt7615_mcu_sta_ba(dev, params, enable, false);
+}
+
+static const struct mt7615_mcu_ops sta_update_ops = {
+       .add_tx_ba = mt7615_mcu_sta_tx_ba,
+       .add_rx_ba = mt7615_mcu_sta_rx_ba,
+};
+
 static int mt7615_mcu_send_firmware(struct mt7615_dev *dev, const void *data,
                                    int len)
 {
                 sizeof(dev->mt76.hw->wiphy->fw_version),
                 "%.10s-%.15s", hdr->fw_ver, hdr->build_date);
 
-       if (!strncmp(hdr->fw_ver, "2.0", sizeof(hdr->fw_ver)))
+       if (!strncmp(hdr->fw_ver, "2.0", sizeof(hdr->fw_ver))) {
                dev->fw_ver = MT7615_FIRMWARE_V2;
-       else
+               dev->mcu_ops = &sta_update_ops;
+       } else {
                dev->fw_ver = MT7615_FIRMWARE_V1;
+               dev->mcu_ops = &wtbl_update_ops;
+       }
 
 out:
        release_firmware(fw);
        return __mt76_mcu_send_msg(&dev->mt76, cmd, &req, sizeof(req), true);
 }
 
-int mt7615_mcu_set_tx_ba(struct mt7615_dev *dev,
-                        struct ieee80211_ampdu_params *params,
-                        bool add)
-{
-       struct mt7615_sta *msta = (struct mt7615_sta *)params->sta->drv_priv;
-       struct mt7615_vif *mvif = msta->vif;
-       struct {
-               struct sta_req_hdr hdr;
-               struct sta_rec_ba ba;
-               u8 buf[MT7615_WTBL_UPDATE_MAX_SIZE];
-       } __packed req = {
-               .hdr = {
-                       .bss_idx = mvif->idx,
-                       .wlan_idx = msta->wcid.idx,
-                       .tlv_num = cpu_to_le16(1),
-                       .is_tlv_append = 1,
-                       .muar_idx = mvif->omac_idx,
-               },
-               .ba = {
-                       .tag = cpu_to_le16(STA_REC_BA),
-                       .len = cpu_to_le16(sizeof(struct sta_rec_ba)),
-                       .tid = params->tid,
-                       .ba_type = MT_BA_TYPE_ORIGINATOR,
-                       .amsdu = params->amsdu,
-                       .ba_en = add << params->tid,
-                       .ssn = cpu_to_le16(params->ssn),
-                       .winsize = cpu_to_le16(params->buf_size),
-               },
-       };
-       struct sta_rec_wtbl *wtbl = NULL;
-       struct wtbl_req_hdr *wtbl_hdr;
-       struct wtbl_ba *wtbl_ba;
-       u8 *buf = req.buf;
-
-       if (dev->fw_ver > MT7615_FIRMWARE_V1) {
-               req.hdr.tlv_num = cpu_to_le16(2);
-               wtbl = (struct sta_rec_wtbl *)buf;
-               wtbl->tag = cpu_to_le16(STA_REC_WTBL);
-               buf += sizeof(*wtbl);
-       }
-
-       wtbl_hdr = (struct wtbl_req_hdr *)buf;
-       buf += sizeof(*wtbl_hdr);
-       wtbl_hdr->wlan_idx = msta->wcid.idx;
-       wtbl_hdr->operation = WTBL_SET;
-       wtbl_hdr->tlv_num = cpu_to_le16(1);
-
-       wtbl_ba = (struct wtbl_ba *)buf;
-       buf += sizeof(*wtbl_ba);
-       wtbl_ba->tag = cpu_to_le16(WTBL_BA);
-       wtbl_ba->len = cpu_to_le16(sizeof(*wtbl_ba));
-       wtbl_ba->tid = params->tid;
-       wtbl_ba->ba_type = MT_BA_TYPE_ORIGINATOR;
-       wtbl_ba->sn = add ? cpu_to_le16(params->ssn) : 0;
-       wtbl_ba->ba_en = add;
-
-       if (add) {
-               u8 idx, ba_range[] = { 4, 8, 12, 24, 36, 48, 54, 64 };
-
-               for (idx = 7; idx > 0; idx--) {
-                       if (params->buf_size >= ba_range[idx])
-                               break;
-               }
-
-               wtbl_ba->ba_winsize_idx = idx;
-       }
-
-       if (wtbl)
-               wtbl->len = cpu_to_le16(buf - (u8 *)wtbl_hdr);
-
-       return mt7615_mcu_send_sta_rec(dev, (u8 *)&req, (u8 *)wtbl_hdr,
-                                      buf - (u8 *)wtbl_hdr, true);
-}
-
-int mt7615_mcu_set_rx_ba(struct mt7615_dev *dev,
-                        struct ieee80211_ampdu_params *params,
-                        bool add)
-{
-       struct mt7615_sta *msta = (struct mt7615_sta *)params->sta->drv_priv;
-       struct mt7615_vif *mvif = msta->vif;
-       struct {
-               struct sta_req_hdr hdr;
-               struct sta_rec_ba ba;
-               u8 buf[MT7615_WTBL_UPDATE_MAX_SIZE];
-       } __packed req = {
-               .hdr = {
-                       .bss_idx = mvif->idx,
-                       .wlan_idx = msta->wcid.idx,
-                       .tlv_num = cpu_to_le16(1),
-                       .is_tlv_append = 1,
-                       .muar_idx = mvif->omac_idx,
-               },
-               .ba = {
-                       .tag = cpu_to_le16(STA_REC_BA),
-                       .len = cpu_to_le16(sizeof(struct sta_rec_ba)),
-                       .tid = params->tid,
-                       .ba_type = MT_BA_TYPE_RECIPIENT,
-                       .amsdu = params->amsdu,
-                       .ba_en = add << params->tid,
-                       .ssn = cpu_to_le16(params->ssn),
-                       .winsize = cpu_to_le16(params->buf_size),
-               },
-       };
-       struct sta_rec_wtbl *wtbl = NULL;
-       struct wtbl_req_hdr *wtbl_hdr;
-       struct wtbl_ba *wtbl_ba;
-       u8 *buf = req.buf;
-
-       if (dev->fw_ver > MT7615_FIRMWARE_V1) {
-               req.hdr.tlv_num = cpu_to_le16(2);
-               wtbl = (struct sta_rec_wtbl *)buf;
-               wtbl->tag = cpu_to_le16(STA_REC_WTBL);
-               buf += sizeof(*wtbl);
-       }
-
-       wtbl_hdr = (struct wtbl_req_hdr *)buf;
-       buf += sizeof(*wtbl_hdr);
-       wtbl_hdr->wlan_idx = msta->wcid.idx;
-       wtbl_hdr->operation = WTBL_SET;
-       wtbl_hdr->tlv_num = cpu_to_le16(1);
-
-       wtbl_ba = (struct wtbl_ba *)buf;
-       buf += sizeof(*wtbl_ba);
-       wtbl_ba->tag = cpu_to_le16(WTBL_BA);
-       wtbl_ba->len = cpu_to_le16(sizeof(*wtbl_ba));
-       wtbl_ba->tid = params->tid;
-       wtbl_ba->ba_type = MT_BA_TYPE_RECIPIENT;
-       wtbl_ba->rst_ba_tid = params->tid;
-       wtbl_ba->rst_ba_sel = RST_BA_MAC_TID_MATCH;
-       wtbl_ba->rst_ba_sb = 1;
-
-       memcpy(wtbl_ba->peer_addr, params->sta->addr, ETH_ALEN);
-
-       if (wtbl)
-               wtbl->len = cpu_to_le16(buf - (u8 *)wtbl_hdr);
-
-       return mt7615_mcu_send_sta_rec(dev, (u8 *)&req, (u8 *)wtbl_hdr,
-                                      buf - (u8 *)wtbl_hdr, add);
-}
-
 int mt7615_mcu_get_temperature(struct mt7615_dev *dev, int index)
 {
        struct {