Files
foxhunt/crates/ml-explainability
jgrusewski 952302149e feat(explainability): GPU-resident Integrated Gradients kernel
Implements the CUDA kernels (`interpolate_input`, `perturb_dimension`)
and wires `compute_gpu` in `integrated_gradients.rs` to run the full
IG algorithm on-device. Replaces the stub that previously returned
`MLError::ModelError("IG GPU kernels not available: cubins not yet
wired")` and fell back to CPU.

Algorithm: for each of `num_steps` interpolation points along the
baseline->input path, compute central-finite-difference gradients for
all `num_features` dimensions via two GPU forward passes per feature.
Single cubin (`ig_kernels.cubin`) compiled via build.rs (following the
crates/ml/build.rs pattern — nvcc, `-arch=sm_\${CUDA_COMPUTE_CAP}`, O3,
f32) and embedded with `include_bytes!`. The `forward_fn` consumers
operate on `GpuTensor`, so no dtoh round-trips inside the inner loop —
only the scalar `[1]` tensor output of each forward pass is pulled to
host (via `to_scalar`) per gradient sample. The three scratch buffers
(`interpolated`, `x_plus`, `x_minus`) are allocated once outside the
step loop and reused across all steps and features. No atomicAdd —
the kernels are trivial 1-D element-wise writes.

Tests: existing CPU tests pass unchanged. Added GPU smoke test
`test_ig_compute_gpu_linear_model` (gated `#[cfg(feature = \"cuda\")]`
+ `#[ignore]`) that builds a linear model as a `GpuTensor`-native
forward (elementwise mul + mean), verifies the completeness axiom
within 1%, and cross-checks GPU attributions against the CPU path
within 5%. Passes locally on RTX 3050 Ti (sm_86).

Removes 3 TODO markers and the embedded `_IG_CUDA_SRC` const that
were awaiting this work, along with the `#[allow(unused_variables)]`
stub attribute on `compute_gpu`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-23 08:32:09 +02:00
..