Swap in PyTorch kernels for MHA

.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__/

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.

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

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`.
 

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

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

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");