#define IWL_TX_CMD_OFFLD_MH_MASK       0x1f
 #define IWL_TX_CMD_OFFLD_IP_HDR_MASK   0x3f
 
+enum iwl_tx_offload_assist_bz {
+       IWL_TX_CMD_OFFLD_BZ_RESULT_OFFS         = 0x000003ff,
+       IWL_TX_CMD_OFFLD_BZ_START_OFFS          = 0x001ff800,
+       IWL_TX_CMD_OFFLD_BZ_MH_LEN              = 0x07c00000,
+       IWL_TX_CMD_OFFLD_BZ_MH_PAD              = 0x08000000,
+       IWL_TX_CMD_OFFLD_BZ_AMSDU               = 0x10000000,
+       IWL_TX_CMD_OFFLD_BZ_ZERO2ONES           = 0x20000000,
+       IWL_TX_CMD_OFFLD_BZ_ENABLE_CSUM         = 0x40000000,
+       IWL_TX_CMD_OFFLD_BZ_PARTIAL_CSUM        = 0x80000000,
+};
+
 /* TODO: complete documentation for try_cnt and btkill_cnt */
 /**
  * struct iwl_tx_cmd - TX command struct to FW
 
 #define OPT_HDR(type, skb, off) \
        (type *)(skb_network_header(skb) + (off))
 
-static u16 iwl_mvm_tx_csum(struct iwl_mvm *mvm, struct sk_buff *skb,
-                          struct ieee80211_hdr *hdr,
-                          struct ieee80211_tx_info *info,
-                          bool amsdu)
+static u16 iwl_mvm_tx_csum_pre_bz(struct iwl_mvm *mvm, struct sk_buff *skb,
+                                 struct ieee80211_tx_info *info, bool amsdu)
 {
+       struct ieee80211_hdr *hdr = (void *)skb->data;
        u16 offload_assist = 0;
 #if IS_ENABLED(CONFIG_INET)
        u16 mh_len = ieee80211_hdrlen(hdr->frame_control);
        return offload_assist;
 }
 
+u32 iwl_mvm_tx_csum_bz(struct iwl_mvm *mvm, struct sk_buff *skb, bool amsdu)
+{
+       struct ieee80211_hdr *hdr = (void *)skb->data;
+       u32 offload_assist = IWL_TX_CMD_OFFLD_BZ_PARTIAL_CSUM;
+       unsigned int hdrlen = ieee80211_hdrlen(hdr->frame_control);
+       unsigned int csum_start = skb_checksum_start_offset(skb);
+
+       offload_assist |= u32_encode_bits(hdrlen / 2,
+                                         IWL_TX_CMD_OFFLD_BZ_MH_LEN);
+       if (amsdu)
+               offload_assist |= IWL_TX_CMD_OFFLD_BZ_AMSDU;
+       else if (hdrlen % 4)
+               /* padding is inserted later in transport */
+               offload_assist |= IWL_TX_CMD_OFFLD_BZ_MH_PAD;
+
+       if (skb->ip_summed != CHECKSUM_PARTIAL)
+               return offload_assist;
+
+       offload_assist |= IWL_TX_CMD_OFFLD_BZ_ENABLE_CSUM |
+                         IWL_TX_CMD_OFFLD_BZ_ZERO2ONES;
+
+       /*
+        * mac80211 will always calculate checksum in software for
+        * non-fast-xmit, and so we can only do offloaded checksum
+        * for fast-xmit frames. In this case, we always have the
+        * RFC 1042 header present. skb_checksum_start_offset()
+        * returns the offset from the beginning, but the hardware
+        * needs it from after the header & SNAP header.
+        */
+       csum_start -= hdrlen + 8;
+
+       offload_assist |= u32_encode_bits(csum_start,
+                                         IWL_TX_CMD_OFFLD_BZ_START_OFFS);
+       offload_assist |= u32_encode_bits(csum_start + skb->csum_offset,
+                                         IWL_TX_CMD_OFFLD_BZ_RESULT_OFFS);
+
+       return offload_assist;
+}
+
+static u32 iwl_mvm_tx_csum(struct iwl_mvm *mvm, struct sk_buff *skb,
+                          struct ieee80211_tx_info *info,
+                          bool amsdu)
+{
+       if (mvm->trans->trans_cfg->device_family < IWL_DEVICE_FAMILY_BZ)
+               return iwl_mvm_tx_csum_pre_bz(mvm, skb, info, amsdu);
+       return iwl_mvm_tx_csum_bz(mvm, skb, amsdu);
+}
+
 /*
  * Sets most of the Tx cmd's fields
  */
        tx_cmd->sta_id = sta_id;
 
        tx_cmd->offload_assist =
-               cpu_to_le16(iwl_mvm_tx_csum(mvm, skb, hdr, info, amsdu));
+               cpu_to_le16(iwl_mvm_tx_csum_pre_bz(mvm, skb, info, amsdu));
 }
 
 static u32 iwl_mvm_get_tx_ant(struct iwl_mvm *mvm,
        dev_cmd->hdr.cmd = TX_CMD;
 
        if (iwl_mvm_has_new_tx_api(mvm)) {
-               u16 offload_assist;
                u32 rate_n_flags = 0;
                u16 flags = 0;
                struct iwl_mvm_sta *mvmsta = sta ?
                        amsdu = *qc & IEEE80211_QOS_CTL_A_MSDU_PRESENT;
                }
 
-               offload_assist = iwl_mvm_tx_csum(mvm, skb, hdr, info, amsdu);
-
                if (!info->control.hw_key)
                        flags |= IWL_TX_FLAGS_ENCRYPT_DIS;
 
                if (mvm->trans->trans_cfg->device_family >=
                    IWL_DEVICE_FAMILY_AX210) {
                        struct iwl_tx_cmd_gen3 *cmd = (void *)dev_cmd->payload;
+                       u32 offload_assist = iwl_mvm_tx_csum(mvm, skb,
+                                                            info, amsdu);
 
                        cmd->offload_assist = cpu_to_le32(offload_assist);
 
                        cmd->rate_n_flags = cpu_to_le32(rate_n_flags);
                } else {
                        struct iwl_tx_cmd_gen2 *cmd = (void *)dev_cmd->payload;
+                       u16 offload_assist = iwl_mvm_tx_csum_pre_bz(mvm, skb,
+                                                                   info,
+                                                                   amsdu);
 
                        cmd->offload_assist = cpu_to_le16(offload_assist);