Fix CUDA illegal memory access in backward with mixed requires_grad (#1056)

[nv-slang-bot]

Summary

Fixes a hard process abort (CUDA_ERROR_ILLEGAL_ADDRESS) when a [Differentiable] Slang function with multiple IDiffTensor parameters is called through the PyTorch integration and only some of its inputs have requires_grad=True.

• Status: Fix implemented, built, and verified on CUDA + Vulkan (native and fallback torch-bridge modes). Draft pending human review.
• Link: Closes [Bug] CUDA illegal memory access (process abort) in backward when a [Differentiable] function has two tensor params and only one requires grad #1056.
• Verdict: Root cause confirmed by reproducing the abort on the unfixed tree; Approach A (bind a scratch grad buffer for no-grad input pairs in backward) resolves all three grad-mismatch cases with no kernel/ABI/cache-key change.
• Next-action: Human maintainer review; promote from draft when satisfied.
• Blocker: None.

Root cause

In the backward pass every IDiffTensor parameter binds as a Slang DiffTensor — this is keyed on the call mode (backward), not on each tensor's requires_grad (slangpy/builtin/tensorcommon.py). A DiffTensor always carries an _grad_out atomic accumulator, and the generated backward path for its tensor loads scatters into _grad_out with no per-tensor requires_grad guard (slangpy/slang/difftensor.slang).

But NativeCallData::autograd_backward (src/slangpy_ext/utils/slangpy.cpp) only allocated a gradient buffer for inputs with requires_grad=True, leaving pair->grad = None otherwise. The marshall then skips writing that parameter's _grad_out fields, so on CUDA — where the atomic target is a raw device pointer — the scatter writes through a dangling pointer and faults. This is why yes/yes worked while yes/no and no/yes aborted.

Fix

In autograd_backward, for an input pair whose primal has requires_grad=False, bind a throwaway zeroed buffer (instead of None) so the kernel's unconditional scatter lands in valid memory. The buffer is discarded, and None is still returned to torch for that leaf, so autograd correctly reports no gradient for it. This mirrors the reported t.detach().requires_grad_(True) workaround, done internally.

The buffer is bound on all backends, not just CUDA: on CUDA it prevents the illegal-address fault; on other backends type resolution rejects a diff tensor with no associated output gradient at kernel-generation time (see the != DeviceType.cuda TypeError guard in slangpy/builtin/tensorcommon.py), so the buffer is required there too — a CUDA-only gate would leave the non-CUDA mixed-grad path raising ResolveException. The bind is therefore unconditional.

Testing

Adds test_mixed_requires_grad_idifftensor to slangpy/tests/slangpy_tests/test_torchintegration.py, covering all three cases (yes/yes, yes/no, no/yes) with two IDiffTensor<float,1> inputs and asserting gradients appear only for requires_grad=True leaves. It runs under the autouse torch_bridge_mode fixture, exercising both the native and Python-fallback bridge paths.

Verified locally on an L40S (CUDA + Vulkan):

• Regression test: all 3 cases PASS on CUDA + Vulkan, native + fallback (12/12). On the unfixed tree the yes/no and no/yes cases abort the process with CUDA_ERROR_ILLEGAL_ADDRESS.
• Full test_torchintegration.py: 342 passed, 84 skipped, 0 failed (matches the pre-change baseline).
• pre-commit run on the changed files: clean.

Scope

Scoped to the torch-integration backward path. The plain-Tensor (non-torch) CUDA guard-bypass at tensorcommon.py (the != DeviceType.cuda gate that suppresses the missing-grad TypeError on CUDA) is a separate latent hardening item and is intentionally not folded in here — naively un-gating it would raise at cache time and regress the valid mixed-grad case this fix enables.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}

[ccummingsNV]

Adding a potentially significant memory allocation to solve a null pointer seems overkill. Is there no way to deal with the null pointer correctly?

[nv-slang-bot]

@ccummingsNV The allocation is one buffer per no-grad input, sized like its primal and discarded after the backward pass — so for a large frozen/detached input it is wasted memory every backward. That cost is real; below is why the fix currently binds a buffer unconditionally, and the cleaner alternative.

The fault has two different forms depending on backend:

• CUDA — AtomicTensor's storage is a raw T* (slangpy/slang/atomics.slang), and the generated backward kernel scatters _grad_out unconditionally for every IDiffTensor param. With no buffer bound the pointer is dangling → CUDA_ERROR_ILLEGAL_ADDRESS.
• Other backends — type resolution rejects an IDiffTensor bound in a diff pass with no output gradient at kernel-generation time (slangpy/builtin/tensorcommon.py, the has_grad_out and d_out is None check), so it never reaches dispatch.

A CUDA-side null-check on the atomic write would address only the first; the non-CUDA path fails earlier, at generation. So a null-guard alone isn't sufficient across backends.

The fix that removes the allocation entirely is to bind a no-grad input as PrimalTensor instead of DiffTensor. PrimalTensor already conforms to IDiffTensor, carries no _grad_out buffer, and compiles _write_grad_each / _accumulate_grad_each to no-ops (slangpy/slang/primaltensor.slang) — it exists precisely for "an IDiffTensor where gradients aren't needed." No allocation, no scatter, no null pointer, and it satisfies both backends (no _grad_out ⇒ no non-CUDA rejection).

Why it isn't a one-liner: the DiffTensor-vs-PrimalTensor choice is currently made by call mode (tensorcommon.py), not per-tensor requires_grad, and the torch tensor's kernel signature is [Dn,Sm] (ndim + scalar type) — it doesn't encode requires_grad. A mixed-grad call and an all-grad call therefore hash to the same signature and reuse one backward kernel in which every input is a DiffTensor. Binding a no-grad input as PrimalTensor means threading requires_grad into type resolution and the signature so the backward kernel specializes per grad-config.

That's more involved than this PR, but it's the correct fix and avoids both the allocation and the null pointer. I can switch this PR to the PrimalTensor approach if you'd prefer it over the scratch-buffer workaround.

_{🤖 Generated by an automated SlangPy coworker — may be inaccurate. A human maintainer should verify.}