From c1a8a0c2a461f3cb49e75750ad30ad3ce7f09e89 Mon Sep 17 00:00:00 2001 From: wadkisson Date: Tue, 7 Jul 2026 12:58:30 -0700 Subject: [PATCH 1/3] implement libtorch calls for MHA --- .gitignore | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.gitignore b/.gitignore index d26e6b79..8cb462e3 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,10 @@ NN/Examples/Data/*.npy _external/ third_party/interval/ +# Local LibTorch unpack (CUDA prerequisite; see README) +/libtorch/ +.lake/build/libtorch.path + # Python caches and artifacts __pycache__/ **/__pycache__/ From b339f82936416c1e5a03d4a3cc8f5e2c9c5a2da5 Mon Sep 17 00:00:00 2001 From: wadkisson Date: Tue, 7 Jul 2026 12:58:41 -0700 Subject: [PATCH 2/3] implement libtorch attention kernel --- NN/Runtime/Autograd/Engine/Cuda/Kernels.lean | 23 +- .../Autograd/Engine/Cuda/Ops/Attention.lean | 64 ++-- README.md | 13 +- csrc/cuda/README.md | 3 + csrc/cuda/common/torchlean_cuda_buffer.h | 13 + csrc/cuda/kernels/torchlean_cuda_kernels.cu | 286 +----------------- .../kernels/torchlean_cuda_kernels_stub.c | 61 +--- .../kernels/torchlean_flash_attention_fwd.cpp | 49 +++ lakefile.lean | 115 +++++-- 9 files changed, 247 insertions(+), 380 deletions(-) create mode 100644 csrc/cuda/kernels/torchlean_flash_attention_fwd.cpp diff --git a/NN/Runtime/Autograd/Engine/Cuda/Kernels.lean b/NN/Runtime/Autograd/Engine/Cuda/Kernels.lean index 06737f93..72af7adc 100644 --- a/NN/Runtime/Autograd/Engine/Cuda/Kernels.lean +++ b/NN/Runtime/Autograd/Engine/Cuda/Kernels.lean @@ -231,8 +231,8 @@ Inputs are row-major buffers with shapes: - `mask`: `(batch, n, n)` encoded as `0.0/1.0` when `hasMask != 0`; otherwise ignored. Output has shape `(batch, n, d)` and computes the same no-dropout masked attention semantics as: -`hardMaskedSoftmax((Q Kᵀ) * scale, mask) V`. Blocked mask entries contribute zero softmax -numerator; no finite sentinel is inserted. +`softmax((Q Kᵀ) * scale + maskFill) V`, where blocked mask entries use TorchLean's +`-1000.0` fill convention. This is a fused native runtime primitive. The proof layer contract is `Spec.flashAttention` in `NN/Spec/Layers/FlashAttention.lean`; the native kernel is part of the CUDA runtime boundary. @@ -241,20 +241,11 @@ This is a fused native runtime primitive. The proof layer contract is `Spec.flas opaque flashAttentionFwd (Q K V mask : Buffer) (hasMask batch n d : UInt32) (scale : Float) : Buffer -/-- Fused VJP component `∂L/∂Q` for `flashAttentionFwd`. -/ -@[extern "torchlean_cuda_buffer_flash_attention_bwd_q"] -opaque flashAttentionBwdQ - (Q K V mask dOut : Buffer) (hasMask batch n d : UInt32) (scale : Float) : Buffer - -/-- Fused VJP component `∂L/∂K` for `flashAttentionFwd`. -/ -@[extern "torchlean_cuda_buffer_flash_attention_bwd_k"] -opaque flashAttentionBwdK - (Q K V mask dOut : Buffer) (hasMask batch n d : UInt32) (scale : Float) : Buffer - -/-- Fused VJP component `∂L/∂V` for `flashAttentionFwd`. -/ -@[extern "torchlean_cuda_buffer_flash_attention_bwd_v"] -opaque flashAttentionBwdV - (Q K V mask dOut : Buffer) (hasMask batch n d : UInt32) (scale : Float) : Buffer +/-- Fused VJP `(dQ, dK, dV)` for `flashAttentionFwd`. -/ +@[extern "torchlean_cuda_buffer_flash_attention_bwd"] +opaque flashAttentionBwd + (Q K V mask dOut : Buffer) (hasMask batch n d : UInt32) (scale : Float) : + Buffer × Buffer × Buffer /-- Row-major transpose of a 2D buffer. diff --git a/NN/Runtime/Autograd/Engine/Cuda/Ops/Attention.lean b/NN/Runtime/Autograd/Engine/Cuda/Ops/Attention.lean index b42b52d6..3930b315 100644 --- a/NN/Runtime/Autograd/Engine/Cuda/Ops/Attention.lean +++ b/NN/Runtime/Autograd/Engine/Cuda/Ops/Attention.lean @@ -33,8 +33,9 @@ Forward structure matches `Spec.MultiHeadAttention.forward`: 4. combine heads, then output projection `@ Wo` Masking: -- If `mask` is provided, we upload it as a float32 `{0,1}` matrix and apply hard-mask semantics: - blocked entries contribute zero softmax numerator, and their score gradients are zero. +- If `mask` is provided, we upload it as a float32 `{0,1}` matrix and apply the same semantics as + the spec: masked logits are replaced by `-1000.0` before softmax, and gradients through masked + entries are zeroed. - This incurs a host-to-device copy for the mask (since the mask is a host `Tensor Bool`). -/ @@ -72,7 +73,8 @@ def multiHeadAttention let Kh := Buffer.swapAdjacentAtDepth K dimsView depth0 -- (numHeads,n,headDim) let Vh := Buffer.swapAdjacentAtDepth V dimsView depth0 -- (numHeads,n,headDim) let scale : Float := 1.0 / Float.sqrt (Float.ofNat headDim) - -- Optional mask: `mask[i,j]=true` means allowed; false entries contribute zero numerator. + -- Optional mask: `mask[i,j]=true` means allowed; false positions are set to a large negative + -- constant before softmax. let (maskB, hasMask) : Buffer × UInt32 := match mask with | none => (Buffer.zeros 0, 0) @@ -93,13 +95,23 @@ def multiHeadAttention let KhT := Buffer.swapAdjacentAtDepth Kh dimsHead depth1 let scores := Buffer.bmm Qh KhT h32 n32 head32 n32 let scaled0 := Buffer.scale scores scale - let rowsFold32 ← u32 (numHeads * n) - let attnOwned := + let (scaled, maskCleanup) ← match mask with - | none => rowSoftmaxForward scaled0 rowsFold32 n32 - | some _ => rowHardMaskedSoftmaxForward scaled0 maskB rowsFold32 n32 + | none => pure (scaled0, ([] : List Buffer)) + | some _ => do + let total : Nat := numHeads * n * n + let total32 ← u32 total + let ones := Buffer.full total32 1.0 + let invMask := Buffer.sub ones maskB + let fill := Buffer.full total32 (-1000.0) + let scaledMask := Buffer.mul scaled0 maskB + let fillMask := Buffer.mul fill invMask + let scaled := Buffer.add scaledMask fillMask + pure (scaled, [ones, invMask, fill, scaledMask, fillMask]) + let rowsFold32 ← u32 (numHeads * n) + let attnOwned := rowSoftmaxForward scaled rowsFold32 n32 let outHeads := Buffer.bmm attnOwned.value Vh h32 n32 n32 head32 - pure (outHeads, [KhT, scores, scaled0, attnOwned.value] ++ + pure (outHeads, [KhT, scores, scaled0, scaled, attnOwned.value] ++ maskCleanup ++ attnOwned.workspace) -- combine heads: swap to (n,numHeads,headDim), then reshape to (n,projDim) let swapped := Buffer.swapAdjacentAtDepth outHeads dimsHead depth0 -- (n,numHeads,headDim) @@ -127,23 +139,28 @@ def multiHeadAttention let dOutHeads := Buffer.swapAdjacentAtDepth dSwapped dimsView depth0 -- (numHeads,n,headDim) let (dQh, dKh, dVh) ← if useFlash then - -- Fused VJP for the native attention primitive. The kernels recompute the row-wise - -- softmax summaries instead of materializing/storing the full attention matrix. - let dQh := Buffer.flashAttentionBwdQ Qh Kh Vh maskB dOutHeads hasMask h32 n32 head32 scale - let dKh := Buffer.flashAttentionBwdK Qh Kh Vh maskB dOutHeads hasMask h32 n32 head32 scale - let dVh := Buffer.releaseThen dOutHeads <| - Buffer.flashAttentionBwdV Qh Kh Vh maskB dOutHeads hasMask h32 n32 head32 scale - pure (dQh, dKh, dVh) + pure <| Buffer.flashAttentionBwd Qh Kh Vh maskB dOutHeads hasMask h32 n32 head32 scale else let dimsAttn : Array Nat := #[numHeads, n, n] let KhT := Buffer.swapAdjacentAtDepth Kh dimsHead depth1 let scores := Buffer.bmm Qh KhT h32 n32 head32 n32 let scaled0 := Buffer.scale scores scale - let rowsFold32 ← u32 (numHeads * n) - let attnOwned := + let scaled ← match mask with - | none => rowSoftmaxForward scaled0 rowsFold32 n32 - | some _ => rowHardMaskedSoftmaxForward scaled0 maskB rowsFold32 n32 + | none => pure scaled0 + | some _ => do + let total : Nat := numHeads * n * n + let total32 ← u32 total + let ones := Buffer.full total32 1.0 + let invMask := Buffer.sub ones maskB + let fill := Buffer.full total32 (-1000.0) + let scaledMask := Buffer.mul scaled0 maskB + let fillMask := Buffer.mul fill invMask + pure <| Buffer.releaseThen ones <| Buffer.releaseThen invMask <| + Buffer.releaseThen fill <| Buffer.releaseThen scaledMask <| + Buffer.releaseThen fillMask <| Buffer.add scaledMask fillMask + let rowsFold32 ← u32 (numHeads * n) + let attnOwned := rowSoftmaxForward scaled rowsFold32 n32 let VhT := Buffer.swapAdjacentAtDepth Vh dimsHead depth1 let dAttn := Buffer.releaseThen VhT <| Buffer.bmm dOutHeads VhT h32 n32 head32 n32 @@ -153,16 +170,19 @@ def multiHeadAttention Buffer.bmm attnT dOutHeads h32 n32 n32 head32 let dScaled := rowSoftmaxBwd attnOwned.value dAttn rowsFold32 n32 let dScoresMasked := Buffer.scale dScaled scale - let dScores := dScoresMasked + let dScores := + match mask with + | none => dScoresMasked + | some _ => Buffer.mul dScoresMasked maskB let dQh := Buffer.bmm dScores Kh h32 n32 n32 head32 let dScoresT := Buffer.swapAdjacentAtDepth dScores dimsAttn depth1 let dKh := Buffer.releaseThen dScoresT <| Buffer.bmm dScoresT Qh h32 n32 n32 head32 let dQh := Buffer.releaseThen KhT <| Buffer.releaseThen scores <| - Buffer.releaseThen scaled0 <| + Buffer.releaseThen scaled0 <| Buffer.releaseThen scaled <| Buffer.releaseThen attnOwned.value <| Buffer.releaseThen dAttn <| Buffer.releaseThen dScaled <| Buffer.releaseThen dScoresMasked <| - attnOwned.releaseWorkspaceThen dQh + Buffer.releaseThen dScores <| attnOwned.releaseWorkspaceThen dQh pure (dQh, dKh, dVh) -- Backprop split-head permutations: swap back to (n,numHeads,headDim), then view as (n,projDim). let dQ := Buffer.releaseThen dQh <| Buffer.swapAdjacentAtDepth dQh dimsHead depth0 diff --git a/README.md b/README.md index 5b2a2d70..e1b13d24 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,9 @@ lake build lake exe torchlean quickstart_mlp --cpu --steps 10 --dtype float32 --backend eager lake exe torchlean quickstart_mlp --cpu --steps 10 --dtype float --backend eager -# Optional CUDA run, if the CUDA toolkit and an NVIDIA GPU are available: +# Optional CUDA run (requires CUDA toolkit, LibTorch, and an NVIDIA GPU): +# https://pytorch.org/get-started/locally/ → LibTorch → extract to ./libtorch +# or set TORCHLEAN_LIBTORCH_HOME=/path/to/libtorch lake build -K cuda=true lake exe -K cuda=true torchlean mlp --cuda --fast-kernels --steps 1000 ``` @@ -35,6 +37,15 @@ native GPU runtime path and checks that the CUDA backend is available. Theorem statements that mention CUDA cite the native-runtime boundary in `TRUST_BOUNDARIES.md` instead of treating a kernel launch as Lean proof evidence. +CUDA builds require [LibTorch](https://pytorch.org/get-started/locally/) in addition to the NVIDIA +CUDA toolkit. Pick the **LibTorch** download (not Pip) with a CUDA version that matches your +toolkit. Extract it to `./libtorch` at the repo root, or point the build at it with +`TORCHLEAN_LIBTORCH_HOME` or `-K libtorch_home=/path/to/libtorch`. + +On `lake build -K cuda=true`, TorchLean looks for `include/torch/torch.h` under those locations. +If LibTorch is missing, the build stops with download instructions. When found, the resolved path +is written to `.lake/build/libtorch.path`. + TorchLean is pinned by `lean-toolchain` and currently builds with `leanprover/lean4:v4.31.0`. diff --git a/csrc/cuda/README.md b/csrc/cuda/README.md index dada300b..40370961 100644 --- a/csrc/cuda/README.md +++ b/csrc/cuda/README.md @@ -32,6 +32,9 @@ used by default so `lake build` works without a CUDA toolkit. Real CUDA builds a lake build -R -K cuda=true -K cuda_home=/usr/local/cuda ``` +CUDA builds also require LibTorch; see the root `README.md`. The path is resolved by +`scripts/setup/resolve_libtorch.sh` before native backends compile. + ## CUDA Graph Status TorchLean's current CUDA path is eager: each autograd step records a Lean runtime tape and dispatches diff --git a/csrc/cuda/common/torchlean_cuda_buffer.h b/csrc/cuda/common/torchlean_cuda_buffer.h index f8c0768e..1a344dfa 100644 --- a/csrc/cuda/common/torchlean_cuda_buffer.h +++ b/csrc/cuda/common/torchlean_cuda_buffer.h @@ -126,6 +126,19 @@ LEAN_EXPORT uint64_t torchlean_cuda_allocator_free_count(uint32_t u); LEAN_EXPORT uint64_t torchlean_cuda_allocator_device_free_bytes(uint32_t u); LEAN_EXPORT uint64_t torchlean_cuda_allocator_device_total_bytes(uint32_t u); +// Lean `Buffer × Buffer × Buffer` as nested pairs. +static inline lean_object* torchlean_cuda_box_three_buffers( + torchlean_cuda_buffer* first, torchlean_cuda_buffer* second, + torchlean_cuda_buffer* third) { + lean_object* tail = lean_alloc_ctor(0, 2, 0); + lean_ctor_set(tail, 0, torchlean_cuda_buffer_box(second)); + lean_ctor_set(tail, 1, torchlean_cuda_buffer_box(third)); + lean_object* out = lean_alloc_ctor(0, 2, 0); + lean_ctor_set(out, 0, torchlean_cuda_buffer_box(first)); + lean_ctor_set(out, 1, tail); + return out; +} + #ifdef __cplusplus } // extern "C" #endif diff --git a/csrc/cuda/kernels/torchlean_cuda_kernels.cu b/csrc/cuda/kernels/torchlean_cuda_kernels.cu index 7ef6b2dc..a388fc7c 100644 --- a/csrc/cuda/kernels/torchlean_cuda_kernels.cu +++ b/csrc/cuda/kernels/torchlean_cuda_kernels.cu @@ -16,9 +16,20 @@ #include #include -// CUDA kernels above the raw buffer ops: reductions, broadcasting, gather/scatter, batched matmul, -// selective scan, FFT helpers, and the simple attention path. -// Everything is flat row-major; wrappers check ranks and shapes before launch. +// CUDA implementation of TorchLean's general float32 tensor kernels. +// +// This file covers kernels that sit above raw elementwise buffer ops and below model-specific +// layers: reductions over axes, broadcasting, gathers/scatters, batched matmul, selective scan, and +// the correctness-first fused attention path. Host `LEAN_EXPORT` wrappers validate Lean-provided +// shape metadata before launching device kernels; device kernels assume that validation succeeded. +// +// Layout conventions: +// - all tensor buffers are flat row-major arrays; +// - rank-polymorphic kernels are capped at `kMaxRank` so stack coordinate arrays stay bounded; +// - atomic accumulation has deterministic alternatives where reproducibility matters; +// - cuBLAS calls use the row-major-as-transposed-column-major convention documented at call sites. +// - cuFFT R2C/C2R kernels expose spectra as packed real/imag float32 buffers: +// `(batch, n/2+1, 2)`, last channel `[re, im]`. static constexpr int kBlockSize = 256; static constexpr int kMaxRank = 8; @@ -1646,271 +1657,6 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_selective_scan_diag_va return torchlean_cuda_buffer_box(out); } -// Simple attention kernel. -// -// Match TorchLean's scaled-dot-product attention spec: scores = QK^T * scale, optional dense -// Boolean mask, stable row softmax, then multiplication by V. This is deliberately direct. Backward -// recomputes the row softmax statistics instead of storing the full attention matrix, so this is not -// the production IO-tiled FlashAttention schedule. -// -// Mask convention: zero means disallowed. Disallowed entries contribute literal zero softmax -// numerator, matching TorchLean's hard-mask spec rather than a finite sentinel. -__device__ inline bool flash_attention_allowed(const float* mask, uint32_t hasMask, - uint32_t batchIdx, uint32_t i, uint32_t j, - uint32_t n) { - if (hasMask == 0) return true; - const size_t idx = ((size_t)batchIdx * (size_t)n + (size_t)i) * (size_t)n + (size_t)j; - return mask[idx] != 0.0f; -} - -__device__ inline float flash_attention_score(const float* Q, const float* K, const float* mask, - uint32_t hasMask, uint32_t batchIdx, uint32_t i, - uint32_t j, uint32_t n, uint32_t d, float scale) { - float dot = 0.0f; - const size_t qBase = ((size_t)batchIdx * (size_t)n + (size_t)i) * (size_t)d; - const size_t kBase = ((size_t)batchIdx * (size_t)n + (size_t)j) * (size_t)d; - for (uint32_t k = 0; k < d; ++k) { - dot += Q[qBase + (size_t)k] * K[kBase + (size_t)k]; - } - return dot * scale; -} - -__device__ inline void flash_attention_row_stats(const float* Q, const float* K, - const float* mask, uint32_t hasMask, - uint32_t batchIdx, uint32_t i, uint32_t n, - uint32_t d, float scale, float* rowMax, - float* denom) { - float m = -INFINITY; - for (uint32_t j = 0; j < n; ++j) { - if (!flash_attention_allowed(mask, hasMask, batchIdx, i, j, n)) continue; - float s = flash_attention_score(Q, K, mask, hasMask, batchIdx, i, j, n, d, scale); - if (s > m) m = s; - } - - float z = 0.0f; - for (uint32_t j = 0; j < n; ++j) { - if (!flash_attention_allowed(mask, hasMask, batchIdx, i, j, n)) continue; - float s = flash_attention_score(Q, K, mask, hasMask, batchIdx, i, j, n, d, scale); - z += expf(s - m); - } - *rowMax = m; - *denom = z; -} - -__device__ inline float flash_attention_prob(const float* Q, const float* K, const float* mask, - uint32_t hasMask, uint32_t batchIdx, uint32_t i, - uint32_t j, uint32_t n, uint32_t d, float scale, - float rowMax, float denom) { - if (!flash_attention_allowed(mask, hasMask, batchIdx, i, j, n) || denom == 0.0f) { - return 0.0f; - } - float s = flash_attention_score(Q, K, mask, hasMask, batchIdx, i, j, n, d, scale); - return expf(s - rowMax) / denom; -} - -__device__ inline float flash_attention_d_attn(const float* V, const float* dOut, - uint32_t batchIdx, uint32_t i, uint32_t j, - uint32_t n, uint32_t d) { - float acc = 0.0f; - const size_t dOutBase = ((size_t)batchIdx * (size_t)n + (size_t)i) * (size_t)d; - const size_t vBase = ((size_t)batchIdx * (size_t)n + (size_t)j) * (size_t)d; - for (uint32_t k = 0; k < d; ++k) { - acc += dOut[dOutBase + (size_t)k] * V[vBase + (size_t)k]; - } - return acc; -} - -__global__ void flash_attention_fwd_f32(const float* Q, const float* K, const float* V, - const float* mask, uint32_t hasMask, uint32_t batch, - uint32_t n, uint32_t d, float scale, float* out) { - const size_t idx = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; - const size_t total = (size_t)batch * (size_t)n * (size_t)d; - if (idx >= total) return; - - const uint32_t dv = (uint32_t)(idx % (size_t)d); - const uint32_t i = (uint32_t)((idx / (size_t)d) % (size_t)n); - const uint32_t b = (uint32_t)(idx / ((size_t)n * (size_t)d)); - - float rowMax = 0.0f; - float denom = 0.0f; - flash_attention_row_stats(Q, K, mask, hasMask, b, i, n, d, scale, &rowMax, &denom); - - float acc = 0.0f; - for (uint32_t j = 0; j < n; ++j) { - const float p = flash_attention_prob(Q, K, mask, hasMask, b, i, j, n, d, scale, rowMax, denom); - const size_t vIdx = ((size_t)b * (size_t)n + (size_t)j) * (size_t)d + (size_t)dv; - acc += p * V[vIdx]; - } - out[idx] = acc; -} - -__global__ void flash_attention_bwd_q_f32(const float* Q, const float* K, const float* V, - const float* mask, const float* dOut, uint32_t hasMask, - uint32_t batch, uint32_t n, uint32_t d, float scale, - float* dQ) { - const size_t idx = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; - const size_t total = (size_t)batch * (size_t)n * (size_t)d; - if (idx >= total) return; - - const uint32_t k = (uint32_t)(idx % (size_t)d); - const uint32_t i = (uint32_t)((idx / (size_t)d) % (size_t)n); - const uint32_t b = (uint32_t)(idx / ((size_t)n * (size_t)d)); - - float rowMax = 0.0f; - float denom = 0.0f; - flash_attention_row_stats(Q, K, mask, hasMask, b, i, n, d, scale, &rowMax, &denom); - - float rowDot = 0.0f; - for (uint32_t j = 0; j < n; ++j) { - const float p = flash_attention_prob(Q, K, mask, hasMask, b, i, j, n, d, scale, rowMax, denom); - rowDot += p * flash_attention_d_attn(V, dOut, b, i, j, n, d); - } - - float acc = 0.0f; - for (uint32_t j = 0; j < n; ++j) { - if (!flash_attention_allowed(mask, hasMask, b, i, j, n)) continue; - const float p = flash_attention_prob(Q, K, mask, hasMask, b, i, j, n, d, scale, rowMax, denom); - const float dAttn = flash_attention_d_attn(V, dOut, b, i, j, n, d); - const float dScore = p * (dAttn - rowDot) * scale; - const size_t kIdx = ((size_t)b * (size_t)n + (size_t)j) * (size_t)d + (size_t)k; - acc += dScore * K[kIdx]; - } - dQ[idx] = acc; -} - -__global__ void flash_attention_bwd_k_f32(const float* Q, const float* K, const float* V, - const float* mask, const float* dOut, uint32_t hasMask, - uint32_t batch, uint32_t n, uint32_t d, float scale, - float* dK) { - const size_t idx = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; - const size_t total = (size_t)batch * (size_t)n * (size_t)d; - if (idx >= total) return; - - const uint32_t k = (uint32_t)(idx % (size_t)d); - const uint32_t j = (uint32_t)((idx / (size_t)d) % (size_t)n); - const uint32_t b = (uint32_t)(idx / ((size_t)n * (size_t)d)); - - float acc = 0.0f; - for (uint32_t i = 0; i < n; ++i) { - float rowMax = 0.0f; - float denom = 0.0f; - flash_attention_row_stats(Q, K, mask, hasMask, b, i, n, d, scale, &rowMax, &denom); - - float rowDot = 0.0f; - for (uint32_t t = 0; t < n; ++t) { - const float p = flash_attention_prob(Q, K, mask, hasMask, b, i, t, n, d, scale, rowMax, - denom); - rowDot += p * flash_attention_d_attn(V, dOut, b, i, t, n, d); - } - - if (!flash_attention_allowed(mask, hasMask, b, i, j, n)) continue; - const float p = flash_attention_prob(Q, K, mask, hasMask, b, i, j, n, d, scale, rowMax, denom); - const float dAttn = flash_attention_d_attn(V, dOut, b, i, j, n, d); - const float dScore = p * (dAttn - rowDot) * scale; - const size_t qIdx = ((size_t)b * (size_t)n + (size_t)i) * (size_t)d + (size_t)k; - acc += dScore * Q[qIdx]; - } - dK[idx] = acc; -} - -__global__ void flash_attention_bwd_v_f32(const float* Q, const float* K, const float* V, - const float* mask, const float* dOut, uint32_t hasMask, - uint32_t batch, uint32_t n, uint32_t d, float scale, - float* dV) { - const size_t idx = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; - const size_t total = (size_t)batch * (size_t)n * (size_t)d; - if (idx >= total) return; - - const uint32_t dv = (uint32_t)(idx % (size_t)d); - const uint32_t j = (uint32_t)((idx / (size_t)d) % (size_t)n); - const uint32_t b = (uint32_t)(idx / ((size_t)n * (size_t)d)); - - float acc = 0.0f; - for (uint32_t i = 0; i < n; ++i) { - float rowMax = 0.0f; - float denom = 0.0f; - flash_attention_row_stats(Q, K, mask, hasMask, b, i, n, d, scale, &rowMax, &denom); - const float p = flash_attention_prob(Q, K, mask, hasMask, b, i, j, n, d, scale, rowMax, denom); - const size_t dOutIdx = ((size_t)b * (size_t)n + (size_t)i) * (size_t)d + (size_t)dv; - acc += p * dOut[dOutIdx]; - } - dV[idx] = acc; -} - -extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_fwd( - b_lean_obj_arg QObj, b_lean_obj_arg KObj, b_lean_obj_arg VObj, b_lean_obj_arg MaskObj, - uint32_t hasMask, uint32_t batch, uint32_t n, uint32_t d, double scaleHost) { - torchlean_cuda_buffer* Q = torchlean_cuda_buffer_unbox(QObj); - torchlean_cuda_buffer* K = torchlean_cuda_buffer_unbox(KObj); - torchlean_cuda_buffer* V = torchlean_cuda_buffer_unbox(VObj); - torchlean_cuda_buffer* mask = torchlean_cuda_buffer_unbox(MaskObj); - const size_t qkvSz = - checked_mul3_size((size_t)batch, (size_t)n, (size_t)d, - "torchlean_cuda_buffer_flash_attention_fwd: Q/K/V size overflow"); - const size_t maskSz = - checked_mul3_size((size_t)batch, (size_t)n, (size_t)n, - "torchlean_cuda_buffer_flash_attention_fwd: mask size overflow"); - if (Q->size != qkvSz || K->size != qkvSz || V->size != qkvSz) { - lean_internal_panic("torchlean_cuda_buffer_flash_attention_fwd: Q/K/V size mismatch"); - } - if (hasMask != 0 && mask->size != maskSz) { - lean_internal_panic("torchlean_cuda_buffer_flash_attention_fwd: mask size mismatch"); - } - - torchlean_cuda_buffer* out = torchlean_cuda_buffer_alloc(qkvSz); - if (qkvSz == 0) { - return torchlean_cuda_buffer_box(out); - } - dim3 blocks = blocks_for(qkvSz); - dim3 threads = dim3(kBlockSize); - flash_attention_fwd_f32<<>>(Q->data, K->data, V->data, mask->data, hasMask, - batch, n, d, (float)scaleHost, out->data); - checkCuda(cudaGetLastError(), "cuda flashAttention forward kernel launch failed"); - return torchlean_cuda_buffer_box(out); -} - -#define TORCHLEAN_FLASH_BWD_EXPORT(name, kernel, label) \ - extern "C" LEAN_EXPORT lean_obj_res name( \ - b_lean_obj_arg QObj, b_lean_obj_arg KObj, b_lean_obj_arg VObj, b_lean_obj_arg MaskObj, \ - b_lean_obj_arg DOutObj, uint32_t hasMask, uint32_t batch, uint32_t n, uint32_t d, \ - double scaleHost) { \ - torchlean_cuda_buffer* Q = torchlean_cuda_buffer_unbox(QObj); \ - torchlean_cuda_buffer* K = torchlean_cuda_buffer_unbox(KObj); \ - torchlean_cuda_buffer* V = torchlean_cuda_buffer_unbox(VObj); \ - torchlean_cuda_buffer* mask = torchlean_cuda_buffer_unbox(MaskObj); \ - torchlean_cuda_buffer* dOut = torchlean_cuda_buffer_unbox(DOutObj); \ - const size_t qkvSz = checked_mul3_size((size_t)batch, (size_t)n, (size_t)d, \ - label ": Q/K/V/dOut size overflow"); \ - const size_t maskSz = checked_mul3_size((size_t)batch, (size_t)n, (size_t)n, \ - label ": mask size overflow"); \ - if (Q->size != qkvSz || K->size != qkvSz || V->size != qkvSz || dOut->size != qkvSz) { \ - lean_internal_panic(label ": Q/K/V/dOut size mismatch"); \ - } \ - if (hasMask != 0 && mask->size != maskSz) { \ - lean_internal_panic(label ": mask size mismatch"); \ - } \ - torchlean_cuda_buffer* out = torchlean_cuda_buffer_alloc(qkvSz); \ - if (qkvSz == 0) return torchlean_cuda_buffer_box(out); \ - dim3 blocks = blocks_for(qkvSz); \ - dim3 threads = dim3(kBlockSize); \ - kernel<<>>(Q->data, K->data, V->data, mask->data, dOut->data, hasMask, batch, \ - n, d, (float)scaleHost, out->data); \ - checkCuda(cudaGetLastError(), label " kernel launch failed"); \ - return torchlean_cuda_buffer_box(out); \ - } - -TORCHLEAN_FLASH_BWD_EXPORT(torchlean_cuda_buffer_flash_attention_bwd_q, - flash_attention_bwd_q_f32, - "cuda flashAttention backward Q") -TORCHLEAN_FLASH_BWD_EXPORT(torchlean_cuda_buffer_flash_attention_bwd_k, - flash_attention_bwd_k_f32, - "cuda flashAttention backward K") -TORCHLEAN_FLASH_BWD_EXPORT(torchlean_cuda_buffer_flash_attention_bwd_v, - flash_attention_bwd_v_f32, - "cuda flashAttention backward V") - -#undef TORCHLEAN_FLASH_BWD_EXPORT - extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_transpose2d(b_lean_obj_arg BObj, uint32_t rows, uint32_t cols) { torchlean_cuda_buffer* b = torchlean_cuda_buffer_unbox(BObj); @@ -2081,7 +1827,7 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_broadcast_to(b_lean_ob lean_internal_panic("torchlean_cuda_buffer_broadcast_to: input size mismatch"); } - // Check broadcast shape agreement before launch. + // Check broadcast shape agreement at the FFI boundary before launching the kernel. for (size_t ax = 0; ax < rankOut; ++ax) { uint32_t mv = h.axisMap[ax]; if (mv == 0) continue; @@ -2206,7 +1952,7 @@ extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_reduce_from_broadcast( lean_internal_panic("torchlean_cuda_buffer_reduce_from_broadcast: dOut size mismatch"); } - // Use the same broadcast checks as the forward path. + // Check the same broadcast contract used by the forward path. for (size_t ax = 0; ax < rankOut; ++ax) { uint32_t mv = h.axisMap[ax]; if (mv == 0) continue; diff --git a/csrc/cuda/kernels/torchlean_cuda_kernels_stub.c b/csrc/cuda/kernels/torchlean_cuda_kernels_stub.c index e2926a5a..6eb2bc1c 100644 --- a/csrc/cuda/kernels/torchlean_cuda_kernels_stub.c +++ b/csrc/cuda/kernels/torchlean_cuda_kernels_stub.c @@ -9,8 +9,11 @@ #include #include -// CPU version of the general tensor-kernel FFI symbols. -// Keep its shape checks and edge cases aligned with `torchlean_cuda_kernels.cu`. +// CPU fallback for TorchLean's general tensor-kernel FFI surface. +// +// This file mirrors `torchlean_cuda_kernels.cu` symbol-for-symbol so the Lean runtime can link on +// machines without CUDA. Keep the shape checks and edge-case conventions in lockstep with the CUDA +// file; these stubs are also useful as readable reference implementations for tests and audits. static const float kTorchLeanPiF = 3.14159265358979323846f; @@ -673,6 +676,7 @@ static inline float flash_attention_score_stub(const torchlean_cuda_buffer* Q, const torchlean_cuda_buffer* mask, uint32_t hasMask, size_t batchIdx, size_t i, size_t j, size_t n, size_t d, float scale) { + if (!flash_attention_allowed_stub(mask, hasMask, batchIdx, i, j, n)) return -1000.0f; float dot = 0.0f; const size_t qBase = (batchIdx * n + i) * d; const size_t kBase = (batchIdx * n + j) * d; @@ -688,13 +692,11 @@ static inline void flash_attention_row_stats_stub(const torchlean_cuda_buffer* Q float* rowMax, float* denom) { float m = -INFINITY; for (size_t j = 0; j < n; ++j) { - if (!flash_attention_allowed_stub(mask, hasMask, batchIdx, i, j, n)) continue; float s = flash_attention_score_stub(Q, K, mask, hasMask, batchIdx, i, j, n, d, scale); if (s > m) m = s; } float z = 0.0f; for (size_t j = 0; j < n; ++j) { - if (!flash_attention_allowed_stub(mask, hasMask, batchIdx, i, j, n)) continue; float s = flash_attention_score_stub(Q, K, mask, hasMask, batchIdx, i, j, n, d, scale); z += expf(s - m); } @@ -708,9 +710,6 @@ static inline float flash_attention_prob_stub(const torchlean_cuda_buffer* Q, uint32_t hasMask, size_t batchIdx, size_t i, size_t j, size_t n, size_t d, float scale, float rowMax, float denom) { - if (!flash_attention_allowed_stub(mask, hasMask, batchIdx, i, j, n) || denom == 0.0f) { - return 0.0f; - } float s = flash_attention_score_stub(Q, K, mask, hasMask, batchIdx, i, j, n, d, scale); return expf(s - rowMax) / denom; } @@ -781,7 +780,7 @@ LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_fwd( return torchlean_cuda_buffer_box(out); } -LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd_q( +LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd( b_lean_obj_arg QObj, b_lean_obj_arg KObj, b_lean_obj_arg VObj, b_lean_obj_arg MaskObj, b_lean_obj_arg DOutObj, uint32_t hasMask, uint32_t batch, uint32_t n, uint32_t d, double scaleHost) { @@ -790,12 +789,14 @@ LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd_q( torchlean_cuda_buffer* V = torchlean_cuda_buffer_unbox(VObj); torchlean_cuda_buffer* mask = torchlean_cuda_buffer_unbox(MaskObj); torchlean_cuda_buffer* dOut = torchlean_cuda_buffer_unbox(DOutObj); - flash_attention_check_stub("torchlean_cuda_buffer_flash_attention_bwd_q_stub: size mismatch", + flash_attention_check_stub("torchlean_cuda_buffer_flash_attention_bwd_stub: size mismatch", Q, K, V, mask, dOut, hasMask, batch, n, d); const size_t qkvSz = checked_mul3_size((size_t)batch, (size_t)n, (size_t)d, - "torchlean_cuda_buffer_flash_attention_bwd_k_stub: output size overflow"); + "torchlean_cuda_buffer_flash_attention_bwd_stub: output size overflow"); torchlean_cuda_buffer* dQ = torchlean_cuda_buffer_alloc(qkvSz); + torchlean_cuda_buffer* dK = torchlean_cuda_buffer_alloc(qkvSz); + torchlean_cuda_buffer* dV = torchlean_cuda_buffer_alloc(qkvSz); const float scale = (float)scaleHost; for (size_t b = 0; b < batch; ++b) { for (size_t i = 0; i < n; ++i) { @@ -820,25 +821,6 @@ LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd_q( } } } - return torchlean_cuda_buffer_box(dQ); -} - -LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd_k( - b_lean_obj_arg QObj, b_lean_obj_arg KObj, b_lean_obj_arg VObj, b_lean_obj_arg MaskObj, - b_lean_obj_arg DOutObj, uint32_t hasMask, uint32_t batch, uint32_t n, uint32_t d, - double scaleHost) { - torchlean_cuda_buffer* Q = torchlean_cuda_buffer_unbox(QObj); - torchlean_cuda_buffer* K = torchlean_cuda_buffer_unbox(KObj); - torchlean_cuda_buffer* V = torchlean_cuda_buffer_unbox(VObj); - torchlean_cuda_buffer* mask = torchlean_cuda_buffer_unbox(MaskObj); - torchlean_cuda_buffer* dOut = torchlean_cuda_buffer_unbox(DOutObj); - flash_attention_check_stub("torchlean_cuda_buffer_flash_attention_bwd_k_stub: size mismatch", - Q, K, V, mask, dOut, hasMask, batch, n, d); - const size_t qkvSz = - checked_mul3_size((size_t)batch, (size_t)n, (size_t)d, - "torchlean_cuda_buffer_flash_attention_bwd_v_stub: output size overflow"); - torchlean_cuda_buffer* dK = torchlean_cuda_buffer_alloc(qkvSz); - const float scale = (float)scaleHost; for (size_t b = 0; b < batch; ++b) { for (size_t j = 0; j < n; ++j) { for (size_t k = 0; k < d; ++k) { @@ -862,25 +844,6 @@ LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd_k( } } } - return torchlean_cuda_buffer_box(dK); -} - -LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd_v( - b_lean_obj_arg QObj, b_lean_obj_arg KObj, b_lean_obj_arg VObj, b_lean_obj_arg MaskObj, - b_lean_obj_arg DOutObj, uint32_t hasMask, uint32_t batch, uint32_t n, uint32_t d, - double scaleHost) { - torchlean_cuda_buffer* Q = torchlean_cuda_buffer_unbox(QObj); - torchlean_cuda_buffer* K = torchlean_cuda_buffer_unbox(KObj); - torchlean_cuda_buffer* V = torchlean_cuda_buffer_unbox(VObj); - torchlean_cuda_buffer* mask = torchlean_cuda_buffer_unbox(MaskObj); - torchlean_cuda_buffer* dOut = torchlean_cuda_buffer_unbox(DOutObj); - flash_attention_check_stub("torchlean_cuda_buffer_flash_attention_bwd_v_stub: size mismatch", - Q, K, V, mask, dOut, hasMask, batch, n, d); - const size_t qkvSz = - checked_mul3_size((size_t)batch, (size_t)n, (size_t)d, - "torchlean_cuda_buffer_flash_attention_bwd_v_stub: output size overflow"); - torchlean_cuda_buffer* dV = torchlean_cuda_buffer_alloc(qkvSz); - const float scale = (float)scaleHost; for (size_t b = 0; b < batch; ++b) { for (size_t j = 0; j < n; ++j) { for (size_t dv = 0; dv < d; ++dv) { @@ -896,7 +859,7 @@ LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd_v( } } } - return torchlean_cuda_buffer_box(dV); + return torchlean_cuda_box_three_buffers(dQ, dK, dV); } LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_transpose2d(b_lean_obj_arg BObj, uint32_t rows, diff --git a/csrc/cuda/kernels/torchlean_flash_attention_fwd.cpp b/csrc/cuda/kernels/torchlean_flash_attention_fwd.cpp new file mode 100644 index 00000000..dfc16b4b --- /dev/null +++ b/csrc/cuda/kernels/torchlean_flash_attention_fwd.cpp @@ -0,0 +1,49 @@ +// LibTorch SDPA forward/backward (g++; nvcc cannot parse torch headers). +#include +#include +#include "torchlean_cuda_buffer.h" +#include + +static at::Tensor view3(b_lean_obj_arg o, uint32_t batch, int64_t n, int64_t d) { + return at::from_blob(torchlean_cuda_buffer_unbox(o)->data, {batch, n, d}, at::kFloat).cuda(); +} + +static c10::optional attnMask(b_lean_obj_arg M, uint32_t hasMask, uint32_t batch, + uint32_t n) { + if (!hasMask) return c10::nullopt; + return c10::optional((1 - view3(M, batch, n, n)) * -1000.f); +} + +extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_fwd( + b_lean_obj_arg Q, b_lean_obj_arg K, b_lean_obj_arg V, b_lean_obj_arg M, + uint32_t hasMask, uint32_t batch, uint32_t n, uint32_t d, double scale) { + auto mask = attnMask(M, hasMask, batch, n); + auto y = at::scaled_dot_product_attention(view3(Q, batch, n, d), view3(K, batch, n, d), + view3(V, batch, n, d), mask, 0., false, scale) + .contiguous(); + auto* out = torchlean_cuda_buffer_alloc(y.numel()); + cudaMemcpy(out->data, y.data_ptr(), y.numel() * sizeof(float), cudaMemcpyDeviceToDevice); + return torchlean_cuda_buffer_box(out); +} + +extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_flash_attention_bwd( + b_lean_obj_arg Q, b_lean_obj_arg K, b_lean_obj_arg V, b_lean_obj_arg M, b_lean_obj_arg DOut, + uint32_t hasMask, uint32_t batch, uint32_t n, uint32_t d, double scale) { + at::AutoGradMode enable_grad(true); + auto mask = attnMask(M, hasMask, batch, n); + auto q = view3(Q, batch, n, d).detach().requires_grad_(true); + auto k = view3(K, batch, n, d).detach().requires_grad_(true); + auto v = view3(V, batch, n, d).detach().requires_grad_(true); + auto grad_out = view3(DOut, batch, n, d); + auto out = at::scaled_dot_product_attention(q, k, v, mask, 0., false, scale); + out.backward(grad_out); + const size_t numel = (size_t)batch * (size_t)n * (size_t)d; + at::Tensor grads[3] = {q.grad(), k.grad(), v.grad()}; + torchlean_cuda_buffer* bufs[3]; + for (int i = 0; i < 3; ++i) bufs[i] = torchlean_cuda_buffer_alloc(numel); + for (int i = 0; i < 3; ++i) { + auto t = grads[i].contiguous(); + cudaMemcpy(bufs[i]->data, t.data_ptr(), numel * sizeof(float), cudaMemcpyDeviceToDevice); + } + return torchlean_cuda_box_three_buffers(bufs[0], bufs[1], bufs[2]); +} diff --git a/lakefile.lean b/lakefile.lean index 373892bc..dbca7551 100644 --- a/lakefile.lean +++ b/lakefile.lean @@ -5,6 +5,7 @@ Authors: TorchLean Team -/ import Lake +import Lake.Util.Proc open Lake DSL open System @@ -30,14 +31,22 @@ private def cudaHome : String := | some p => cleanCudaHome p | none => "/usr/local/cuda" +/-- Optional explicit LibTorch root from `-K libtorch_home=...`. -/ +private def libtorchHomeConfig : Option String := + match get_config? libtorch_home with + | some p => + let t := p.trimAscii.toString + if t.isEmpty then none else some t + | none => none + /-- Native link flags selected by the `cuda` Lake option. -/ private def nativeLinkArgs : Array String := if cudaEnabled then + let lt := match libtorchHomeConfig with | some h => h | none => "libtorch" + -- LibTorch is linked into `torchlean_flash_attention_fwd` shared lib (g++), not here. #[ - "-L", s!"{cudaHome}/lib64", - "-lcudart", "-lcublas", "-lcufft", - "-lstdc++", - "-Wl,-rpath," ++ s!"{cudaHome}/lib64" + "-L", s!"{cudaHome}/lib64", "-lcudart", "-lcublas", "-lcufft", + "-Wl,-rpath," ++ s!"{cudaHome}/lib64", "-Wl,-rpath," ++ s!"{lt}/lib" ] else if Platform.isWindows || Platform.isOSX then -- Windows and macOS provide libm via the default C runtime @@ -64,24 +73,6 @@ package TorchLean where ⟨`backward.privateInPublic.warn, false⟩] moreLinkArgs := nativeLinkArgs -@[default_target] -lean_lib NN where - -- `NN:docs` should document the whole maintained Lean surface, including examples and CLI - -- dispatchers. Keep tests out of this library surface; they build through `nn_tests_suite`. - roots := #[ - `NN, - `NN.Examples.Zoo, - `NN.CI.SlowProofs, - `NN.Examples.Models.Runner, - `NN.Verification.CLI - ] - globs := #[ - .one `NN, - .one `NN.Library, - .submodules `NN.Examples, - .submodules `NN.Verification - ] - /-! ## Native backend libraries @@ -96,6 +87,35 @@ private structure NativeBackendLib where cudaSrc : String stubSrc : String +/-- LibTorch root for `-I` / `-L` (must match `resolve_libtorch.sh`). -/ +private def libtorchHome (pkg : Package) : String := + match libtorchHomeConfig with + | some h => h + | none => (pkg.dir / "libtorch").toString + +/-- g++ compile flags for LibTorch C++ sources. -/ +private def libtorchCppCompileArgs (pkg : Package) (lean : LeanInstall) (lt : String) : Array String := + #[ + "-I", lean.includeDir.toString, + "-I", s!"{pkg.dir}/csrc/cuda/common", + "-I", s!"{cudaHome}/include", + "-I", s!"{lt}/include", + "-I", s!"{lt}/include/torch/csrc/api/include", + "-c", "-O2", "-fPIC", "-std=c++17", "-D_GLIBCXX_USE_CXX11_ABI=1" + ] + +/-- g++ link flags for the LibTorch flash-attention forward shared library. -/ +private def libtorchFlashFwdLinkArgs (lt : String) : Array String := + #[ + "-L", s!"{lt}/lib", + "-Wl,--no-as-needed", + "-ltorch", "-ltorch_cpu", "-ltorch_cuda", "-lc10", "-lc10_cuda", + "-L", s!"{cudaHome}/lib64", "-lcudart", + "-lstdc++", + "-Wl,-rpath," ++ s!"{lt}/lib", + "-Wl,-rpath," ++ s!"{cudaHome}/lib64" + ] + /-- Include paths shared by the CUDA implementations and the portable C stubs. -/ private def nativeIncludeArgs (pkg : Package) : Array String := #[ @@ -103,6 +123,57 @@ private def nativeIncludeArgs (pkg : Package) : Array String := "-I", (pkg.dir / "csrc/cuda/conv_pool").toString ] +/-- Resolve LibTorch; caches `.lake/build/libtorch.path`. -/ +private def libtorchResolveJob (pkg : Package) : SpawnM (Job FilePath) := do + let stamp := pkg.buildDir / "libtorch.path" + let resolver := pkg.dir / "scripts" / "setup" / "resolve_libtorch.sh" + let resolverJob ← inputFile resolver false + let args := + match libtorchHomeConfig with + | some home => #[resolver.toString, home] + | none => #[resolver.toString] + buildFileAfterDep stamp resolverJob fun _ => + proc { cmd := "bash", args := args, cwd := some pkg.dir } + +/-- LibTorch SDPA forward as a shared library (g++ links libstdc++ + LibTorch). -/ +private def buildLibtorchFlashAttentionFwdSo (pkg : Package) := do + let lean ← getLeanInstall + let _ ← libtorchResolveJob pkg + let lt := libtorchHome pkg + let cppJob ← inputFile (pkg.dir / "csrc/cuda/kernels/torchlean_flash_attention_fwd.cpp") false + let cppO := pkg.buildDir / "torchlean_flash_attention_fwd.o" + let cppOJob ← buildO cppO cppJob (libtorchCppCompileArgs pkg lean lt) #[] "c++" + let soFile := pkg.buildDir / nameToSharedLib "torchlean_flash_attention_fwd" + cppOJob.mapM fun o => do + let art ← buildArtifactUnlessUpToDate soFile (ext := sharedLibExt) (restore := true) do + compileSharedLib soFile (#[o.toString] ++ libtorchFlashFwdLinkArgs lt) "g++" + return art.path + +target torchlean_flash_attention_fwd_so pkg : FilePath := + if cudaEnabled then + buildLibtorchFlashAttentionFwdSo pkg + else + pure (Job.pure (pkg.buildDir / "torchlean_flash_attention_fwd_skipped")) + +@[default_target] +lean_lib NN where + moreLinkObjs := if cudaEnabled then #[torchlean_flash_attention_fwd_so] else #[] + -- `NN:docs` should document the whole maintained Lean surface, including examples and CLI + -- dispatchers. Keep tests out of this library surface; they build through `nn_tests_suite`. + roots := #[ + `NN, + `NN.Examples.Zoo, + `NN.CI.SlowProofs, + `NN.Examples.Models.Runner, + `NN.Verification.CLI + ] + globs := #[ + .one `NN, + .one `NN.Library, + .submodules `NN.Examples, + .submodules `NN.Verification + ] + /-- Build one native backend library for the current Lake configuration. -/ private def buildNativeBackendLib (pkg : Package) (spec : NativeBackendLib) := do let lean ← getLeanInstall From 8589278c8c08087c0da0595770419eec62334fdb Mon Sep 17 00:00:00 2001 From: wadkisson Date: Tue, 7 Jul 2026 14:17:47 -0700 Subject: [PATCH 3/3] delete duplicate header --- .../conv_pool/torchlean_cuda_conv_pool_common.h | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/csrc/cuda/conv_pool/torchlean_cuda_conv_pool_common.h b/csrc/cuda/conv_pool/torchlean_cuda_conv_pool_common.h index c0369d65..096649f2 100644 --- a/csrc/cuda/conv_pool/torchlean_cuda_conv_pool_common.h +++ b/csrc/cuda/conv_pool/torchlean_cuda_conv_pool_common.h @@ -18,19 +18,6 @@ static inline void checkBufSize(torchlean_cuda_buffer* b, size_t elems, const ch } } -static inline lean_object* torchlean_cuda_box_three_buffers( - torchlean_cuda_buffer* first, - torchlean_cuda_buffer* second, - torchlean_cuda_buffer* third) { - lean_object* tail = lean_alloc_ctor(0, 2, 0); - lean_ctor_set(tail, 0, torchlean_cuda_buffer_box(second)); - lean_ctor_set(tail, 1, torchlean_cuda_buffer_box(third)); - lean_object* out = lean_alloc_ctor(0, 2, 0); - lean_ctor_set(out, 0, torchlean_cuda_buffer_box(first)); - lean_ctor_set(out, 1, tail); - return out; -} - static inline uint32_t outDim(uint32_t in, uint32_t k, uint32_t stride, uint32_t padding) { if (stride == 0) { lean_internal_panic("torchlean_cuda_conv_pool: stride must be > 0");