Custom CUDA kernels with AMP

username: Skylar_Wurster
cooked: <p>To speed up throughput of my NeRF-like model for 3D scientific data, I wrote some custom CUDA kernels for the encoding step (the decoding step is just a small MLP).<br>
This worked fine and dandy, improving my encoding throughput.<br>
Native PyTorch:112 mil points/sec, 8GB VRAM<br>
CUDA kernels: <strong>190 mil points/sec</strong>, 4GB VRAM.</p>
<p>Then I thought about AMP and how that could probably speed up both significantly. After re-writing my C++/kernels to use e.g. <code>packed_accessor32&lt;scalar_t,dims,torch::RestrictPtrTraits&gt;()</code> with templated functions and kernels, my throughput remained the same for <code>float32</code>, but doesn’t improve much for <code>float16</code> operations. PyTorch is faster again! Here are the stats using the same code as above, but in <code>with torch.autocast(device_type='cuda', enabled=True, dtype=torch.float16)</code>:<br>
Native PyTorch: <strong>236 mil points/sec</strong>, 7.16 GB<br>
CUDA kernels: 225 mil points/sec, 2.23 GB</p>
<p>As compared to float32, PyTorch is realizing almost a perfect 100% improvement, while my kernel is only improving by ~20%. I have verified that the <code>dtype</code>s of the data are <code>torch.float16</code>, casted properly with the CUDA functions returning that <code>dtype</code> as well.</p>
<p>I decorated the tested function as shown below:</p>
<pre><code class="lang-auto">import torch
from torch.amp import custom_fwd, custom_bwd
class EncodeCoordinates(torch.autograd.Function):
    @staticmethod
    @custom_fwd(device_type="cuda", cast_inputs=torch.float16)
    def forward(ctx, query_coordinates, rotations, scales, translations, feature_grids):
        
        feature_vectors = _C.encodeForward(query_coordinates, rotations, 
                                        scales, translations, feature_grids)

        ctx.save_for_backward(query_coordinates, rotations, 
                        scales, translations, feature_grids)
        return feature_vectors

    @staticmethod
    @custom_bwd(device_type="cuda")
    def backward(ctx, grad_output):
        query_coordinates, rotations, scales, \
            translations, feature_grids = ctx.saved_tensors

        grad_feature_grids = _C.encodeBackward(query_coordinates, 
            rotations, scales, translations, feature_grids, grad_output)
        return None, None, None, None, grad_feature_grids
</code></pre>
<p>The C++ function bound with pybind looks like this:</p>
<pre><code class="lang-auto">torch::Tensor encode_forward(
    const torch::Tensor&amp; query_points,
    const torch::Tensor&amp; rotations, 
    const torch::Tensor&amp; scales, 
    const torch::Tensor&amp; translations,
    const torch::Tensor&amp; feature_grids)
{
    const auto num_points = query_points.size(0);
    const auto num_grids = feature_grids.size(0);
    const auto features_per_grid = feature_grids.size(1);
    const auto D = feature_grids.size(2);
    const auto H = feature_grids.size(3);
    const auto W = feature_grids.size(4);    

    auto options = torch::TensorOptions().dtype(query_points.dtype()).device(query_points.device());

    torch::Tensor out_features = torch::empty({num_points, num_grids*features_per_grid}, options);

    AT_DISPATCH_FLOATING_TYPES_AND_HALF(query_points.scalar_type(), "launch_encode_forward", ([&amp;] {
        launch_encode_forward&lt;scalar_t&gt;(
            query_points, 
            rotations,
            scales, 
            translations, 
            feature_grids, 
            out_features
        );
    }));

    return out_features;
}
</code></pre>
<p>And the relevant CUDA code looks like this</p>
<pre><code class="lang-auto">template &lt;typename scalar_t&gt;
struct scalar_t3 {
    scalar_t x;
    scalar_t y;
    scalar_t z;

