struct ath10k_pci *ar_pci = ath10k_pci_priv(ar);
        int ret = 0;
        u32 *buf;
-       unsigned int completed_nbytes, orig_nbytes, remaining_bytes;
+       unsigned int completed_nbytes, alloc_nbytes, remaining_bytes;
        struct ath10k_ce_pipe *ce_diag;
        /* Host buffer address in CE space */
        u32 ce_data;
         *   1) 4-byte alignment
         *   2) Buffer in DMA-able space
         */
-       orig_nbytes = nbytes;
+       alloc_nbytes = min_t(unsigned int, nbytes, DIAG_TRANSFER_LIMIT);
+
        data_buf = (unsigned char *)dma_alloc_coherent(ar->dev,
-                                                      orig_nbytes,
+                                                      alloc_nbytes,
                                                       &ce_data_base,
                                                       GFP_ATOMIC);
 
                ret = -ENOMEM;
                goto done;
        }
-       memset(data_buf, 0, orig_nbytes);
+       memset(data_buf, 0, alloc_nbytes);
 
-       remaining_bytes = orig_nbytes;
+       remaining_bytes = nbytes;
        ce_data = ce_data_base;
        while (remaining_bytes) {
                nbytes = min_t(unsigned int, remaining_bytes,
                }
 
                remaining_bytes -= nbytes;
+
+               if (ret) {
+                       ath10k_warn(ar, "failed to read diag value at 0x%x: %d\n",
+                                   address, ret);
+                       break;
+               }
+               memcpy(data, data_buf, nbytes);
+
                address += nbytes;
-               ce_data += nbytes;
+               data += nbytes;
        }
 
 done:
-       if (ret == 0)
-               memcpy(data, data_buf, orig_nbytes);
-       else
-               ath10k_warn(ar, "failed to read diag value at 0x%x: %d\n",
-                           address, ret);
 
        if (data_buf)
-               dma_free_coherent(ar->dev, orig_nbytes, data_buf,
+               dma_free_coherent(ar->dev, alloc_nbytes, data_buf,
                                  ce_data_base);
 
        spin_unlock_bh(&ar_pci->ce_lock);