ops: transformer op fwd/bwd CUDA kernels + Tensor wrappers

add/mul/add_bias(+sum_rows)/rms_norm/silu/rope/softmax/cross_entropy,
each with its analytic backward, in csrc/ops/nn.cu (inlined warp/block
reductions). FFI declarations + nn.cu in build.rs (no_cuda gated). Tensor
gains the matching thin wrappers; DType grows I32 for cross-entropy targets.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-15 15:44:09 +08:00
parent 88fbe0a85d
commit 5aef3742d6
5 changed files with 815 additions and 0 deletions

View File

@@ -32,6 +32,7 @@ fn main() {
.file("../../csrc/test/vecadd.cu")
.file("../../csrc/ops/elementwise.cu")
.file("../../csrc/ops/gemm.cu")
.file("../../csrc/ops/nn.cu")
.compile("xtrain_cuda_kernels");
}