    __device__ __host__ scalar_t3() : x(0), y(0), z(0) {}
    __device__ __host__ scalar_t3(scalar_t x_, scalar_t y_, scalar_t z_) : x(x_), y(y_), z(z_) {}
};
template &lt;typename scalar_t&gt;
__host__ __device__ __forceinline__ scalar_t3&lt;scalar_t&gt; make_scalar_t3(scalar_t x, scalar_t y, scalar_t z) {
    return scalar_t3&lt;scalar_t&gt;(x, y, z);
}
template &lt;typename scalar_t&gt;
void launch_encode_forward(
    const torch::Tensor&amp; query_points,
    const torch::Tensor&amp; rotations,
    const torch::Tensor&amp; scales,
    const torch::Tensor&amp; translations,
    const torch::Tensor&amp; feature_grids,
    torch::Tensor&amp; output_features)
{
    const int num_points = query_points.size(0);
    const int num_grids = rotations.size(0);

    // Allocate memory for rotation matrices
    scalar_t* rotation_matrices;
    cudaMalloc(&amp;rotation_matrices, num_grids * 3 * 3 * sizeof(scalar_t));

    auto blocksPerGrid = (num_grids + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
    quaternionScaleToRotationMatrix&lt;&lt;&lt;blocksPerGrid, THREADS_PER_BLOCK&gt;&gt;&gt;(
        rotations.packed_accessor32&lt;scalar_t,2,torch::RestrictPtrTraits&gt;(),
        scales.packed_accessor32&lt;scalar_t,2,torch::RestrictPtrTraits&gt;(),
        rotation_matrices
    );

    dim3 threadsPerBlock(THREADS_PER_BLOCK, 1);
    dim3 numBlocks((num_points + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK, num_grids);
    encodeForwardKernel&lt;&lt;&lt;numBlocks, threadsPerBlock&gt;&gt;&gt;(
        query_points.packed_accessor32&lt;scalar_t,2,torch::RestrictPtrTraits&gt;(),
        rotation_matrices,
        translations.packed_accessor32&lt;scalar_t,2,torch::RestrictPtrTraits&gt;(),
        feature_grids.packed_accessor32&lt;scalar_t,5,torch::RestrictPtrTraits&gt;(),
        output_features.packed_accessor32&lt;scalar_t,2,torch::RestrictPtrTraits&gt;()
    );

    // Free the allocated memory
    cudaFree(rotation_matrices);
}
template &lt;typename scalar_t&gt;
__global__ void encodeForwardKernel(
    const at::PackedTensorAccessor32&lt;scalar_t,2,torch::RestrictPtrTraits&gt; query_points,
    const scalar_t* rotation_matrices,
    const at::PackedTensorAccessor32&lt;scalar_t,2,torch::RestrictPtrTraits&gt; translations,
    const at::PackedTensorAccessor32&lt;scalar_t,5,torch::RestrictPtrTraits&gt; feature_grids,
    at::PackedTensorAccessor32&lt;scalar_t,2,torch::RestrictPtrTraits&gt; output_features) {

    const auto point_idx = blockIdx.x * blockDim.x + threadIdx.x;
    const auto grid_idx = blockIdx.y;

    if (grid_idx &gt;= feature_grids.size(0) || point_idx &gt;= query_points.size(0)) return;

    scalar_t3&lt;scalar_t&gt; point = make_scalar_t3&lt;scalar_t&gt;(query_points[point_idx][0], query_points[point_idx][1], query_points[point_idx][2]);

    scalar_t3&lt;scalar_t&gt; point_t = transformPoint&lt;scalar_t&gt;(grid_idx, rotation_matrices, translations, point);
    
    trilinearInterpolate&lt;scalar_t&gt;(
        grid_idx,
        point_idx,
        feature_grids,
        point_t,
        output_features
    );
    
}
</code></pre>
<p>I also made sure the various precisions are being compiled in the CUDA code with</p>
<pre><code class="lang-auto">template void launch_encode_forward&lt;float&gt;(const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, at::Tensor&amp;);
template void launch_encode_forward&lt;double&gt;(const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, at::Tensor&amp;);
template void launch_encode_forward&lt;c10::Half&gt;(const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, const at::Tensor&amp;, at::Tensor&amp;);
</code></pre>
<p>added at the end of mu <code>.cu</code> file with all the CUDA code.</p>
<p>All of my code is <a href="https://github.com/skywolf829/AMGSRN/tree/main/AMGSRN/Models/AMGSRN/src" rel="noopener nofollow ugc">here</a>, not included in this post for brevity (and because I suspect it might just come down to writing better kernels for fp16).</p>
<p>I had trouble finding documentation covering AMP with custom CUDA operations, and most of my support was with Claude through Cursor, so I’m wondering if there is something I’m missing to make sure the AMP code is working properly. If everything looks correct here, are there tips/resources/examples of CUDA kernels written to work with AMP I could adjust my code to follow?</p>