If you only use basic ops (as opposed to custom kernels), you could see whether you can use automatic differentiation (edit: I tried, works out of the box).
In that case, however, it might not be necessary to use C++ extension at all unless you do really funny stuff, the Python speed has been rather decent for me. (I tried that with torch.nn.functional.bilinear once.)
Now, all that is theorizing, I don’t know what is best for your use case.
Best regards
Thomas