You don't need a heatsink for most NVMe drives and even then it's best to make sure it only makes contact with the controller. You can usually do that by removing part of the thermal pad so it only makes contact with the controller chip.
The controllers themselves can handle fairly high temperatures before they start throttling, even if they get to 70C~80C. Even when they throttle and get to >90C, it should still be fine since that's within spec for most of them.
As for the NAND, you shouldn't put a heatsink on that because it decreases its long term reliability, NAND works best when it's hot.
Also if you think putting a heatsink on your drive will help with write speeds, make sure your writes are decreasing because of thermal throttling first. Write speeds will drop considerable on most drives once they exhaust their pSLC cache.