Skip to content

fix(pyutils): raise on cudaFuncSetAttribute and launch errors in bind_kernel - #215

Open
TanishGudise wants to merge 1 commit into
HazyResearch:mainfrom
TanishGudise:c5-bind-kernel-checks
Open

TanishGudise wants to merge 1 commit into
HazyResearch:mainfrom
TanishGudise:c5-bind-kernel-checks

Conversation

@TanishGudise

Copy link
Copy Markdown

kittens::py::bind_kernel ignores the result of cudaFuncSetAttribute and never checks the launch. When a launch can't run (dynamic shared memory above the device's opt-in limit, an invalid block size, or a module with no code for the current GPU), the Python call returns normally and the output buffers are left untouched.

Repro

A minimal binding: one gl<float, 1, 1, -1, -1> output, which the Python side pre-fills with NaN, with dynamic shared memory and threads per block passed in.

// c5_ext.cu
#include "kittens.cuh"
#include "pyutils/pyutils.cuh"
using namespace kittens;

using out_gl = gl<float, 1, 1, -1, -1>;
struct c5_globals {
    out_gl out;
    int smem;      // requested dynamic shared memory (bytes)
    int threads;   // threads per block
    dim3 grid() const { return dim3(1); }
    dim3 block() const { return dim3(threads); }
    int dynamic_shared_memory() const { return smem; }
};

__global__ void c5_fill(const __grid_constant__ c5_globals g) {
    int i = threadIdx.x;
    if (i < (int)g.out.cols()) g.out.raw_ptr[i] = 1.0f;
}

PYBIND11_MODULE(c5, m) {
    kittens::py::bind_kernel<c5_fill>(m, "fill", &c5_globals::out, &c5_globals::smem, &c5_globals::threads);
}

TK is a ThunderKittens checkout:

nvcc -std=c++20 -O2 -shared -Xcompiler -fPIC --expt-extended-lambda --expt-relaxed-constexpr -DKITTENS_SM90 \
  -gencode arch=compute_90a,code=sm_90a -I$TK/include -I$TK/prototype \
  $(python3 -m pybind11 --includes) c5_ext.cu -o c5$(python3-config --extension-suffix) -lcuda
import torch, c5
out = torch.full((32,), float("nan"), device="cuda")
c5.fill(out, 233472, 32)   # 1 KB over the H100's opt-in shared memory limit
torch.cuda.synchronize()
print(out[:4])

On main, fill and torch.cuda.synchronize() both return and out is still NaN. With this branch, fill raises:

RuntimeError: fill: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=233472) failed: invalid argument

Results

H100 80GB HBM3, driver 595.91.07, nvcc 13.0, built as above (nvcc's default static cudart), 20 calls per case:

case main this branch
1 KB smem, 32 threads output written output written
smem = opt-in limit + 1 KB (233472 B) returns, output still NaN RuntimeError: fill: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=233472) failed: invalid argument
2048 threads per block returns, output still NaN RuntimeError: fill: kernel launch failed: invalid argument
1 KB smem, 32 threads, after the two failures output written output written
each case above, module built for sm_100a (no code for this GPU)* returns, output still NaN RuntimeError: fill: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=<smem>) failed: no kernel image is available for execution on the device, before the launch

* Built with nvcc 12.8.93 and -gencode arch=compute_100a,code=sm_100a, otherwise as above.

Megakernels' low-latency Llama-1B demo behaves the same with and without this branch (20 runs each): no run with all-"!" output, one megakernel launch per step, median step 0.949 ms (main) vs 0.948 ms (this branch). Setup: head of HazyResearch/Megakernels#14, which pins this repo at 67845f5, built with nvcc 13.0 and -gencode arch=compute_90a,code=sm_90a. The Makefile's -arch=sm_90a doesn't build on CUDA 13; see HazyResearch/Megakernels#10. Output varies from run to run in both builds (5 and 7 distinct outputs in 20 runs).

I ran these on 2026-10-23 against 67845f5, and main hasn't moved since, so I reused those results rather than re-running them. On 2026-10-26 I checked that the repro still compiles against main and this branch with nvcc 12.8.93 and 13.4.92 for sm_90a (compile only, no warnings).

Change

bind_kernel now checks both calls. On failure it throws std::runtime_error (pybind11 turns it into RuntimeError) with the binding name and the CUDA error string. A failed cudaFuncSetAttribute throws before the launch and clears the error it set. The binding name is copied into the lambda, so the message doesn't depend on the lifetime of the name argument.

Notes

  • This catches launch-configuration errors reported by the launch call itself, not faults inside a running kernel.
  • Like C10_CUDA_KERNEL_LAUNCH_CHECK, the cudaGetLastError() check also picks up an error left pending by earlier CUDA work in the same runtime. That includes a sticky fault from an earlier kernel (e.g. an illegal memory access), which would now be raised here, under this binding's name.
  • Without the check, whether the failure stays silent depends on which cudart the extension uses. With its own static cudart (nvcc's default; Megakernels links -lcudart_static), nothing else in the process sees the error; that's the setup measured above. If the extension shares PyTorch's libcudart, the error stays pending and can surface at PyTorch's next launch check, under the wrong operation. I haven't tested that case.
  • No kernel in this repo calls bind_kernel. In-tree kernels use launch_kernel from torchutils.cuh, which already wraps both calls in CUDACHECK (print and exit), or launch directly. The callers are downstream projects such as Megakernels, which pins an older revision and would get this only after a submodule bump.

AI disclosure: I used Claude while writing this fix and its repro; I reviewed every line, ran the tests on a rented GPU, and can answer questions about any of it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant