{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE NoFieldSelectors #-}

-- | A port of Karpathy's microgpt.py (vendored in microgpt/microgpt.py).
--
-- * test suite (@cabal test@): checks the port against reference values
--   produced by microgpt/reference.py
-- * benchmark (@cabal bench@): same workloads as microgpt/bench.py
-- * executable (@cabal run microgpt@): the full microgpt.py script
--   (shuffle, 1000 training steps, 20 samples)
module Main (main, benchMain, microGPT, parameter, pow, exp', log', relu) where

import Control.DeepSeq (NFData, deepseq, force)
import Control.Exception (evaluate)
import Control.Monad (when)
import Control.Monad.ST (ST, runST, stToIO)
import Control.Monad.State.Strict (State, get, put, runState, state)
import qualified Criterion.Main as Criterion
import Data.Array.Accelerate ((:.) (..))
import qualified Data.Array.Accelerate as A
import qualified Data.Array.Accelerate.LLVM.Native as CPU
import qualified Data.Char as Char
import qualified Data.Foldable as Foldable
import qualified Data.Functor as Functor
import qualified Data.IntMap.Strict as IntMap
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.Maybe as Maybe
import qualified Data.STRef as STRef
import qualified Data.Set as Set
import qualified Data.Traversable as Traversable
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import qualified Data.Vector.Mutable as MVector
import qualified Data.Vector.Unboxed as UVector
import qualified Data.Vector.Unboxed.Mutable as UMVector
import GHC.Generics (Generic)
import GHC.Records (HasField (..))
import qualified Numeric
import System.IO (hFlush, stdout)
import System.Random (StdGen)
import qualified System.Random as Random
import Test.Hspec

main :: IO ()
main = hspec spec

spec :: Spec
spec = do
  describe "Autograd Engine" $ do
    it "computes gradients for f(a, b) = a * b + a" $ do
      let graph = do
            a <- constant 2.0
            b <- constant 3.0
            ab <- mul a b
            add ab a

          (finalNode, tape) = runState graph initialTape
          grads = backward finalNode tape

      -- Forward pass result
      finalNode.node.val `shouldBe` 8.0

      let (NodeIdxMap gMap) = grads
      IntMap.lookup 0 gMap `shouldBe` Just 4.0 -- df/da
      IntMap.lookup 1 gMap `shouldBe` Just 2.0 -- df/db
  describe "MicroGPT Training Step" $ do
    it "computes loss and gradients for a single word" $ do
      let tok = mkTokenizer (Vector.singleton "abcdefghijklmnopqrstuvwxyz")
          cfg = microConfig tok.vocabSize
          w = deterministicWeights cfg

          (loss, grads) = trainingStep cfg w (encode tok "hello")

      loss `shouldSatisfy` (> 0.0)
      -- Verify gradient map is not empty
      let (NodeIdxMap gMap) = grads
      IntMap.size gMap `shouldSatisfy` (> 0)
    it "handles empty documents (BOS -> BOS) and degenerate single-token inputs" $ do
      let tok = mkTokenizer (Vector.singleton "abcdefghijklmnopqrstuvwxyz")
          cfg = microConfig tok.vocabSize
          w = deterministicWeights cfg
          (emptyDocLoss, _) = trainingStep cfg w (encode tok "")
          (singleTokLoss, _) = trainingStep cfg w (Vector.singleton tok.bos)
      emptyDocLoss `shouldSatisfy` (> 0.0)
      singleTokLoss `shouldBe` 0.0
  describe "Equivalence with microgpt.py (values from microgpt/reference.py)" $ do
    let tok = mkTokenizer (Vector.singleton "abcdefghijklmnopqrstuvwxyz")
        cfg = microConfig tok.vocabSize
        w = deterministicWeights cfg
    it "has the same vocabulary and number of parameters" $ do
      tok.vocabSize `shouldBe` 27
      length w `shouldBe` 4192
    it "computes the same loss" $ do
      let (loss, _) = trainingStep cfg w (encode tok "emma")
      loss `shouldBeApprox` 3.506449207715624
    it "computes the same gradients" $ do
      let (_, grads) = trainingStep cfg w (encode tok "emma")
          paramGrads = Vector.fromList (Foldable.toList (parameterGradients grads w))
      Vector.forM_ referenceGradients $ \(i, expected) ->
        (paramGrads Vector.!? i)
          `shouldSatisfy` maybe False (approxEq expected)
    it "follows the same training trajectory (Adam with linear decay)" $ do
      let docs = Vector.map (encode tok) (Vector.fromList ["emma", "olivia", "ava", "isabella", "sophia"])
          (losses, final) = train cfg microHyperparams (initTrainState w) docs
          finalParams = Vector.fromList (Foldable.toList final.params)
      Vector.length losses `shouldBe` 5
      Vector.zipWithM_
        shouldBeApprox
        losses
        (Vector.fromList [3.506449207715624, 3.2498071277112257, 2.7546668273691433, 3.496728826446422, 3.056047453583676])
      fmap (.value) (finalParams Vector.!? 0)
        `shouldSatisfy` maybe False (approxEq (-0.00267390585765643))
  describe "Inference" $ do
    it "samples documents made of known tokens, at most block_size long" $ do
      let tok = mkTokenizer (Vector.singleton "abcdefghijklmnopqrstuvwxyz")
          cfg = microConfig tok.vocabSize
          (sample, _) = generate cfg tok.bos 0.5 (deterministicWeights cfg) (Random.mkStdGen 42)
      Vector.length sample `shouldSatisfy` (<= cfg.blockSize)
      sample `shouldSatisfy` Vector.all (\t -> t >= 0 && t < tok.bos)

-- | (parameter index, gradient) for "emma", from microgpt/reference.py
referenceGradients :: Vector (Int, Double)
referenceGradients =
  Vector.fromList
    [ (64, 0.5013475135105153),
      (432, -0.10409231751964869),
      (700, 0.12256221981304923),
      (1200, -0.003786914738571059),
      (1500, -0.0012771713881593858),
      (1700, 0.015920736906842876),
      (2000, 0.018551333841752494),
      (2500, -0.0005752369491939385),
      (4191, 0.014184402670265658)
    ]

approxEq :: Double -> Double -> Bool
approxEq expected actual = abs (actual - expected) <= 1e-9 * max 1 (abs expected)

shouldBeApprox :: Double -> Double -> Expectation
shouldBeApprox actual expected = actual `shouldSatisfy` approxEq expected

-- | Same file as microgpt.py (makemore's names.txt)
datasetPath :: FilePath
datasetPath = "microgpt/input.txt"

-- | Loads documents like microgpt.py: one per non-blank line, stripped
loadDocs :: FilePath -> IO (Vector String)
loadDocs path = do
  contents <- readFile path
  let strip = List.dropWhileEnd Char.isSpace . dropWhile Char.isSpace
  pure $ Vector.fromList (filter (not . null) (map strip (lines contents)))

-- | Same workloads as microgpt/bench.py: the first X documents (unshuffled),
-- with deterministic weights
benchMain :: IO ()
benchMain = do
  docs <- loadDocs datasetPath
  let tok = mkTokenizer docs
      cfg = microConfig tok.vocabSize
      n = 100
  w <- evaluate (force (deterministicWeights cfg))
  batch <- evaluate (force (Vector.map (encode tok) (Vector.take n docs)))

  Criterion.defaultMain
    [ Criterion.bgroup
        "microGPT"
        [ Criterion.bench ("forward (" <> show n <> " docs)") $
            Criterion.nf (Vector.map (forwardLoss cfg w)) batch,
          Criterion.bench ("forward+backward (" <> show n <> " docs)") $
            Criterion.nf (Vector.map (trainingStep cfg w)) batch,
          Criterion.bench "train (10 steps)" $
            Criterion.nf
              ( \ds ->
                  let (losses, final) = train cfg microHyperparams (initTrainState w) ds
                   in (losses, fmap (.value) final.params)
              )
              batch
        ]
    ]

-- 1. Scalar AD engine

newtype NodeId
  = NodeId Int
  deriving newtype (Eq, Ord, Num, Enum, Show)

initialNodeId :: NodeId
initialNodeId = NodeId 0

nextNodeId :: NodeId -> NodeId
nextNodeId (NodeId x) = NodeId $ x + 1

newtype LocalPartialDerivative a
  = LocalPartialDerivative a
  deriving newtype (Eq, Ord, Num, Show)

newtype NodeIdxMap a
  = NodeIdxMap (IntMap.IntMap a)
  deriving newtype (NFData)

nimEmpty :: NodeIdxMap a
nimEmpty = NodeIdxMap IntMap.empty

nimInsert :: NodeId -> a -> NodeIdxMap a -> NodeIdxMap a
nimInsert (NodeId idx) x (NodeIdxMap m) = NodeIdxMap $ IntMap.insert idx x m

nimLookup :: NodeId -> NodeIdxMap a -> Maybe a
nimLookup (NodeId k) (NodeIdxMap m) = IntMap.lookup k m

data Node a = Node
  { val :: a,
    deps :: Vector (NodeId, LocalPartialDerivative a)
  }

data NodeIdx a = NodeIdx
  { idx :: NodeId,
    node :: Node a
  }

data Tape a = Tape
  { nodes :: NodeIdxMap (Node a),
    nextId :: NodeId
  }

initialTape :: Tape a
initialTape = Tape {nodes = nimEmpty, nextId = initialNodeId}

type Gradients = NodeIdxMap Double

nadd :: (Num a) => NodeIdx a -> NodeIdx a -> Node a
nadd x y =
  Node
    { val = x.node.val + y.node.val,
      deps = Vector.fromListN 2 [(x.idx, LocalPartialDerivative 1), (y.idx, LocalPartialDerivative 1)]
    }

nmul :: (Num a) => NodeIdx a -> NodeIdx a -> Node a
nmul x y =
  Node
    { val = x.node.val * y.node.val,
      deps = Vector.fromListN 2 [(x.idx, nvLpd y), (y.idx, nvLpd x)]
    }
  where
    nvLpd z = LocalPartialDerivative z.node.val

npow :: (Floating a) => NodeIdx a -> a -> Node a
npow x p =
  Node
    { val = x.node.val ** p,
      deps = Vector.singleton (x.idx, LocalPartialDerivative $ p * (x.node.val ** (p - 1)))
    }

nexp :: (Floating a) => NodeIdx a -> Node a
nexp x =
  Node
    { deps = Vector.singleton (x.idx, LocalPartialDerivative val),
      ..
    }
  where
    val = exp x.node.val

nlog :: (Floating a) => NodeIdx a -> Node a
nlog x =
  Node
    { val = log x.node.val,
      deps = Vector.singleton (x.idx, LocalPartialDerivative $ 1 / x.node.val)
    }

nrelu :: (Ord a, Num a) => NodeIdx a -> Node a
nrelu x =
  Node
    { val = max 0 x.node.val,
      deps = Vector.singleton (x.idx, LocalPartialDerivative $ if x.node.val > 0 then 1 else 0)
    }

type GraphBuilder = State (Tape Double)

-- | Automatically handles ID assignment and tape insertion
pushNode :: Node Double -> GraphBuilder (NodeIdx Double)
pushNode node = do
  tape <- get
  let currentId = tape.nextId
      newIdx = NodeIdx {idx = currentId, node = node}
  put $
    tape
      { nodes = nimInsert currentId node tape.nodes,
        nextId = nextNodeId currentId
      }
  pure newIdx

constant :: Double -> GraphBuilder (NodeIdx Double)
constant x = pushNode Node {val = x, deps = Vector.empty}

-- | Registers a learnable parameter. In the graph, it acts as a leaf node.
parameter :: Double -> GraphBuilder (NodeIdx Double)
parameter = constant

add :: NodeIdx Double -> NodeIdx Double -> GraphBuilder (NodeIdx Double)
add x y = pushNode (nadd x y)

mul :: NodeIdx Double -> NodeIdx Double -> GraphBuilder (NodeIdx Double)
mul x y = pushNode (nmul x y)

pow :: NodeIdx Double -> Double -> GraphBuilder (NodeIdx Double)
pow x p = pushNode (npow x p)

exp' :: NodeIdx Double -> GraphBuilder (NodeIdx Double)
exp' x = pushNode (nexp x)

log' :: NodeIdx Double -> GraphBuilder (NodeIdx Double)
log' x = pushNode (nlog x)

relu :: NodeIdx Double -> GraphBuilder (NodeIdx Double)
relu x = pushNode (nrelu x)

-- | Reverse-mode backpropagation accumulating gradients in-place in an unboxed 'ST' buffer.
backward :: NodeIdx Double -> Tape Double -> Gradients
backward finalNode tape =
  let NodeId startId = finalNode.idx
   in if startId < 0
        then nimEmpty
        else runST $ do
          let len = startId + 1
          gradsBuf <- UMVector.replicate len 0.0
          UMVector.unsafeWrite gradsBuf startId 1.0
          let stepNode i
                | i < 0 = pure ()
                | otherwise = do
                    currentGrad <- UMVector.unsafeRead gradsBuf i
                    when (currentGrad /= 0.0) $
                      case nimLookup (NodeId i) tape.nodes of
                        Nothing -> pure ()
                        Just node ->
                          Vector.forM_ node.deps $ \(NodeId depIdx, LocalPartialDerivative lpd) ->
                            when (depIdx >= 0 && depIdx < len) $
                              stAddAt gradsBuf depIdx (currentGrad * lpd)
                    stepNode (i - 1)
          stepNode startId
          frozen <- UVector.unsafeFreeze gradsBuf
          let collect i
                | i >= len = []
                | otherwise =
                    let !g = UVector.unsafeIndex frozen i
                     in if g == 0.0 && i /= startId
                          then collect (i + 1)
                          else (i, g) : collect (i + 1)
          pure (NodeIdxMap (IntMap.fromDistinctAscList (collect 0)))

-- 2. Tensor Primitives (Accelerate & Mutable Unboxed ST)

-- | Stored as [nout][nin], like microgpt's state_dict
type Matrix a = Vector (Vector a)

scalarOf :: (A.Elt e) => e -> A.Scalar e
scalarOf x = A.fromFunction A.Z (const x)

unScalar :: (A.Elt e) => A.Scalar e -> e
unScalar s = A.indexArray s A.Z

uvecToAcc :: UVector.Vector Double -> A.Vector Double
uvecToAcc uvec =
  A.fromFunction (A.Z :. UVector.length uvec) (\(A.Z :. i) -> UVector.unsafeIndex uvec i)

accToUVec :: A.Vector Double -> UVector.Vector Double
accToUVec arr =
  let A.Z :. len = A.arrayShape arr
   in UVector.generate len (\i -> A.indexArray arr (A.Z :. i))

-- | Slices a row-major @rows x cols@ matrix starting at @offset@ from a flat parameter vector.
sliceMatrix ::
  A.Acc (A.Vector Double) ->
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Matrix Double)
sliceMatrix flat offset rows cols =
  A.generate
    (A.I2 rows cols)
    (\(A.I2 r c) -> flat A.! A.I1 (offset + r * cols + c))

-- | Linear forward pass @y = x * w^T@ for @x@ of shape @(m, k)@ and @w@ of shape @(n, k)@.
accLinearFwd ::
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double)
accLinearFwd x w =
  let A.I2 m _k1 = A.shape x
      A.I2 n _k2 = A.shape w
      xRep = A.replicate (A.lift (A.Z :. A.All :. n :. A.All)) x
      wRep = A.replicate (A.lift (A.Z :. m :. A.All :. A.All)) w
   in A.sum (A.zipWith (*) xRep wRep)

-- | Linear backward wrt input @dx = dy * w@ for @dy@ of shape @(m, n)@ and @w@ of shape @(n, k)@.
accLinearBwdX ::
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double)
accLinearBwdX dy w = accLinearFwd dy (A.transpose w)

-- | Linear backward wrt weights @dw = dy^T * x@ for @dy@ of shape @(m, n)@ and @x@ of shape @(m, k)@.
accLinearBwdW ::
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double)
accLinearBwdW dy x = accLinearFwd (A.transpose dy) (A.transpose x)

-- | Matrix-vector product @y = w * x@ for single-token inference.
accMatVec ::
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double)
accMatVec w x =
  let A.I2 nout _nin = A.shape w
      xRep = A.replicate (A.lift (A.Z :. nout :. A.All)) x
   in A.sum (A.zipWith (*) w xRep)

-- | Row-wise RMSNorm returning both normalized rows and @meanSquare + eps@ per row.
accRmsNorm2DWithMs ::
  A.Acc (A.Matrix Double) ->
  (A.Acc (A.Matrix Double), A.Acc (A.Vector Double))
accRmsNorm2DWithMs x =
  let A.I2 _tLen dDim = A.shape x
      invD = 1.0 / A.fromIntegral dDim
      sumSq = A.sum (A.map (\v -> v * v) x)
      msPlusEps = A.map (\s -> s * invD + 1.0e-5) sumSq
      scales = A.map (** (-0.5)) msPlusEps
      scalesRep = A.replicate (A.lift (A.Z :. A.All :. dDim)) scales
   in (A.zipWith (*) x scalesRep, msPlusEps)

-- | Row-wise RMSNorm backward pass given input @x@, saved @ms + eps@, and upstream @dy@.
accRmsNorm2DBwdFromMs ::
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double)
accRmsNorm2DBwdFromMs x msPlusEps dy =
  let A.I2 _tLen dDim = A.shape x
      invD = 1.0 / A.fromIntegral dDim
      scales = A.map (** (-0.5)) msPlusEps
      dotDyX = A.sum (A.zipWith (*) dy x)
      xCoeff =
        A.zipWith
          (\dotVal u -> 2.0 * (dotVal * (-(0.5 * (u ** (-1.5)))) * invD))
          dotDyX
          msPlusEps
      scalesRep = A.replicate (A.lift (A.Z :. A.All :. dDim)) scales
      xCoeffRep = A.replicate (A.lift (A.Z :. A.All :. dDim)) xCoeff
   in A.zipWith (+) (A.zipWith (*) dy scalesRep) (A.zipWith (*) x xCoeffRep)

-- | 1D RMSNorm for single-token inference.
accRmsNorm1D :: A.Acc (A.Vector Double) -> A.Acc (A.Vector Double)
accRmsNorm1D x =
  let A.I1 dDim = A.shape x
      invD = 1.0 / A.fromIntegral dDim
      sumSq = A.sum (A.map (\v -> v * v) x)
      scale = (A.the sumSq * invD + 1.0e-5) ** (-0.5)
   in A.map (* scale) x

-- | Zeroes a slice @[off .. off + len - 1]@ of a mutable unboxed vector in-place.
{-# INLINE stZeroSlice #-}
stZeroSlice :: UMVector.MVector s Double -> Int -> Int -> ST s ()
stZeroSlice buf off len = UMVector.set (UMVector.unsafeSlice off len buf) 0.0

-- | Strictly adds @delta@ in-place at @idx@ without allocating a modifier closure or thunk.
{-# INLINE stAddAt #-}
stAddAt :: UMVector.MVector s Double -> Int -> Double -> ST s ()
stAddAt buf idx !delta = do
  old <- UMVector.unsafeRead buf idx
  let !new = old + delta
  UMVector.unsafeWrite buf idx new

-- | In-place row-wise RMSNorm in 'ST', saving per-row normalization scales.
{-# INLINE stRmsNormFwd #-}
stRmsNormFwd ::
  Int ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stRmsNormFwd dDim tLen src srcOff dst dstOff scaleBuf scaleOff =
  let !invD = 1.0 / fromIntegral dDim
      loopP p
        | p >= tLen = pure ()
        | otherwise = do
            let !rowSrc = srcOff + p * dDim
                !rowDst = dstOff + p * dDim
                sumSq d !acc
                  | d >= dDim = pure acc
                  | otherwise = do
                      v <- UMVector.unsafeRead src (rowSrc + d)
                      sumSq (d + 1) (acc + v * v)
            sq <- sumSq 0 0.0
            let !scale = 1.0 / sqrt (sq * invD + 1.0e-5)
            UMVector.unsafeWrite scaleBuf (scaleOff + p) scale
            let writeNorm d
                  | d >= dDim = pure ()
                  | otherwise = do
                      v <- UMVector.unsafeRead src (rowSrc + d)
                      UMVector.unsafeWrite dst (rowDst + d) (v * scale)
                      writeNorm (d + 1)
            writeNorm 0
            loopP (p + 1)
   in loopP 0

-- | In-place 1D RMSNorm from @src@ into @dst@ for single-token inference.
{-# INLINE stRmsNorm1D #-}
stRmsNorm1D :: Int -> UMVector.MVector s Double -> UMVector.MVector s Double -> ST s ()
stRmsNorm1D dDim src dst = do
  let !invD = 1.0 / fromIntegral dDim
      sumSq d !acc
        | d >= dDim = pure acc
        | otherwise = do
            v <- UMVector.unsafeRead src d
            sumSq (d + 1) (acc + v * v)
  sq <- sumSq 0 0.0
  let !scale = 1.0 / sqrt (sq * invD + 1.0e-5)
      writeNorm d
        | d >= dDim = pure ()
        | otherwise = do
            v <- UMVector.unsafeRead src d
            UMVector.unsafeWrite dst d (v * scale)
            writeNorm (d + 1)
  writeNorm 0

-- | Accumulates RMSNorm backward gradients directly (@+=@) into @dxDst@ to fuse residual connections.
{-# INLINE stRmsNormBwdAdd #-}
stRmsNormBwdAdd ::
  Int ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stRmsNormBwdAdd dDim tLen xBuf xOff scaleBuf scaleOff dyBuf dyOff dxDst dxOff =
  let !invD = 1.0 / fromIntegral dDim
      loopP p
        | p >= tLen = pure ()
        | otherwise = do
            let !rowX = xBufOff p xOff
                !rowDy = xBufOff p dyOff
                !rowDx = xBufOff p dxOff
            scale <- UMVector.unsafeRead scaleBuf (scaleOff + p)
            let dotRow d !acc
                  | d >= dDim = pure acc
                  | otherwise = do
                      dy <- UMVector.unsafeRead dyBuf (rowDy + d)
                      xv <- UMVector.unsafeRead xBuf (rowX + d)
                      dotRow (d + 1) (acc + dy * xv)
            dotVal <- dotRow 0 0.0
            let !xCoeff = -(dotVal * (scale * scale * scale) * invD)
                accumRow d
                  | d >= dDim = pure ()
                  | otherwise = do
                      dy <- UMVector.unsafeRead dyBuf (rowDy + d)
                      xv <- UMVector.unsafeRead xBuf (rowX + d)
                      stAddAt dxDst (rowDx + d) (scale * dy + xCoeff * xv)
                      accumRow (d + 1)
            accumRow 0
            loopP (p + 1)
      xBufOff p base = base + p * dDim
   in loopP 0

-- | In-place linear forward @y = x * w^T@ over @tLen@ sequence positions.
{-# INLINE stLinearFwd #-}
stLinearFwd ::
  Int ->
  Int ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stLinearFwd tLen nin nout xBuf xOff wBuf wOff yBuf yOff =
  let loopP p
        | p >= tLen = pure ()
        | otherwise = do
            let !rowX = xOff + p * nin
                !rowY = yOff + p * nout
                loopR r
                  | r >= nout = pure ()
                  | otherwise = do
                      let !rowW = wOff + r * nin
                          dotC c !acc
                            | c >= nin = pure acc
                            | otherwise = do
                                xv <- UMVector.unsafeRead xBuf (rowX + c)
                                wv <- UMVector.unsafeRead wBuf (rowW + c)
                                dotC (c + 1) (acc + xv * wv)
                      val <- dotC 0 0.0
                      UMVector.unsafeWrite yBuf (rowY + r) val
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

-- | Fused linear forward + ReLU activation @y = max(0, x * w^T)@ in 'ST'.
{-# INLINE stLinearFwdRelu #-}
stLinearFwdRelu ::
  Int ->
  Int ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stLinearFwdRelu tLen nin nout xBuf xOff wBuf wOff yBuf yOff =
  let loopP p
        | p >= tLen = pure ()
        | otherwise = do
            let !rowX = xOff + p * nin
                !rowY = yOff + p * nout
                loopR r
                  | r >= nout = pure ()
                  | otherwise = do
                      let !rowW = wOff + r * nin
                          dotC c !acc
                            | c >= nin = pure acc
                            | otherwise = do
                                xv <- UMVector.unsafeRead xBuf (rowX + c)
                                wv <- UMVector.unsafeRead wBuf (rowW + c)
                                dotC (c + 1) (acc + xv * wv)
                      val <- dotC 0 0.0
                      UMVector.unsafeWrite yBuf (rowY + r) (max 0.0 val)
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

-- | In-place linear forward @y = x * w^T@ reading immutable unboxed weights directly.
{-# INLINE stLinearFwdImmW #-}
stLinearFwdImmW ::
  Int ->
  Int ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UVector.Vector Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stLinearFwdImmW tLen nin nout xBuf xOff wVec wOff yBuf yOff =
  let loopP p
        | p >= tLen = pure ()
        | otherwise = do
            let !rowX = xOff + p * nin
                !rowY = yOff + p * nout
                loopR r
                  | r >= nout = pure ()
                  | otherwise = do
                      let !rowW = wOff + r * nin
                          dotC c !acc
                            | c >= nin = pure acc
                            | otherwise = do
                                xv <- UMVector.unsafeRead xBuf (rowX + c)
                                let !wv = UVector.unsafeIndex wVec (rowW + c)
                                dotC (c + 1) (acc + xv * wv)
                      val <- dotC 0 0.0
                      UMVector.unsafeWrite yBuf (rowY + r) val
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

-- | Fused linear forward + ReLU reading immutable unboxed weights directly.
{-# INLINE stLinearFwdReluImmW #-}
stLinearFwdReluImmW ::
  Int ->
  Int ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UVector.Vector Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stLinearFwdReluImmW tLen nin nout xBuf xOff wVec wOff yBuf yOff =
  let loopP p
        | p >= tLen = pure ()
        | otherwise = do
            let !rowX = xOff + p * nin
                !rowY = yOff + p * nout
                loopR r
                  | r >= nout = pure ()
                  | otherwise = do
                      let !rowW = wOff + r * nin
                          dotC c !acc
                            | c >= nin = pure acc
                            | otherwise = do
                                xv <- UMVector.unsafeRead xBuf (rowX + c)
                                let !wv = UVector.unsafeIndex wVec (rowW + c)
                                dotC (c + 1) (acc + xv * wv)
                      val <- dotC 0 0.0
                      UMVector.unsafeWrite yBuf (rowY + r) (max 0.0 val)
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

-- | In-place linear backward accumulating both @dw += dy^T * x@ and @dx += dy * w@ in unit-stride loops.
{-# INLINE stLinearBwdAccum #-}
stLinearBwdAccum ::
  Int ->
  Int ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stLinearBwdAccum tLen nout nin dyBuf dyOff xBuf xOff wBuf wOff dwBuf dxBuf dxOff =
  let loopP p
        | p >= tLen = pure ()
        | otherwise = do
            let !rowDy = dyOff + p * nout
                !rowX = xOff + p * nin
                !rowDx = dxOff + p * nin
                loopR r
                  | r >= nout = pure ()
                  | otherwise = do
                      g <- UMVector.unsafeRead dyBuf (rowDy + r)
                      when (g /= 0.0) $ do
                        let !rowW = wOff + r * nin
                            loopC c
                              | c >= nin = pure ()
                              | otherwise = do
                                  xv <- UMVector.unsafeRead xBuf (rowX + c)
                                  wv <- UMVector.unsafeRead wBuf (rowW + c)
                                  stAddAt dwBuf (rowW + c) (g * xv)
                                  stAddAt dxBuf (rowDx + c) (g * wv)
                                  loopC (c + 1)
                        loopC 0
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

{-# INLINE stLinearBwd #-}
stLinearBwd ::
  Int ->
  Int ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  Int ->
  UMVector.MVector s Double ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stLinearBwd tLen nout nin dyBuf dyOff xBuf xOff wBuf wOff dwBuf dxBuf dxOff = do
  stZeroSlice dxBuf dxOff (tLen * nin)
  stLinearBwdAccum tLen nout nin dyBuf dyOff xBuf xOff wBuf wOff dwBuf dxBuf dxOff

-- | Matrix-vector product @y = w * x@ from immutable unboxed @w@ into mutable @y@ for inference.
{-# INLINE stMatVecImmW #-}
stMatVecImmW ::
  Int ->
  Int ->
  UVector.Vector Double ->
  Int ->
  UMVector.MVector s Double ->
  UMVector.MVector s Double ->
  Int ->
  ST s ()
stMatVecImmW nout nin wVec wOff xBuf yBuf yOff =
  let loopR r
        | r >= nout = pure ()
        | otherwise = do
            let !rowW = wOff + r * nin
                dotC c !acc
                  | c >= nin = pure acc
                  | otherwise = do
                      xv <- UMVector.unsafeRead xBuf c
                      let !wv = UVector.unsafeIndex wVec (rowW + c)
                      dotC (c + 1) (acc + xv * wv)
            val <- dotC 0 0.0
            UMVector.unsafeWrite yBuf (yOff + r) val
            loopR (r + 1)
   in loopR 0

-- 3. Model: GPT-2 with rmsnorm, no biases and ReLU (like microgpt)

data Config = Config
  { vocabSize :: Int,
    nLayer :: Int,
    nEmbd :: Int,
    blockSize :: Int,
    nHead :: Int
  }

microConfig :: Int -> Config
microConfig vocab =
  Config
    { vocabSize = vocab,
      nLayer = 1,
      nEmbd = 16,
      blockSize = 16,
      nHead = 4
    }

headDim :: Config -> Int
headDim cfg = cfg.nEmbd `div` cfg.nHead

expectedParamCount :: Config -> Int
expectedParamCount cfg =
  (2 * cfg.vocabSize + cfg.blockSize) * cfg.nEmbd
    + cfg.nLayer * 12 * cfg.nEmbd * cfg.nEmbd

isValidModelShape :: Config -> Int -> Bool
isValidModelShape cfg nParams =
  cfg.vocabSize > 0
    && cfg.blockSize > 0
    && cfg.nEmbd > 0
    && cfg.nHead > 0
    && cfg.nLayer >= 0
    && cfg.nEmbd == cfg.nHead * headDim cfg
    && nParams >= expectedParamCount cfg

-- | Selects multicore LLVM Accelerate kernels for large compute-bound embeddings
-- and low-latency in-place 'ST' kernels for small cache-resident shapes.
shouldUseAccelerate :: Config -> Bool
shouldUseAccelerate cfg = cfg.nEmbd >= 128

-- | Parameters of a transformer layer, matrices are [nout][nin]
data Layer a = Layer
  { attnWq :: Matrix a,
    attnWk :: Matrix a,
    attnWv :: Matrix a,
    attnWo :: Matrix a,
    mlpFc1 :: Matrix a,
    mlpFc2 :: Matrix a
  }
  deriving stock (Generic, Functor, Foldable, Traversable)
  deriving anyclass (NFData)

-- | The root microGPT structure.
-- The 'Foldable' order matches microgpt's flattened @params@ list.
data Weights a = Weights
  { wte :: Matrix a,
    wpe :: Matrix a,
    lmHead :: Matrix a,
    layers :: Vector (Layer a)
  }
  deriving stock (Generic, Functor, Foldable, Traversable)
  deriving anyclass (NFData)

-- | Shape of microgpt's state_dict
weightsShape :: Config -> Weights ()
weightsShape cfg =
  Weights
    { wte = matrix cfg.vocabSize cfg.nEmbd,
      wpe = matrix cfg.blockSize cfg.nEmbd,
      lmHead = matrix cfg.vocabSize cfg.nEmbd,
      layers = Vector.replicate cfg.nLayer layer
    }
  where
    matrix nout nin = Vector.replicate nout (Vector.replicate nin ())
    layer =
      Layer
        { attnWq = matrix cfg.nEmbd cfg.nEmbd,
          attnWk = matrix cfg.nEmbd cfg.nEmbd,
          attnWv = matrix cfg.nEmbd cfg.nEmbd,
          attnWo = matrix cfg.nEmbd cfg.nEmbd,
          mlpFc1 = matrix (4 * cfg.nEmbd) cfg.nEmbd,
          mlpFc2 = matrix cfg.nEmbd (4 * cfg.nEmbd)
        }

-- | Gaussian initialization (std 0.08), like microgpt
randomWeights :: Config -> StdGen -> (Weights Double, StdGen)
randomWeights cfg = runState (traverse (const (gauss 0 0.08)) (weightsShape cfg))

-- | Box-Muller Gaussian sampler
gauss :: Double -> Double -> State StdGen Double
gauss mu sigma = do
  u1 <- state Random.random
  u2 <- state Random.random
  pure $ mu + sigma * cos (2 * pi * u1) * sqrt (-(2 * log (max 1e-300 (1 - u2))))

-- | Deterministic initialization shared with microgpt/harness.py to compare both implementations
deterministicWeights :: Config -> Weights Double
deterministicWeights cfg =
  snd $ Traversable.mapAccumL (\i () -> (i + 1, 0.08 * sin (fromIntegral i * 0.37 + 0.11))) (0 :: Int) (weightsShape cfg)

matrixElemCount :: Matrix a -> Int
matrixElemCount = Vector.foldl' (\acc row -> acc + Vector.length row) 0

layerElemCount :: Layer a -> Int
layerElemCount layer =
  matrixElemCount layer.attnWq
    + matrixElemCount layer.attnWk
    + matrixElemCount layer.attnWv
    + matrixElemCount layer.attnWo
    + matrixElemCount layer.mlpFc1
    + matrixElemCount layer.mlpFc2

weightsElemCount :: Weights a -> Int
weightsElemCount w =
  matrixElemCount w.wte
    + matrixElemCount w.wpe
    + matrixElemCount w.lmHead
    + Vector.foldl' (\acc l -> acc + layerElemCount l) 0 w.layers

-- | Flattens 'Weights' directly into a single unboxed vector via in-place 'ST' writes.
{-# NOINLINE flattenWeightsUVec #-}
flattenWeightsUVec :: Weights Double -> UVector.Vector Double
flattenWeightsUVec w = runST $ do
  let total = weightsElemCount w
  buf <- UMVector.new total
  let writeMatrix !off =
        Vector.foldM
          ( \cur row -> do
              let rLen = Vector.length row
                  writeRow j
                    | j >= rLen = pure (cur + rLen)
                    | otherwise = do
                        UMVector.unsafeWrite buf (cur + j) (Vector.unsafeIndex row j)
                        writeRow (j + 1)
              writeRow 0
          )
          off
      writeLayer !off layer = do
        off1 <- writeMatrix off layer.attnWq
        off2 <- writeMatrix off1 layer.attnWk
        off3 <- writeMatrix off2 layer.attnWv
        off4 <- writeMatrix off3 layer.attnWo
        off5 <- writeMatrix off4 layer.mlpFc1
        writeMatrix off5 layer.mlpFc2
  off1 <- writeMatrix 0 w.wte
  off2 <- writeMatrix off1 w.wpe
  off3 <- writeMatrix off2 w.lmHead
  Vector.foldM_ writeLayer off3 w.layers
  UVector.unsafeFreeze buf

-- | Packs model configuration dimensions into a 6-element Accelerate vector.
packConfig :: Config -> A.Vector Int
packConfig cfg =
  A.fromFunction
    (A.Z :. 6)
    ( \case
        A.Z :. 0 -> cfg.vocabSize
        A.Z :. 1 -> cfg.nLayer
        A.Z :. 2 -> cfg.nEmbd
        A.Z :. 3 -> cfg.blockSize
        A.Z :. 4 -> cfg.nHead
        _ -> headDim cfg
    )

unpackAccConfig ::
  A.Acc (A.Vector Int) ->
  (A.Exp Int, A.Exp Int, A.Exp Int, A.Exp Int, A.Exp Int, A.Exp Int)
unpackAccConfig cfgVec =
  ( cfgVec A.! A.I1 0,
    cfgVec A.! A.I1 1,
    cfgVec A.! A.I1 2,
    cfgVec A.! A.I1 3,
    cfgVec A.! A.I1 4,
    cfgVec A.! A.I1 5
  )

-- | Packs 'Weights' into a flat Accelerate vector matching 'Foldable' parameter order.
{-# NOINLINE packWeights #-}
packWeights :: Weights Double -> A.Vector Double
packWeights = uvecToAcc . flattenWeightsUVec

sliceLayerMatrices ::
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Vector Double) ->
  A.Exp Int ->
  ( A.Acc (A.Matrix Double),
    A.Acc (A.Matrix Double),
    A.Acc (A.Matrix Double),
    A.Acc (A.Matrix Double)
  )
sliceLayerMatrices vSize bSize dDim flatW lIdx =
  let dd = dDim * dDim
      base = (2 * vSize + bSize) * dDim + lIdx * 12 * dd
      mlpDim = 4 * dDim
   in ( sliceMatrix flatW base (3 * dDim) dDim,
        sliceMatrix flatW (base + 3 * dd) dDim dDim,
        sliceMatrix flatW (base + 4 * dd) mlpDim dDim,
        sliceMatrix flatW (base + 8 * dd) dDim mlpDim
      )

-- | Online log-sum-exp combiner computing @(maxVal, sum(exp(x - maxVal)))@ in a single reduction pass.
combineMaxSumExp ::
  A.Exp (Double, Double) ->
  A.Exp (Double, Double) ->
  A.Exp (Double, Double)
combineMaxSumExp (A.T2 m1 s1) (A.T2 m2 s2) =
  let m = A.max m1 m2
   in A.T2 m (s1 * exp (m1 - m) + s2 * exp (m2 - m))

-- | Extracts @(H, T, hd)@ head slices starting at @colOffset@ from fused @qkv@ of shape @(T, 3 * D)@.
splitHeadsAt ::
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Array A.DIM3 Double)
splitHeadsAt colOffset nHead hd mat =
  let A.I2 tLen _cols = A.shape mat
   in A.generate (A.I3 nHead tLen hd) $ \(A.I3 h p c) ->
        mat A.! A.I2 p (colOffset + h * hd + c)

-- | Reshapes @(T, H * hd)@ into @(H, T, hd)@.
splitHeads ::
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Array A.DIM3 Double)
splitHeads = splitHeadsAt 0

-- | Reshapes @(H, T, hd)@ back into @(T, H * hd)@.
mergeHeads ::
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Matrix Double)
mergeHeads nHead hd arr =
  let A.I3 _h tLen _hd = A.shape arr
   in A.generate (A.I2 tLen (nHead * hd)) $ \(A.I2 p d) ->
        arr A.! A.I3 (d `A.quot` hd) p (d `A.rem` hd)

-- | Merges @(dQH, dKH, dVH)@ of shape @(H, T, hd)@ into a single @(T, 3 * D)@ gradient matrix.
mergeQKVHeads ::
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Matrix Double)
mergeQKVHeads dDim hd dQH dKH dVH =
  let A.I3 _h tLen _hd = A.shape dQH
   in A.generate (A.I2 tLen (3 * dDim)) $ \(A.I2 p col) ->
        let part = col `A.quot` dDim
            remD = col `A.rem` dDim
            h = remD `A.quot` hd
            c = remD `A.rem` hd
         in (part A.== 0)
              A.? ( dQH A.! A.I3 h p c,
                    (part A.== 1)
                      A.? (dKH A.! A.I3 h p c, dVH A.! A.I3 h p c)
                  )

-- | Batched contraction @out[h, p, s] = sum_c a[h, p, c] * b[h, s, c]@.
batchedScoreContract ::
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double)
batchedScoreContract a b =
  let A.I3 _h1 pLen _c1 = A.shape a
      A.I3 _h2 sLen _c2 = A.shape b
      aRep = A.replicate (A.lift (A.Z :. A.All :. A.All :. sLen :. A.All)) a
      bRep = A.replicate (A.lift (A.Z :. A.All :. pLen :. A.All :. A.All)) b
   in A.sum (A.zipWith (*) aRep bRep)

-- | Batched contraction @out[h, p, c] = sum_s weights[h, p, s] * vals[h, s, c]@.
batchedWeightedSumS ::
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double)
batchedWeightedSumS weights vals =
  let A.I3 hCnt pLen sLen = A.shape weights
      A.I3 _h2 _sLen2 cDim = A.shape vals
      wRep = A.replicate (A.lift (A.Z :. A.All :. A.All :. cDim :. A.All)) weights
      valsPerm = A.backpermute (A.I3 hCnt cDim sLen) (\(A.I3 h c s) -> A.I3 h s c) vals
      vRep = A.replicate (A.lift (A.Z :. A.All :. pLen :. A.All :. A.All)) valsPerm
   in A.sum (A.zipWith (*) wRep vRep)

-- | Batched contraction @out[h, s, c] = sum_p weights[h, p, s] * grads[h, p, c]@.
batchedWeightedSumP ::
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double)
batchedWeightedSumP weights grads =
  let A.I3 hCnt pLen sLen = A.shape weights
      weightsT = A.backpermute (A.I3 hCnt sLen pLen) (\(A.I3 h s p) -> A.I3 h p s) weights
   in batchedWeightedSumS weightsT grads

-- | Packs layer parameter gradients into a flat @12 * D * D@ vector without intermediate concatenations.
packLayerGrads ::
  A.Exp Int ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Vector Double)
packLayerGrads dDim dWQKV dWo dFc1 dFc2 =
  let dd = dDim * dDim
      total = 12 * dd
      fWQKV = A.flatten dWQKV
      fWo = A.flatten dWo
      fFc1 = A.flatten dFc1
      fFc2 = A.flatten dFc2
   in A.generate (A.I1 total) $ \(A.I1 idx) ->
        (idx A.< 3 * dd)
          A.? ( fWQKV A.! A.I1 (A.min (3 * dd - 1) idx),
                (idx A.< 4 * dd)
                  A.? ( fWo A.! A.I1 (A.max 0 (A.min (dd - 1) (idx - 3 * dd))),
                        (idx A.< 8 * dd)
                          A.? ( fFc1 A.! A.I1 (A.max 0 (A.min (4 * dd - 1) (idx - 4 * dd))),
                                fFc2 A.! A.I1 (A.max 0 (A.min (4 * dd - 1) (idx - 8 * dd)))
                              )
                      )
              )

-- | Packs embedding, LM head, and layer gradients into a single flat parameter gradient vector.
packAllGrads ::
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double)
packAllGrads vSize bSize dDim dWte dWpe dLmHead dLayers =
  let nWte = vSize * dDim
      nWpe = bSize * dDim
      nLm = vSize * dDim
      offWpe = nWte
      offLm = nWte + nWpe
      offLayers = offLm + nLm
      A.I1 nLayers = A.shape dLayers
      total = offLayers + nLayers
      fWte = A.flatten dWte
      fWpe = A.flatten dWpe
      fLm = A.flatten dLmHead
   in A.generate (A.I1 total) $ \(A.I1 idx) ->
        (idx A.< offWpe)
          A.? ( fWte A.! A.I1 (A.min (nWte - 1) idx),
                (idx A.< offLm)
                  A.? ( fWpe A.! A.I1 (A.max 0 (A.min (nWpe - 1) (idx - offWpe))),
                        (idx A.< offLayers)
                          A.? ( fLm A.! A.I1 (A.max 0 (A.min (nLm - 1) (idx - offLm))),
                                (nLayers A.> 0)
                                  A.? (dLayers A.! A.I1 (A.max 0 (A.min (nLayers - 1) (idx - offLayers))), 0.0)
                              )
                      )
              )

-- | Saved forward activations of one transformer layer for analytical backpropagation.
data LayerCache = LayerCache
  { xIn :: A.Acc (A.Matrix Double),
    ms1 :: A.Acc (A.Vector Double),
    xNorm1 :: A.Acc (A.Matrix Double),
    wQKV :: A.Acc (A.Matrix Double),
    wo :: A.Acc (A.Matrix Double),
    fc1 :: A.Acc (A.Matrix Double),
    fc2 :: A.Acc (A.Matrix Double),
    qH :: A.Acc (A.Array A.DIM3 Double),
    kH :: A.Acc (A.Array A.DIM3 Double),
    vH :: A.Acc (A.Array A.DIM3 Double),
    attnW :: A.Acc (A.Array A.DIM3 Double),
    attnConcat :: A.Acc (A.Matrix Double),
    xMid :: A.Acc (A.Matrix Double),
    ms2 :: A.Acc (A.Vector Double),
    xNorm2 :: A.Acc (A.Matrix Double),
    preRelu :: A.Acc (A.Matrix Double),
    hidden :: A.Acc (A.Matrix Double)
  }

-- | Parallel causal transformer layer forward pass across all sequence positions simultaneously.
accLayerFwdWithCache ::
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Vector Double) ->
  A.Exp Int ->
  A.Acc (A.Matrix Double) ->
  (A.Acc (A.Matrix Double), LayerCache)
accLayerFwdWithCache vSize bSize dDim nHead hd flatW lIdx x =
  let A.I2 tLen _d = A.shape x
      invScale = sqrt (A.fromIntegral hd) ** (-1.0)
      (wQKV, wo, fc1, fc2) = sliceLayerMatrices vSize bSize dDim flatW lIdx
      (xNorm1, ms1) = accRmsNorm2DWithMs x
      qkv = accLinearFwd xNorm1 wQKV
      qH = splitHeadsAt 0 nHead hd qkv
      kH = splitHeadsAt dDim nHead hd qkv
      vH = splitHeadsAt (2 * dDim) nHead hd qkv
      rawScores = A.map (* invScale) (batchedScoreContract qH kH)
      lseInit =
        A.generate (A.I3 nHead tLen tLen) $ \(A.I3 h p s) ->
          (s A.<= p) A.? (A.T2 (rawScores A.! A.I3 h p s) 1.0, A.T2 (-1.0e308) 0.0)
      lse = A.fold1 combineMaxSumExp lseInit
      attnW =
        A.generate (A.I3 nHead tLen tLen) $ \(A.I3 h p s) ->
          let A.T2 m sumE = lse A.! A.I2 h p
           in (s A.<= p) A.? (exp (rawScores A.! A.I3 h p s - m) * (sumE ** (-1.0)), 0.0)
      headOut = batchedWeightedSumS attnW vH
      attnConcat = mergeHeads nHead hd headOut
      attnOut = accLinearFwd attnConcat wo
      xMid = A.zipWith (+) attnOut x
      (xNorm2, ms2) = accRmsNorm2DWithMs xMid
      preRelu = accLinearFwd xNorm2 fc1
      hidden = A.map (\z -> (z A.> 0.0) A.? (z, 0.0)) preRelu
      mlpOut = accLinearFwd hidden fc2
      xOut = A.zipWith (+) mlpOut xMid
      cache =
        LayerCache
          { xIn = x,
            ms1 = ms1,
            xNorm1 = xNorm1,
            wQKV = wQKV,
            wo = wo,
            fc1 = fc1,
            fc2 = fc2,
            qH = qH,
            kH = kH,
            vH = vH,
            attnW = attnW,
            attnConcat = attnConcat,
            xMid = xMid,
            ms2 = ms2,
            xNorm2 = xNorm2,
            preRelu = preRelu,
            hidden = hidden
          }
   in (xOut, cache)

-- | Analytical backward pass through a transformer layer using its saved forward cache.
accLayerBwdFromCache ::
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  LayerCache ->
  A.Acc (A.Matrix Double) ->
  (A.Acc (A.Matrix Double), A.Acc (A.Vector Double))
accLayerBwdFromCache dDim nHead hd cache dXOut =
  let A.I2 tLen _d = A.shape dXOut
      invScale = sqrt (A.fromIntegral hd) ** (-1.0)
      dFc2 = accLinearBwdW dXOut cache.hidden
      dHidden = accLinearBwdX dXOut cache.fc2
      dPreRelu = A.zipWith (\pre dh -> (pre A.> 0.0) A.? (dh, 0.0)) cache.preRelu dHidden
      dFc1 = accLinearBwdW dPreRelu cache.xNorm2
      dXNorm2 = accLinearBwdX dPreRelu cache.fc1
      dXMidFromMlp = accRmsNorm2DBwdFromMs cache.xMid cache.ms2 dXNorm2
      dXMid = A.zipWith (+) dXOut dXMidFromMlp
      dWo = accLinearBwdW dXMid cache.attnConcat
      dAttnConcat = accLinearBwdX dXMid cache.wo
      dHeadOut = splitHeads nHead hd dAttnConcat
      dAttnW = batchedScoreContract dHeadOut cache.vH
      dVH = batchedWeightedSumP cache.attnW dHeadOut
      wDotDw = A.sum (A.zipWith (*) cache.attnW dAttnW)
      wDotDwRep = A.replicate (A.lift (A.Z :. A.All :. A.All :. tLen)) wDotDw
      ds = A.zipWith3 (\w dw dotVal -> invScale * w * (dw - dotVal)) cache.attnW dAttnW wDotDwRep
      dQH = batchedWeightedSumS ds cache.kH
      dKH = batchedWeightedSumP ds cache.qH
      dQKV = mergeQKVHeads dDim hd dQH dKH dVH
      dWQKV = accLinearBwdW dQKV cache.xNorm1
      dXNorm1 = accLinearBwdX dQKV cache.wQKV
      dXFromAttn = accRmsNorm2DBwdFromMs cache.xIn cache.ms1 dXNorm1
      dXIn = A.zipWith (+) dXMid dXFromAttn
      dLayerFlat = packLayerGrads dDim dWQKV dWo dFc1 dFc2
   in (dXIn, dLayerFlat)

accEmbedFwd ::
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Matrix Double)
accEmbedFwd vSize bSize dDim wte wpe inTokens =
  let A.I1 tLen = A.shape inTokens
   in A.generate (A.I2 tLen dDim) $ \(A.I2 p d) ->
        let tok = inTokens A.! A.I1 p
            safeTok = A.max 0 (A.min (vSize - 1) tok)
            safePos = A.max 0 (A.min (bSize - 1) p)
            tokVal = (0 A.<= tok A.&& tok A.< vSize) A.? (wte A.! A.I2 safeTok d, 0.0)
            posVal = (0 A.<= p A.&& p A.< bSize) A.? (wpe A.! A.I2 safePos d, 0.0)
         in tokVal + posVal

accEmbedBwd ::
  A.Exp Int ->
  A.Exp Int ->
  A.Exp Int ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Matrix Double) ->
  (A.Acc (A.Matrix Double), A.Acc (A.Matrix Double))
accEmbedBwd vSize bSize dDim inTokens dX0 =
  let A.I1 tLen = A.shape inTokens
      dWteTerms =
        A.generate (A.I3 vSize dDim tLen) $ \(A.I3 v d p) ->
          (inTokens A.! A.I1 p A.== v) A.? (dX0 A.! A.I2 p d, 0.0)
      dWte = A.sum dWteTerms
      dWpe =
        A.generate (A.I2 bSize dDim) $ \(A.I2 r d) ->
          let safeR = A.max 0 (A.min (tLen - 1) r)
           in (r A.< tLen) A.? (dX0 A.! A.I2 safeR d, 0.0)
   in (dWte, dWpe)

accLossOnly ::
  A.Exp Int ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Scalar Double)
accLossOnly vSize logits targetTokens =
  let A.I2 tLen _vDim = A.shape logits
      invT = 1.0 / A.fromIntegral tLen
      lse = A.fold1 combineMaxSumExp (A.map (`A.T2` 1.0) logits)
      stepLosses =
        A.generate (A.I1 tLen) $ \(A.I1 p) ->
          let tgt = targetTokens A.! A.I1 p
              safeTgt = A.max 0 (A.min (vSize - 1) tgt)
              valid = 0 A.<= tgt A.&& tgt A.< vSize
              A.T2 m sumE = lse A.! A.I1 p
              prob = exp (logits A.! A.I2 p safeTgt - m) * (sumE ** (-1.0))
           in valid A.? (-log prob, 0.0)
   in A.map (* invT) (A.sum stepLosses)

accLossAndGradLogits ::
  A.Exp Int ->
  A.Acc (A.Matrix Double) ->
  A.Acc (A.Vector Int) ->
  (A.Acc (A.Scalar Double), A.Acc (A.Matrix Double))
accLossAndGradLogits vSize logits targetTokens =
  let A.I2 tLen _vDim = A.shape logits
      invT = 1.0 / A.fromIntegral tLen
      lse = A.fold1 combineMaxSumExp (A.map (`A.T2` 1.0) logits)
      stepLosses =
        A.generate (A.I1 tLen) $ \(A.I1 p) ->
          let tgt = targetTokens A.! A.I1 p
              safeTgt = A.max 0 (A.min (vSize - 1) tgt)
              valid = 0 A.<= tgt A.&& tgt A.< vSize
              A.T2 m sumE = lse A.! A.I1 p
              prob = exp (logits A.! A.I2 p safeTgt - m) * (sumE ** (-1.0))
           in valid A.? (-log prob, 0.0)
      meanLoss = A.map (* invT) (A.sum stepLosses)
      dLogits =
        A.generate (A.I2 tLen vSize) $ \(A.I2 p v) ->
          let tgt = targetTokens A.! A.I1 p
              valid = 0 A.<= tgt A.&& tgt A.< vSize
              A.T2 m sumE = lse A.! A.I1 p
              prob = exp (logits A.! A.I2 p v - m) * (sumE ** (-1.0))
              indicator = (v A.== tgt) A.? (1.0, 0.0)
           in valid A.? (invT * (prob - indicator), 0.0)
   in (meanLoss, dLogits)

-- | Reusable mutable workspace in 'ST' holding forward caches and backward scratch buffers.
data STWorkspace s = STWorkspace
  { x0Buf :: UMVector.MVector s Double,
    scale0Buf :: UMVector.MVector s Double,
    xInBuf :: UMVector.MVector s Double,
    scale1Buf :: UMVector.MVector s Double,
    xNorm1Buf :: UMVector.MVector s Double,
    qBuf :: UMVector.MVector s Double,
    kBuf :: UMVector.MVector s Double,
    vBuf :: UMVector.MVector s Double,
    attnWBuf :: UMVector.MVector s Double,
    attnConcatBuf :: UMVector.MVector s Double,
    xMidBuf :: UMVector.MVector s Double,
    scale2Buf :: UMVector.MVector s Double,
    xNorm2Buf :: UMVector.MVector s Double,
    hiddenBuf :: UMVector.MVector s Double,
    probsBuf :: UMVector.MVector s Double,
    gradParamsBuf :: UMVector.MVector s Double,
    dXCurBuf :: UMVector.MVector s Double,
    dXNormBuf :: UMVector.MVector s Double,
    dHiddenBuf :: UMVector.MVector s Double,
    dAttnConcatBuf :: UMVector.MVector s Double,
    dQBuf :: UMVector.MVector s Double,
    dKBuf :: UMVector.MVector s Double,
    dVBuf :: UMVector.MVector s Double,
    dAttnWBuf :: UMVector.MVector s Double
  }

{-# INLINE newSTWorkspace #-}
newSTWorkspace :: Config -> Int -> ST s (STWorkspace s)
newSTWorkspace cfg nParams = do
  let bSize = max 1 cfg.blockSize
      dDim = max 1 cfg.nEmbd
      lCount = max 0 cfg.nLayer
      nHead = max 1 cfg.nHead
      vSize = max 1 cfg.vocabSize
      bd = bSize * dDim
      inSize = (lCount + 1) * bd
      lbd = max 1 (lCount * bd)
      lb = max 1 (lCount * bSize)
      mlp = bSize * 4 * dDim
      lmlp = max 1 (lCount * mlp)
      attnSize = max 1 (lCount * nHead * bSize * bSize)
      !oX0 = 0
      !oScale0 = oX0 + bd
      !oXIn = oScale0 + bSize
      !oScale1 = oXIn + inSize
      !oXNorm1 = oScale1 + lb
      !oQ = oXNorm1 + lbd
      !oK = oQ + lbd
      !oV = oK + lbd
      !oAttnW = oV + lbd
      !oAttnConcat = oAttnW + attnSize
      !oXMid = oAttnConcat + lbd
      !oScale2 = oXMid + lbd
      !oXNorm2 = oScale2 + lb
      !oHidden = oXNorm2 + lbd
      !oProbs = oHidden + lmlp
      !oDXCur = oProbs + vSize
      !oDXNorm = oDXCur + bd
      !oDHidden = oDXNorm + bd
      !oDAttnConcat = oDHidden + mlp
      !oDQ = oDAttnConcat + bd
      !oDK = oDQ + bd
      !oDV = oDK + bd
      !oDAttnW = oDV + bd
      !totalArena = oDAttnW + bSize
  arena <- UMVector.new totalArena
  gradParamsBuf <- UMVector.new (max 1 nParams)
  let x0Buf = UMVector.unsafeSlice oX0 bd arena
      scale0Buf = UMVector.unsafeSlice oScale0 bSize arena
      xInBuf = UMVector.unsafeSlice oXIn inSize arena
      scale1Buf = UMVector.unsafeSlice oScale1 lb arena
      xNorm1Buf = UMVector.unsafeSlice oXNorm1 lbd arena
      qBuf = UMVector.unsafeSlice oQ lbd arena
      kBuf = UMVector.unsafeSlice oK lbd arena
      vBuf = UMVector.unsafeSlice oV lbd arena
      attnWBuf = UMVector.unsafeSlice oAttnW attnSize arena
      attnConcatBuf = UMVector.unsafeSlice oAttnConcat lbd arena
      xMidBuf = UMVector.unsafeSlice oXMid lbd arena
      scale2Buf = UMVector.unsafeSlice oScale2 lb arena
      xNorm2Buf = UMVector.unsafeSlice oXNorm2 lbd arena
      hiddenBuf = UMVector.unsafeSlice oHidden lmlp arena
      probsBuf = UMVector.unsafeSlice oProbs vSize arena
      dXCurBuf = UMVector.unsafeSlice oDXCur bd arena
      dXNormBuf = UMVector.unsafeSlice oDXNorm bd arena
      dHiddenBuf = UMVector.unsafeSlice oDHidden mlp arena
      dAttnConcatBuf = UMVector.unsafeSlice oDAttnConcat bd arena
      dQBuf = UMVector.unsafeSlice oDQ bd arena
      dKBuf = UMVector.unsafeSlice oDK bd arena
      dVBuf = UMVector.unsafeSlice oDV bd arena
      dAttnWBuf = UMVector.unsafeSlice oDAttnW bSize arena
  pure STWorkspace {..}

-- | Low-latency forward-only sequence loss using a single contiguous 'ST' activation arena.
stForwardLoss :: Config -> UVector.Vector Double -> Vector Int -> Double
stForwardLoss cfg flatW tokens =
  let nPairs = min cfg.blockSize (max 0 (Vector.length tokens - 1))
   in if nPairs <= 0 || not (isValidModelShape cfg (UVector.length flatW))
        then 0.0
        else runST $ do
          let vSize = cfg.vocabSize
              lCount = cfg.nLayer
              dDim = cfg.nEmbd
              bSize = cfg.blockSize
              nHead = cfg.nHead
              hd = headDim cfg
              mlpDim = 4 * dDim
              dd = dDim * dDim
              offWte = 0
              offWpe = vSize * dDim
              offLmHead = (vSize + bSize) * dDim
              baseLayers = (2 * vSize + bSize) * dDim
              !invScale = 1.0 / sqrt (fromIntegral hd)
              !invT = 1.0 / fromIntegral nPairs
              td = nPairs * dDim
              tmlp = nPairs * mlpDim
              !oXCur = 0
              !oXNorm = oXCur + td
              !oScales = oXNorm + td
              !oQ = oScales + nPairs
              !oK = oQ + td
              !oV = oK + td
              !oAttnConcat = oV + td
              !oAttnOut = oAttnConcat + td
              !oHidden = oAttnOut + td
              !oScores = oHidden + tmlp
              !oProbs = oScores + nPairs
              !totalArena = oProbs + vSize
          arena <- UMVector.new totalArena
          let xCur = UMVector.unsafeSlice oXCur td arena
              xNorm = UMVector.unsafeSlice oXNorm td arena
              scales = UMVector.unsafeSlice oScales nPairs arena
              qBuf = UMVector.unsafeSlice oQ td arena
              kBuf = UMVector.unsafeSlice oK td arena
              vBuf = UMVector.unsafeSlice oV td arena
              attnConcat = UMVector.unsafeSlice oAttnConcat td arena
              attnOut = UMVector.unsafeSlice oAttnOut td arena
              hidden = UMVector.unsafeSlice oHidden tmlp arena
              scoresRow = UMVector.unsafeSlice oScores nPairs arena
              probsRow = UMVector.unsafeSlice oProbs vSize arena
          let embedLoop p
                | p >= nPairs = pure ()
                | otherwise = do
                    let !tok = Vector.unsafeIndex tokens p
                        loopD d
                          | d >= dDim = pure ()
                          | otherwise = do
                              let !tv =
                                    if tok >= 0 && tok < vSize
                                      then UVector.unsafeIndex flatW (offWte + tok * dDim + d)
                                      else 0.0
                                  !pv = UVector.unsafeIndex flatW (offWpe + p * dDim + d)
                              UMVector.unsafeWrite xCur (p * dDim + d) (tv + pv)
                              loopD (d + 1)
                    loopD 0
                    embedLoop (p + 1)
          embedLoop 0
          stRmsNormFwd dDim nPairs xCur 0 xCur 0 scales 0
          let layerLoop l
                | l >= lCount = pure ()
                | otherwise = do
                    let !lBase = baseLayers + l * 12 * dd
                        !offWq = lBase
                        !offWk = lBase + dd
                        !offWv = lBase + 2 * dd
                        !offWo = lBase + 3 * dd
                        !offFc1 = lBase + 4 * dd
                        !offFc2 = lBase + 8 * dd
                    stRmsNormFwd dDim nPairs xCur 0 xNorm 0 scales 0
                    stLinearFwdImmW nPairs dDim dDim xNorm 0 flatW offWq qBuf 0
                    stLinearFwdImmW nPairs dDim dDim xNorm 0 flatW offWk kBuf 0
                    stLinearFwdImmW nPairs dDim dDim xNorm 0 flatW offWv vBuf 0
                    stZeroSlice attnConcat 0 (nPairs * dDim)
                    let headLoop h
                          | h >= nHead = pure ()
                          | otherwise = do
                              let posLoop p
                                    | p >= nPairs = pure ()
                                    | otherwise = do
                                        let !qOff = p * dDim + h * hd
                                            scoreS s !maxSc
                                              | s > p = pure maxSc
                                              | otherwise = do
                                                  let !kOff = s * dDim + h * hd
                                                      dotC c !acc
                                                        | c >= hd = pure acc
                                                        | otherwise = do
                                                            qv <- UMVector.unsafeRead qBuf (qOff + c)
                                                            kv <- UMVector.unsafeRead kBuf (kOff + c)
                                                            dotC (c + 1) (acc + qv * kv)
                                                  raw <- dotC 0 0.0
                                                  let !sc = raw * invScale
                                                  UMVector.unsafeWrite scoresRow s sc
                                                  scoreS (s + 1) (max maxSc sc)
                                        maxSc <- scoreS 0 (-1.0e308)
                                        let expS s !sumE
                                              | s > p = pure sumE
                                              | otherwise = do
                                                  sc <- UMVector.unsafeRead scoresRow s
                                                  let !ev = exp (sc - maxSc)
                                                  UMVector.unsafeWrite scoresRow s ev
                                                  expS (s + 1) (sumE + ev)
                                        sumE <- expS 0 0.0
                                        let !invSum = 1.0 / sumE
                                            valS s
                                              | s > p = pure ()
                                              | otherwise = do
                                                  ev <- UMVector.unsafeRead scoresRow s
                                                  let !wS = ev * invSum
                                                      !vOff = s * dDim + h * hd
                                                      addC c
                                                        | c >= hd = pure ()
                                                        | otherwise = do
                                                            vv <- UMVector.unsafeRead vBuf (vOff + c)
                                                            stAddAt attnConcat (qOff + c) (wS * vv)
                                                            addC (c + 1)
                                                  addC 0
                                                  valS (s + 1)
                                        valS 0
                                        posLoop (p + 1)
                              posLoop 0
                              headLoop (h + 1)
                    headLoop 0
                    stLinearFwdImmW nPairs dDim dDim attnConcat 0 flatW offWo attnOut 0
                    let addRes1 i
                          | i >= nPairs * dDim = pure ()
                          | otherwise = do
                              ao <- UMVector.unsafeRead attnOut i
                              stAddAt xCur i ao
                              addRes1 (i + 1)
                    addRes1 0
                    stRmsNormFwd dDim nPairs xCur 0 xNorm 0 scales 0
                    stLinearFwdReluImmW nPairs dDim mlpDim xNorm 0 flatW offFc1 hidden 0
                    stLinearFwdImmW nPairs mlpDim dDim hidden 0 flatW offFc2 attnOut 0
                    let addRes2 i
                          | i >= nPairs * dDim = pure ()
                          | otherwise = do
                              mo <- UMVector.unsafeRead attnOut i
                              stAddAt xCur i mo
                              addRes2 (i + 1)
                    addRes2 0
                    layerLoop (l + 1)
          layerLoop 0
          let lossLoop p !accLoss
                | p >= nPairs = pure (accLoss * invT)
                | otherwise = do
                    let !tgt = Vector.unsafeIndex tokens (p + 1)
                    if tgt < 0 || tgt >= vSize
                      then lossLoop (p + 1) accLoss
                      else do
                        let !rowX = p * dDim
                            logitV v !maxL
                              | v >= vSize = pure maxL
                              | otherwise = do
                                  let !rowW = offLmHead + v * dDim
                                      dotD d !acc
                                        | d >= dDim = pure acc
                                        | otherwise = do
                                            xv <- UMVector.unsafeRead xCur (rowX + d)
                                            let !wv = UVector.unsafeIndex flatW (rowW + d)
                                            dotD (d + 1) (acc + xv * wv)
                                  lv <- dotD 0 0.0
                                  UMVector.unsafeWrite probsRow v lv
                                  logitV (v + 1) (max maxL lv)
                        maxL <- logitV 0 (-1.0e308)
                        let expV v !sumE
                              | v >= vSize = pure sumE
                              | otherwise = do
                                  lv <- UMVector.unsafeRead probsRow v
                                  let !ev = exp (lv - maxL)
                                  UMVector.unsafeWrite probsRow v ev
                                  expV (v + 1) (sumE + ev)
                        sumE <- expV 0 0.0
                        tgtExp <- UMVector.unsafeRead probsRow tgt
                        let !probTgt = tgtExp / sumE
                        lossLoop (p + 1) (accLoss - log probTgt)
          lossLoop 0 0.0

-- | Full forward and analytical backward pass in 'ST', mutating @ws.gradParamsBuf@ in-place.
{-# INLINE stForwardBackward #-}
stForwardBackward ::
  Config ->
  STWorkspace s ->
  UMVector.MVector s Double ->
  Vector Int ->
  Int ->
  ST s Double
stForwardBackward cfg ws valsM tokens tLen = do
  let STWorkspace {..} = ws
      vSize = cfg.vocabSize
      lCount = cfg.nLayer
      dDim = cfg.nEmbd
      bSize = cfg.blockSize
      nHead = cfg.nHead
      hd = headDim cfg
      mlpDim = 4 * dDim
      bd = bSize * dDim
      dd = dDim * dDim
      offWte = 0
      offWpe = vSize * dDim
      offLmHead = (vSize + bSize) * dDim
      baseLayers = (2 * vSize + bSize) * dDim
      totalParams = expectedParamCount cfg
      !invScale = 1.0 / sqrt (fromIntegral hd)
      !invT = 1.0 / fromIntegral tLen
  let embedLoop p
        | p >= tLen = pure ()
        | otherwise = do
            let !tok = Vector.unsafeIndex tokens p
                loopD d
                  | d >= dDim = pure ()
                  | otherwise = do
                      tv <-
                        if tok >= 0 && tok < vSize
                          then UMVector.unsafeRead valsM (offWte + tok * dDim + d)
                          else pure 0.0
                      pv <- UMVector.unsafeRead valsM (offWpe + p * dDim + d)
                      UMVector.unsafeWrite x0Buf (p * dDim + d) (tv + pv)
                      loopD (d + 1)
            loopD 0
            embedLoop (p + 1)
  embedLoop 0
  stRmsNormFwd dDim tLen x0Buf 0 xInBuf 0 scale0Buf 0
  let fwdLayers l
        | l >= lCount = pure ()
        | otherwise = do
            let !lOff = l * bd
                !nextLOff = (l + 1) * bd
                !lScaleOff = l * bSize
                !lMlpOff = l * bSize * mlpDim
                !lBase = baseLayers + l * 12 * dd
                !offWq = lBase
                !offWk = lBase + dd
                !offWv = lBase + 2 * dd
                !offWo = lBase + 3 * dd
                !offFc1 = lBase + 4 * dd
                !offFc2 = lBase + 8 * dd
            stRmsNormFwd dDim tLen xInBuf lOff xNorm1Buf lOff scale1Buf lScaleOff
            stLinearFwd tLen dDim dDim xNorm1Buf lOff valsM offWq qBuf lOff
            stLinearFwd tLen dDim dDim xNorm1Buf lOff valsM offWk kBuf lOff
            stLinearFwd tLen dDim dDim xNorm1Buf lOff valsM offWv vBuf lOff
            stZeroSlice attnConcatBuf lOff (tLen * dDim)
            let headLoop h
                  | h >= nHead = pure ()
                  | otherwise = do
                      let posLoop p
                            | p >= tLen = pure ()
                            | otherwise = do
                                let !qOff = lOff + p * dDim + h * hd
                                    !wRowOff = ((l * nHead + h) * bSize + p) * bSize
                                    scoreS s !maxSc
                                      | s > p = pure maxSc
                                      | otherwise = do
                                          let !kOff = lOff + s * dDim + h * hd
                                              dotC c !acc
                                                | c >= hd = pure acc
                                                | otherwise = do
                                                    qv <- UMVector.unsafeRead qBuf (qOff + c)
                                                    kv <- UMVector.unsafeRead kBuf (kOff + c)
                                                    dotC (c + 1) (acc + qv * kv)
                                          raw <- dotC 0 0.0
                                          let !sc = raw * invScale
                                          UMVector.unsafeWrite attnWBuf (wRowOff + s) sc
                                          scoreS (s + 1) (max maxSc sc)
                                maxSc <- scoreS 0 (-1.0e308)
                                let expS s !sumE
                                      | s > p = pure sumE
                                      | otherwise = do
                                          sc <- UMVector.unsafeRead attnWBuf (wRowOff + s)
                                          let !ev = exp (sc - maxSc)
                                          UMVector.unsafeWrite attnWBuf (wRowOff + s) ev
                                          expS (s + 1) (sumE + ev)
                                sumE <- expS 0 0.0
                                let !invSum = 1.0 / sumE
                                    valS s
                                      | s > p = pure ()
                                      | otherwise = do
                                          ev <- UMVector.unsafeRead attnWBuf (wRowOff + s)
                                          let !wS = ev * invSum
                                              !vOff = lOff + s * dDim + h * hd
                                          UMVector.unsafeWrite attnWBuf (wRowOff + s) wS
                                          let addC c
                                                | c >= hd = pure ()
                                                | otherwise = do
                                                    vv <- UMVector.unsafeRead vBuf (vOff + c)
                                                    stAddAt attnConcatBuf (qOff + c) (wS * vv)
                                                    addC (c + 1)
                                          addC 0
                                          valS (s + 1)
                                valS 0
                                posLoop (p + 1)
                      posLoop 0
                      headLoop (h + 1)
            headLoop 0
            stLinearFwd tLen dDim dDim attnConcatBuf lOff valsM offWo dXCurBuf 0
            let addRes1 i
                  | i >= tLen * dDim = pure ()
                  | otherwise = do
                      ao <- UMVector.unsafeRead dXCurBuf i
                      xi <- UMVector.unsafeRead xInBuf (lOff + i)
                      UMVector.unsafeWrite xMidBuf (lOff + i) (ao + xi)
                      addRes1 (i + 1)
            addRes1 0
            stRmsNormFwd dDim tLen xMidBuf lOff xNorm2Buf lOff scale2Buf lScaleOff
            stLinearFwdRelu tLen dDim mlpDim xNorm2Buf lOff valsM offFc1 hiddenBuf lMlpOff
            stLinearFwd tLen mlpDim dDim hiddenBuf lMlpOff valsM offFc2 dXCurBuf 0
            let addRes2 i
                  | i >= tLen * dDim = pure ()
                  | otherwise = do
                      mo <- UMVector.unsafeRead dXCurBuf i
                      xm <- UMVector.unsafeRead xMidBuf (lOff + i)
                      UMVector.unsafeWrite xInBuf (nextLOff + i) (mo + xm)
                      addRes2 (i + 1)
            addRes2 0
            fwdLayers (l + 1)
  fwdLayers 0
  stZeroSlice gradParamsBuf 0 totalParams
  stZeroSlice dXCurBuf 0 (tLen * dDim)
  let !finalOff = lCount * bd
      lossAndHeadBwd p !accLoss
        | p >= tLen = pure (accLoss * invT)
        | otherwise = do
            let !tgt = Vector.unsafeIndex tokens (p + 1)
            if tgt < 0 || tgt >= vSize
              then lossAndHeadBwd (p + 1) accLoss
              else do
                let !rowX = finalOff + p * dDim
                    !rowDx = p * dDim
                    logitV v !maxL
                      | v >= vSize = pure maxL
                      | otherwise = do
                          let !rowW = offLmHead + v * dDim
                              dotD d !acc
                                | d >= dDim = pure acc
                                | otherwise = do
                                    xv <- UMVector.unsafeRead xInBuf (rowX + d)
                                    wv <- UMVector.unsafeRead valsM (rowW + d)
                                    dotD (d + 1) (acc + xv * wv)
                          lv <- dotD 0 0.0
                          UMVector.unsafeWrite probsBuf v lv
                          logitV (v + 1) (max maxL lv)
                maxL <- logitV 0 (-1.0e308)
                let expV v !sumE
                      | v >= vSize = pure sumE
                      | otherwise = do
                          lv <- UMVector.unsafeRead probsBuf v
                          let !ev = exp (lv - maxL)
                          UMVector.unsafeWrite probsBuf v ev
                          expV (v + 1) (sumE + ev)
                sumE <- expV 0 0.0
                let !invSum = 1.0 / sumE
                tgtExp <- UMVector.unsafeRead probsBuf tgt
                let !probTgt = tgtExp * invSum
                    bwdV v
                      | v >= vSize = pure ()
                      | otherwise = do
                          ev <- UMVector.unsafeRead probsBuf v
                          let !probV = ev * invSum
                              !dl = invT * (probV - (if v == tgt then 1.0 else 0.0))
                              !rowW = offLmHead + v * dDim
                              bwdD d
                                | d >= dDim = pure ()
                                | otherwise = do
                                    xv <- UMVector.unsafeRead xInBuf (rowX + d)
                                    wv <- UMVector.unsafeRead valsM (rowW + d)
                                    stAddAt gradParamsBuf (rowW + d) (dl * xv)
                                    stAddAt dXCurBuf (rowDx + d) (dl * wv)
                                    bwdD (d + 1)
                          bwdD 0
                          bwdV (v + 1)
                bwdV 0
                lossAndHeadBwd (p + 1) (accLoss - log probTgt)
  meanLoss <- lossAndHeadBwd 0 0.0
  let bwdLayers l
        | l < 0 = pure ()
        | otherwise = do
            let !lOff = l * bd
                !lScaleOff = l * bSize
                !lMlpOff = l * bSize * mlpDim
                !lBase = baseLayers + l * 12 * dd
                !offWq = lBase
                !offWk = lBase + dd
                !offWv = lBase + 2 * dd
                !offWo = lBase + 3 * dd
                !offFc1 = lBase + 4 * dd
                !offFc2 = lBase + 8 * dd
            stLinearBwd tLen dDim mlpDim dXCurBuf 0 hiddenBuf lMlpOff valsM offFc2 gradParamsBuf dHiddenBuf 0
            let maskRelu i
                  | i >= tLen * mlpDim = pure ()
                  | otherwise = do
                      hv <- UMVector.unsafeRead hiddenBuf (lMlpOff + i)
                      when (hv <= 0.0) $ UMVector.unsafeWrite dHiddenBuf i 0.0
                      maskRelu (i + 1)
            maskRelu 0
            stLinearBwd tLen mlpDim dDim dHiddenBuf 0 xNorm2Buf lOff valsM offFc1 gradParamsBuf dXNormBuf 0
            stRmsNormBwdAdd dDim tLen xMidBuf lOff scale2Buf lScaleOff dXNormBuf 0 dXCurBuf 0
            stLinearBwd tLen dDim dDim dXCurBuf 0 attnConcatBuf lOff valsM offWo gradParamsBuf dAttnConcatBuf 0
            stZeroSlice dQBuf 0 (tLen * dDim)
            stZeroSlice dKBuf 0 (tLen * dDim)
            stZeroSlice dVBuf 0 (tLen * dDim)
            let bwdHeads h
                  | h >= nHead = pure ()
                  | otherwise = do
                      let bwdPos p
                            | p >= tLen = pure ()
                            | otherwise = do
                                let !baseP = p * dDim + h * hd
                                    !wRowOff = ((l * nHead + h) * bSize + p) * bSize
                                    pass1 s !wDotDw
                                      | s > p = pure wDotDw
                                      | otherwise = do
                                          wS <- UMVector.unsafeRead attnWBuf (wRowOff + s)
                                          let !vOff = lOff + s * dDim + h * hd
                                              !dvOff = s * dDim + h * hd
                                              dotAndDv c !dwAcc
                                                | c >= hd = pure dwAcc
                                                | otherwise = do
                                                    dac <- UMVector.unsafeRead dAttnConcatBuf (baseP + c)
                                                    vv <- UMVector.unsafeRead vBuf (vOff + c)
                                                    stAddAt dVBuf (dvOff + c) (wS * dac)
                                                    dotAndDv (c + 1) (dwAcc + dac * vv)
                                          dwS <- dotAndDv 0 0.0
                                          UMVector.unsafeWrite dAttnWBuf s dwS
                                          pass1 (s + 1) (wDotDw + wS * dwS)
                                wDotDw <- pass1 0 0.0
                                let pass2 s
                                      | s > p = pure ()
                                      | otherwise = do
                                          wS <- UMVector.unsafeRead attnWBuf (wRowOff + s)
                                          dwS <- UMVector.unsafeRead dAttnWBuf s
                                          let !ds = invScale * wS * (dwS - wDotDw)
                                              !kOff = lOff + s * dDim + h * hd
                                              !qOff = lOff + baseP
                                              !dkOff = s * dDim + h * hd
                                              accumQK c
                                                | c >= hd = pure ()
                                                | otherwise = do
                                                    kv <- UMVector.unsafeRead kBuf (kOff + c)
                                                    qv <- UMVector.unsafeRead qBuf (qOff + c)
                                                    stAddAt dQBuf (baseP + c) (ds * kv)
                                                    stAddAt dKBuf (dkOff + c) (ds * qv)
                                                    accumQK (c + 1)
                                          accumQK 0
                                          pass2 (s + 1)
                                pass2 0
                                bwdPos (p + 1)
                      bwdPos 0
                      bwdHeads (h + 1)
            bwdHeads 0
            stZeroSlice dXNormBuf 0 (tLen * dDim)
            stLinearBwdAccum tLen dDim dDim dQBuf 0 xNorm1Buf lOff valsM offWq gradParamsBuf dXNormBuf 0
            stLinearBwdAccum tLen dDim dDim dKBuf 0 xNorm1Buf lOff valsM offWk gradParamsBuf dXNormBuf 0
            stLinearBwdAccum tLen dDim dDim dVBuf 0 xNorm1Buf lOff valsM offWv gradParamsBuf dXNormBuf 0
            stRmsNormBwdAdd dDim tLen xInBuf lOff scale1Buf lScaleOff dXNormBuf 0 dXCurBuf 0
            bwdLayers (l - 1)
  bwdLayers (lCount - 1)
  stZeroSlice dXNormBuf 0 (tLen * dDim)
  stRmsNormBwdAdd dDim tLen x0Buf 0 scale0Buf 0 dXCurBuf 0 dXNormBuf 0
  let embedBwd p
        | p >= tLen = pure ()
        | otherwise = do
            let !tok = Vector.unsafeIndex tokens p
                loopD d
                  | d >= dDim = pure ()
                  | otherwise = do
                      g <- UMVector.unsafeRead dXNormBuf (p * dDim + d)
                      when (tok >= 0 && tok < vSize) $
                        stAddAt gradParamsBuf (offWte + tok * dDim + d) g
                      stAddAt gradParamsBuf (offWpe + p * dDim + d) g
                      loopD (d + 1)
            loopD 0
            embedBwd (p + 1)
  embedBwd 0
  pure meanLoss

-- | In-place Adam parameter and moment buffer update in 'ST'.
{-# INLINE stAdamUpdate #-}
stAdamUpdate ::
  Hyperparams ->
  Double ->
  Int ->
  Int ->
  UMVector.MVector s Double ->
  UMVector.MVector s Double ->
  UMVector.MVector s Double ->
  UMVector.MVector s Double ->
  ST s ()
stAdamUpdate hp lr t nParams gradsM valsM mM vM =
  let !b1 = hp.beta1
      !b2 = hp.beta2
      !oneMinusB1 = 1.0 - b1
      !oneMinusB2 = 1.0 - b2
      !tD = fromIntegral t :: Double
      !b1Corr = 1.0 - (b1 ** tD)
      !b2Corr = 1.0 - (b2 ** tD)
      !eps = hp.epsilon
      loop i
        | i >= nParams = pure ()
        | otherwise = do
            g <- UMVector.unsafeRead gradsM i
            mOld <- UMVector.unsafeRead mM i
            vOld <- UMVector.unsafeRead vM i
            valOld <- UMVector.unsafeRead valsM i
            let !mNew = b1 * mOld + oneMinusB1 * g
                !vNew = b2 * vOld + oneMinusB2 * (g * g)
                !mHat = mNew / b1Corr
                !vHat = vNew / b2Corr
                !valNew = valOld - lr * mHat / (sqrt vHat + eps)
            UMVector.unsafeWrite mM i mNew
            UMVector.unsafeWrite vM i vNew
            UMVector.unsafeWrite valsM i valNew
            loop (i + 1)
   in loop 0

-- 4. Tokenizer

data Tokenizer = Tokenizer
  { uchars :: Vector Char,
    tokenIds :: Map.Map Char Int,
    asciiIds :: UVector.Vector Int,
    bos :: Int,
    vocabSize :: Int
  }

-- | Builds the character vocabulary using an in-place ASCII bitset in 'ST'
-- (with a 'Set' fallback for non-ASCII code points).
mkTokenizer :: Vector String -> Tokenizer
mkTokenizer docs =
  let uchars = runST $ do
        asciiSeen <- UMVector.replicate 128 False
        extraRef <- STRef.newSTRef Set.empty
        Vector.forM_ docs $ \doc ->
          Foldable.forM_ doc $ \c -> do
            let code = Char.ord c
            if code < 128
              then UMVector.unsafeWrite asciiSeen code True
              else STRef.modifySTRef' extraRef (Set.insert c)
        frozenAscii <- UVector.unsafeFreeze asciiSeen
        extraSet <- STRef.readSTRef extraRef
        let collectAscii code
              | code >= 128 = Set.toAscList extraSet
              | UVector.unsafeIndex frozenAscii code = Char.chr code : collectAscii (code + 1)
              | otherwise = collectAscii (code + 1)
        pure (Vector.fromList (collectAscii 0))
      tokenMap = Vector.ifoldl' (\m i c -> Map.insert c i m) Map.empty uchars
      asciiTable = runST $ do
        tbl <- UMVector.replicate 128 (-1)
        Vector.iforM_ uchars $ \i c -> do
          let code = Char.ord c
          when (code < 128) $ UMVector.unsafeWrite tbl code i
        UVector.unsafeFreeze tbl
      nChars = Vector.length uchars
   in Tokenizer
        { uchars = uchars,
          tokenIds = tokenMap,
          asciiIds = asciiTable,
          bos = nChars,
          vocabSize = nChars + 1
        }

-- | Surrounds the document with BOS on both sides (unknown characters are dropped),
-- filling a single mutable token vector in 'ST'.
encode :: Tokenizer -> String -> Vector Int
encode tok doc = runST $ do
  let maxLen = length doc + 2
  buf <- MVector.new maxLen
  MVector.unsafeWrite buf 0 tok.bos
  let fill chars !wIdx = case chars of
        [] -> do
          MVector.unsafeWrite buf wIdx tok.bos
          frozen <- Vector.unsafeFreeze buf
          pure (Vector.take (wIdx + 1) frozen)
        c : cs -> do
          let code = Char.ord c
              mTid
                | code < 128 =
                    let tid = UVector.unsafeIndex tok.asciiIds code
                     in if tid >= 0 then Just tid else Nothing
                | otherwise = Map.lookup c tok.tokenIds
          case mTid of
            Nothing -> fill cs wIdx
            Just tid -> do
              MVector.unsafeWrite buf wIdx tid
              fill cs (wIdx + 1)
  fill doc 1

decode :: Tokenizer -> Vector Int -> String
decode tok = Vector.toList . Vector.mapMaybe (tok.uchars Vector.!?)

-- 5. Training

-- | Single-layer specialized forward pass compiled once to native LLVM machine code.
accFwd1 ::
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Scalar Double)
accFwd1 cfgVec flatW inTokens targetTokens =
  let (vSize, _lCount, dDim, bSize, nHead, hd) = unpackAccConfig cfgVec
      wte = sliceMatrix flatW 0 vSize dDim
      wpe = sliceMatrix flatW (vSize * dDim) bSize dDim
      lmHead = sliceMatrix flatW ((vSize + bSize) * dDim) vSize dDim
      x0 = accEmbedFwd vSize bSize dDim wte wpe inTokens
      (x1, _ms0) = accRmsNorm2DWithMs x0
      (x2, _cache0) = accLayerFwdWithCache vSize bSize dDim nHead hd flatW 0 x1
      logits = accLinearFwd x2 lmHead
   in accLossOnly vSize logits targetTokens

-- | General multi-layer forward pass compiled once to native LLVM machine code.
accFwdGen ::
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Scalar Double)
accFwdGen cfgVec flatW inTokens targetTokens =
  let (vSize, lCount, dDim, bSize, nHead, hd) = unpackAccConfig cfgVec
      wte = sliceMatrix flatW 0 vSize dDim
      wpe = sliceMatrix flatW (vSize * dDim) bSize dDim
      lmHead = sliceMatrix flatW ((vSize + bSize) * dDim) vSize dDim
      x0 = accEmbedFwd vSize bSize dDim wte wpe inTokens
      (x1, _ms0) = accRmsNorm2DWithMs x0
      A.T2 _ xFinal =
        A.awhile
          (\(A.T2 l _) -> A.map (A.< lCount) l)
          ( \(A.T2 l xCur) ->
              let (xNext, _) = accLayerFwdWithCache vSize bSize dDim nHead hd flatW (A.the l) xCur
               in A.T2 (A.map (+ 1) l) xNext
          )
          (A.T2 (A.unit 0) x1)
      logits = accLinearFwd xFinal lmHead
   in accLossOnly vSize logits targetTokens

-- | Core single-layer forward and analytical backward pass returning @(loss, flatGrads)@.
accFwdBwd1Core ::
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Int) ->
  (A.Acc (A.Scalar Double), A.Acc (A.Vector Double))
accFwdBwd1Core cfgVec flatW inTokens targetTokens =
  let (vSize, _lCount, dDim, bSize, nHead, hd) = unpackAccConfig cfgVec
      wte = sliceMatrix flatW 0 vSize dDim
      wpe = sliceMatrix flatW (vSize * dDim) bSize dDim
      lmHead = sliceMatrix flatW ((vSize + bSize) * dDim) vSize dDim
      x0 = accEmbedFwd vSize bSize dDim wte wpe inTokens
      (x1, ms0) = accRmsNorm2DWithMs x0
      (x2, cache0) = accLayerFwdWithCache vSize bSize dDim nHead hd flatW 0 x1
      logits = accLinearFwd x2 lmHead
      (meanLoss, dLogits) = accLossAndGradLogits vSize logits targetTokens
      dLmHead = accLinearBwdW dLogits x2
      dX2 = accLinearBwdX dLogits lmHead
      (dX1, dLayer0) = accLayerBwdFromCache dDim nHead hd cache0 dX2
      dX0 = accRmsNorm2DBwdFromMs x0 ms0 dX1
      (dWte, dWpe) = accEmbedBwd vSize bSize dDim inTokens dX0
      dFlat = packAllGrads vSize bSize dDim dWte dWpe dLmHead dLayer0
   in (meanLoss, dFlat)

accFwdBwd1 ::
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Scalar Double, A.Vector Double)
accFwdBwd1 cfgVec flatW inTokens targetTokens =
  let (meanLoss, dFlat) = accFwdBwd1Core cfgVec flatW inTokens targetTokens
   in A.T2 meanLoss dFlat

-- | General multi-layer forward and analytical backward pass for arbitrary @nLayer@.
accFwdBwdGen ::
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Scalar Double, A.Vector Double)
accFwdBwdGen cfgVec flatW inTokens targetTokens =
  let (vSize, lCount, dDim, bSize, nHead, hd) = unpackAccConfig cfgVec
      A.I1 tLen = A.shape inTokens
      wte = sliceMatrix flatW 0 vSize dDim
      wpe = sliceMatrix flatW (vSize * dDim) bSize dDim
      lmHead = sliceMatrix flatW ((vSize + bSize) * dDim) vSize dDim
      x0 = accEmbedFwd vSize bSize dDim wte wpe inTokens
      (x1, ms0) = accRmsNorm2DWithMs x0
      acts0 =
        A.generate (A.I3 (lCount + 1) tLen dDim) $ \(A.I3 k p d) ->
          (k A.== 0) A.? (x1 A.! A.I2 p d, 0.0)
      A.T2 _ acts =
        A.awhile
          (\(A.T2 l _) -> A.map (A.< lCount) l)
          ( \(A.T2 l curActs) ->
              let lIdx = A.the l
                  xL = A.generate (A.I2 tLen dDim) (\(A.I2 p d) -> curActs A.! A.I3 lIdx p d)
                  (xNext, _) = accLayerFwdWithCache vSize bSize dDim nHead hd flatW lIdx xL
                  nextActs =
                    A.generate (A.I3 (lCount + 1) tLen dDim) $ \(A.I3 k p d) ->
                      (k A.== lIdx + 1) A.? (xNext A.! A.I2 p d, curActs A.! A.I3 k p d)
               in A.T2 (A.map (+ 1) l) nextActs
          )
          (A.T2 (A.unit 0) acts0)
      xFinal = A.generate (A.I2 tLen dDim) (\(A.I2 p d) -> acts A.! A.I3 lCount p d)
      logits = accLinearFwd xFinal lmHead
      (meanLoss, dLogits) = accLossAndGradLogits vSize logits targetTokens
      dLmHead = accLinearBwdW dLogits xFinal
      dXFinal = accLinearBwdX dLogits lmHead
      layerStride = 12 * dDim * dDim
      dLayers0 = A.fill (A.I1 (lCount * layerStride)) 0.0
      A.T3 _ dX1 dLayers =
        A.awhile
          (\(A.T3 l _ _) -> A.map (A.>= 0) l)
          ( \(A.T3 l dXCur dLCur) ->
              let lIdx = A.the l
                  xL = A.generate (A.I2 tLen dDim) (\(A.I2 p d) -> acts A.! A.I3 lIdx p d)
                  (_, cacheL) = accLayerFwdWithCache vSize bSize dDim nHead hd flatW lIdx xL
                  (dXPrev, dLayerL) = accLayerBwdFromCache dDim nHead hd cacheL dXCur
                  dLNext =
                    A.generate (A.I1 (lCount * layerStride)) $ \(A.I1 idx) ->
                      let rel = idx - lIdx * layerStride
                          safeRel = A.max 0 (A.min (layerStride - 1) rel)
                       in (0 A.<= rel A.&& rel A.< layerStride) A.? (dLayerL A.! A.I1 safeRel, dLCur A.! A.I1 idx)
               in A.T3 (A.map (\c -> c - 1) l) dXPrev dLNext
          )
          (A.T3 (A.unit (lCount - 1)) dXFinal dLayers0)
      dX0 = accRmsNorm2DBwdFromMs x0 ms0 dX1
      (dWte, dWpe) = accEmbedBwd vSize bSize dDim inTokens dX0
      dFlat = packAllGrads vSize bSize dDim dWte dWpe dLmHead dLayers
   in A.T2 meanLoss dFlat

-- | Vectorized Adam update across all flattened model parameters.
accAdamUpdate ::
  A.Exp Double ->
  A.Exp Double ->
  A.Exp Double ->
  A.Exp Double ->
  A.Exp Double ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  (A.Acc (A.Vector Double), A.Acc (A.Vector Double), A.Acc (A.Vector Double))
accAdamUpdate lr beta1 beta2 eps tD grads vals ms vs =
  let oneMinusB1 = 1.0 - beta1
      oneMinusB2 = 1.0 - beta2
      b1Corr = 1.0 - (beta1 ** tD)
      b2Corr = 1.0 - (beta2 ** tD)
      ms' = A.zipWith (\m g -> beta1 * m + oneMinusB1 * g) ms grads
      vs' = A.zipWith (\v g -> beta2 * v + oneMinusB2 * (g ** 2.0)) vs grads
      vals' =
        A.zipWith3
          ( \val m v ->
              let mHat = m / b1Corr
                  vHat = v / b2Corr
               in val - lr * mHat / (sqrt vHat + eps)
          )
          vals
          ms'
          vs'
   in (vals', ms', vs')

accAdamStep ::
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double, A.Vector Double, A.Vector Double)
accAdamStep hpVec grads vals ms vs =
  let lr = hpVec A.! A.I1 0
      beta1 = hpVec A.! A.I1 1
      beta2 = hpVec A.! A.I1 2
      eps = hpVec A.! A.I1 3
      tD = hpVec A.! A.I1 4
      (vals', ms', vs') = accAdamUpdate lr beta1 beta2 eps tD grads vals ms vs
   in A.T3 vals' ms' vs'

-- | Fused single-step forward + backward + Adam update inside Accelerate.
accTrainStep1 ::
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Scalar Double, A.Vector Double, A.Vector Double, A.Vector Double)
accTrainStep1 cfgVec hpVec vals ms vs inTokens targetTokens =
  let lr = hpVec A.! A.I1 0
      beta1 = hpVec A.! A.I1 1
      beta2 = hpVec A.! A.I1 2
      eps = hpVec A.! A.I1 3
      tD = hpVec A.! A.I1 4
      (meanLoss, dFlat) = accFwdBwd1Core cfgVec vals inTokens targetTokens
      (vals', ms', vs') = accAdamUpdate lr beta1 beta2 eps tD dFlat vals ms vs
   in A.T4 meanLoss vals' ms' vs'

-- | Fused multi-step training loop executing @awhile@ entirely inside Accelerate.
accTrainBatch1 ::
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Matrix Int) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Vector Double, A.Vector Double, A.Vector Double, A.Vector Double)
accTrainBatch1 cfgVec hpFloats stepInts docsMat pairCounts vals0 ms0 vs0 =
  let baseLr = hpFloats A.! A.I1 0
      beta1 = hpFloats A.! A.I1 1
      beta2 = hpFloats A.! A.I1 2
      eps = hpFloats A.! A.I1 3
      step0 = stepInts A.! A.I1 0
      totalSteps = stepInts A.! A.I1 1
      nDocs = stepInts A.! A.I1 2
      losses0 = A.fill (A.I1 nDocs) 0.0
      A.T5 _ finalLosses finalVals finalMs finalVs =
        A.awhile
          (\(A.T5 k _ _ _ _) -> A.map (A.< nDocs) k)
          ( \(A.T5 k losses vals ms vs) ->
              let docIdx = A.the k
                  tLen = pairCounts A.! A.I1 docIdx
                  curStep = step0 + docIdx
                  lr = baseLr * (1.0 - A.fromIntegral curStep / A.fromIntegral totalSteps)
                  tD = A.fromIntegral (curStep + 1)
                  A.T2 docLoss dFlat =
                    A.acond
                      (tLen A.<= 0)
                      (A.T2 (A.unit 0.0) (A.fill (A.shape vals) 0.0))
                      ( let inToks = A.generate (A.I1 tLen) (\(A.I1 p) -> docsMat A.! A.I2 docIdx p)
                            tgtToks = A.generate (A.I1 tLen) (\(A.I1 p) -> docsMat A.! A.I2 docIdx (p + 1))
                            (l, g) = accFwdBwd1Core cfgVec vals inToks tgtToks
                         in A.T2 l g
                      )
                  (vals', ms', vs') = accAdamUpdate lr beta1 beta2 eps tD dFlat vals ms vs
                  losses' =
                    A.generate (A.I1 nDocs) $ \(A.I1 i) ->
                      (i A.== docIdx) A.? (A.the docLoss, losses A.! A.I1 i)
               in A.T5 (A.map (+ 1) k) losses' vals' ms' vs'
          )
          (A.T5 (A.unit 0) losses0 vals0 ms0 vs0)
   in A.T4 finalLosses finalVals finalMs finalVs

{-# NOINLINE compiledFwd1 #-}
compiledFwd1 ::
  A.Vector Int ->
  A.Vector Double ->
  A.Vector Int ->
  A.Vector Int ->
  A.Scalar Double
compiledFwd1 = CPU.runN accFwd1

{-# NOINLINE compiledFwdGen #-}
compiledFwdGen ::
  A.Vector Int ->
  A.Vector Double ->
  A.Vector Int ->
  A.Vector Int ->
  A.Scalar Double
compiledFwdGen = CPU.runN accFwdGen

{-# NOINLINE compiledFwdBwd1 #-}
compiledFwdBwd1 ::
  A.Vector Int ->
  A.Vector Double ->
  A.Vector Int ->
  A.Vector Int ->
  (A.Scalar Double, A.Vector Double)
compiledFwdBwd1 = CPU.runN accFwdBwd1

{-# NOINLINE compiledFwdBwdGen #-}
compiledFwdBwdGen ::
  A.Vector Int ->
  A.Vector Double ->
  A.Vector Int ->
  A.Vector Int ->
  (A.Scalar Double, A.Vector Double)
compiledFwdBwdGen = CPU.runN accFwdBwdGen

{-# NOINLINE compiledAdamStep #-}
compiledAdamStep ::
  A.Vector Double ->
  A.Vector Double ->
  A.Vector Double ->
  A.Vector Double ->
  A.Vector Double ->
  (A.Vector Double, A.Vector Double, A.Vector Double)
compiledAdamStep = CPU.runN accAdamStep

{-# NOINLINE compiledTrainStep1 #-}
compiledTrainStep1 ::
  A.Vector Int ->
  A.Vector Double ->
  A.Vector Double ->
  A.Vector Double ->
  A.Vector Double ->
  A.Vector Int ->
  A.Vector Int ->
  (A.Scalar Double, A.Vector Double, A.Vector Double, A.Vector Double)
compiledTrainStep1 = CPU.runN accTrainStep1

{-# NOINLINE compiledTrainBatch1 #-}
compiledTrainBatch1 ::
  A.Vector Int ->
  A.Vector Double ->
  A.Vector Int ->
  A.Matrix Int ->
  A.Vector Int ->
  A.Vector Double ->
  A.Vector Double ->
  A.Vector Double ->
  (A.Vector Double, A.Vector Double, A.Vector Double, A.Vector Double)
compiledTrainBatch1 = CPU.runN accTrainBatch1

packTokenPairs :: Int -> Vector Int -> (A.Vector Int, A.Vector Int)
packTokenPairs nPairs tokens =
  let !inToks =
        A.fromFunction
          (A.Z :. nPairs)
          (\(A.Z :. i) -> Maybe.fromMaybe 0 (tokens Vector.!? i))
      !tgtToks =
        A.fromFunction
          (A.Z :. nPairs)
          (\(A.Z :. i) -> Maybe.fromMaybe 0 (tokens Vector.!? (i + 1)))
   in (inToks, tgtToks)

dFlatToGradients :: A.Vector Double -> Gradients
dFlatToGradients dFlat =
  let A.Z :. len = A.arrayShape dFlat
      go i
        | i >= len = []
        | otherwise =
            let !v = A.indexArray dFlat (A.Z :. i)
             in (i, v) : go (i + 1)
   in NodeIdxMap (IntMap.fromDistinctAscList (go 0))

uvecToGradients :: UVector.Vector Double -> Gradients
uvecToGradients uvec =
  let len = UVector.length uvec
      go i
        | i >= len = []
        | otherwise =
            let !v = UVector.unsafeIndex uvec i
             in (i, v) : go (i + 1)
   in NodeIdxMap (IntMap.fromDistinctAscList (go 0))

-- | Loss of a document, forward pass only
{-# INLINE forwardLoss #-}
forwardLoss :: Config -> Weights Double -> Vector Int -> Double
forwardLoss cfg pureWeights
  | shouldUseAccelerate cfg =
      let cfgVec = packConfig cfg
          flatW = packWeights pureWeights
          runFwd = if cfg.nLayer == 1 then compiledFwd1 else compiledFwdGen
       in \tokens ->
            let nPairs = min cfg.blockSize (max 0 (Vector.length tokens - 1))
             in if nPairs <= 0
                  then 0.0
                  else
                    let (!inToks, !tgtToks) = packTokenPairs nPairs tokens
                     in unScalar (runFwd cfgVec flatW inToks tgtToks)
  | otherwise =
      let flatU = flattenWeightsUVec pureWeights
       in stForwardLoss cfg flatU

-- | Runs a full forward and backward pass for a sequence of tokens
{-# INLINE trainingStep #-}
trainingStep :: Config -> Weights Double -> Vector Int -> (Double, Gradients)
trainingStep cfg pureWeights
  | shouldUseAccelerate cfg =
      let cfgVec = packConfig cfg
          flatW = packWeights pureWeights
          runFwdBwd = if cfg.nLayer == 1 then compiledFwdBwd1 else compiledFwdBwdGen
       in \tokens ->
            let nPairs = min cfg.blockSize (max 0 (Vector.length tokens - 1))
             in if nPairs <= 0
                  then (0.0, nimEmpty)
                  else
                    let (!inToks, !tgtToks) = packTokenPairs nPairs tokens
                        (lossS, dFlat) = runFwdBwd cfgVec flatW inToks tgtToks
                     in (unScalar lossS, dFlatToGradients dFlat)
  | otherwise =
      let flatU = flattenWeightsUVec pureWeights
       in \tokens ->
            let nPairs = min cfg.blockSize (max 0 (Vector.length tokens - 1))
                nParams = UVector.length flatU
             in if nPairs <= 0 || not (isValidModelShape cfg nParams)
                  then (0.0, nimEmpty)
                  else
                    let (!loss, !gradUVec) = runST $ do
                          valsM <- UVector.thaw flatU
                          ws <- newSTWorkspace cfg nParams
                          l <- stForwardBackward cfg ws valsM tokens nPairs
                          gFrozen <- UVector.unsafeFreeze ws.gradParamsBuf
                          pure (l, gFrozen)
                     in (loss, uvecToGradients gradUVec)

-- | Pairs each parameter with its gradient (parameters are the first nodes of the tape)
withGradients :: Gradients -> Weights a -> Weights (Double, a)
withGradients grads =
  snd . Traversable.mapAccumL (\i x -> (i + 1, (Maybe.fromMaybe 0 (nimLookup (NodeId i) grads), x))) 0

parameterGradients :: Gradients -> Weights a -> Weights Double
parameterGradients grads = fmap fst . withGradients grads

data Hyperparams = Hyperparams
  { learningRate :: Double,
    beta1 :: Double,
    beta2 :: Double,
    epsilon :: Double,
    numSteps :: Int
  }

microHyperparams :: Hyperparams
microHyperparams =
  Hyperparams
    { learningRate = 0.01,
      beta1 = 0.85,
      beta2 = 0.99,
      epsilon = 1e-8,
      numSteps = 1000
    }

-- | A parameter with its Adam first and second moment buffers
data Param = Param
  { value :: Double,
    m :: Double,
    v :: Double
  }
  deriving stock (Generic)
  deriving anyclass (NFData)

-- | Unboxed training state storing flat parameter and Adam moment buffers.
data TrainState = TrainState
  { shapeTemplate :: Weights (),
    valuesVec :: UVector.Vector Double,
    mVec :: UVector.Vector Double,
    vVec :: UVector.Vector Double,
    step :: Int
  }
  deriving stock (Generic)
  deriving anyclass (NFData)

unpackParamWeights ::
  Weights () ->
  UVector.Vector Double ->
  UVector.Vector Double ->
  UVector.Vector Double ->
  Weights Param
unpackParamWeights tmpl vals ms vs =
  snd $
    Traversable.mapAccumL
      ( \i () ->
          ( i + 1,
            Param
              { value = UVector.unsafeIndex vals i,
                m = UVector.unsafeIndex ms i,
                v = UVector.unsafeIndex vs i
              }
          )
      )
      0
      tmpl

instance HasField "params" TrainState (Weights Param) where
  getField st = unpackParamWeights st.shapeTemplate st.valuesVec st.mVec st.vVec

initTrainState :: Weights Double -> TrainState
initTrainState w =
  let !vals = flattenWeightsUVec w
      !zeros = UVector.replicate (UVector.length vals) 0.0
   in TrainState
        { shapeTemplate = Functor.void w,
          valuesVec = vals,
          mVec = zeros,
          vVec = zeros,
          step = 0
        }

packAdamStepParams :: Double -> Hyperparams -> Int -> A.Vector Double
packAdamStepParams lr hp t =
  A.fromFunction
    (A.Z :. 5)
    ( \case
        A.Z :. 0 -> lr
        A.Z :. 1 -> hp.beta1
        A.Z :. 2 -> hp.beta2
        A.Z :. 3 -> hp.epsilon
        _ -> fromIntegral t
    )

-- | Forward, backward and Adam update on a document, returns the loss before the update
trainStep :: Config -> Hyperparams -> TrainState -> Vector Int -> (Double, TrainState)
trainStep cfg hp st tokens
  | shouldUseAccelerate cfg =
      let !cfgVec = packConfig cfg
          nPairs = min cfg.blockSize (max 0 (Vector.length tokens - 1))
          lr = hp.learningRate * (1.0 - fromIntegral st.step / fromIntegral hp.numSteps)
          !hpVec = packAdamStepParams lr hp (st.step + 1)
          valsAcc = uvecToAcc st.valuesVec
          mAcc = uvecToAcc st.mVec
          vAcc = uvecToAcc st.vVec
       in if nPairs <= 0
            then
              let A.Z :. nParams = A.arrayShape valsAcc
                  !zeroGrads = A.fromFunction (A.Z :. nParams) (const 0.0)
                  (vals', ms', vs') = compiledAdamStep hpVec zeroGrads valsAcc mAcc vAcc
               in ( 0.0,
                    st
                      { valuesVec = accToUVec vals',
                        mVec = accToUVec ms',
                        vVec = accToUVec vs',
                        step = st.step + 1
                      }
                  )
            else
              let (!inToks, !tgtToks) = packTokenPairs nPairs tokens
               in if cfg.nLayer == 1
                    then
                      let (lossS, vals', ms', vs') =
                            compiledTrainStep1 cfgVec hpVec valsAcc mAcc vAcc inToks tgtToks
                       in ( unScalar lossS,
                            st
                              { valuesVec = accToUVec vals',
                                mVec = accToUVec ms',
                                vVec = accToUVec vs',
                                step = st.step + 1
                              }
                          )
                    else
                      let (lossS, dFlat) = compiledFwdBwdGen cfgVec valsAcc inToks tgtToks
                          (vals', ms', vs') = compiledAdamStep hpVec dFlat valsAcc mAcc vAcc
                       in ( unScalar lossS,
                            st
                              { valuesVec = accToUVec vals',
                                mVec = accToUVec ms',
                                vVec = accToUVec vs',
                                step = st.step + 1
                              }
                          )
  | otherwise =
      let nPairs = min cfg.blockSize (max 0 (Vector.length tokens - 1))
          nParams = UVector.length st.valuesVec
          !lr = hp.learningRate * (1.0 - fromIntegral st.step / fromIntegral hp.numSteps)
          !nextStep = st.step + 1
          (!loss, !vals', !ms', !vs') = runST $ do
            valsM <- UVector.thaw st.valuesVec
            mM <- UVector.thaw st.mVec
            vM <- UVector.thaw st.vVec
            docLoss <-
              if nPairs <= 0 || not (isValidModelShape cfg nParams)
                then do
                  zeroG <- UMVector.replicate nParams 0.0
                  stAdamUpdate hp lr nextStep nParams zeroG valsM mM vM
                  pure 0.0
                else do
                  ws <- newSTWorkspace cfg nParams
                  l <- stForwardBackward cfg ws valsM tokens nPairs
                  stAdamUpdate hp lr nextStep nParams ws.gradParamsBuf valsM mM vM
                  pure l
            vFrozen <- UVector.unsafeFreeze valsM
            mFrozen <- UVector.unsafeFreeze mM
            vvFrozen <- UVector.unsafeFreeze vM
            pure (docLoss, vFrozen, mFrozen, vvFrozen)
       in ( loss,
            st
              { valuesVec = vals',
                mVec = ms',
                vVec = vs',
                step = nextStep
              }
          )

-- | One training step per document, reusing mutable 'ST' buffers across the entire batch
-- (or executing an Accelerate @awhile@ loop for large embedding dimensions).
train :: Config -> Hyperparams -> TrainState -> Vector (Vector Int) -> (Vector Double, TrainState)
train cfg hp st0 docs
  | Vector.null docs = (Vector.empty, st0)
  | shouldUseAccelerate cfg && cfg.nLayer == 1 =
      let !cfgVec = packConfig cfg
          numDocs = Vector.length docs
          maxCols = cfg.blockSize + 1
          !docsMat =
            A.fromFunction
              (A.Z :. numDocs :. maxCols)
              (\(A.Z :. i :. j) -> Maybe.fromMaybe 0 ((docs Vector.!? i) >>= (Vector.!? j)))
          !pairCounts =
            A.fromFunction
              (A.Z :. numDocs)
              (\(A.Z :. i) -> maybe 0 (\d -> min cfg.blockSize (max 0 (Vector.length d - 1))) (docs Vector.!? i))
          !hpFloats =
            A.fromFunction
              (A.Z :. 4)
              ( \case
                  A.Z :. 0 -> hp.learningRate
                  A.Z :. 1 -> hp.beta1
                  A.Z :. 2 -> hp.beta2
                  _ -> hp.epsilon
              )
          !stepInts =
            A.fromFunction
              (A.Z :. 3)
              ( \case
                  A.Z :. 0 -> st0.step
                  A.Z :. 1 -> hp.numSteps
                  _ -> numDocs
              )
          (lossesAcc, vals', ms', vs') =
            compiledTrainBatch1
              cfgVec
              hpFloats
              stepInts
              docsMat
              pairCounts
              (uvecToAcc st0.valuesVec)
              (uvecToAcc st0.mVec)
              (uvecToAcc st0.vVec)
          !lossesVec = Vector.generate numDocs (\i -> A.indexArray lossesAcc (A.Z :. i))
          !finalState =
            st0
              { valuesVec = accToUVec vals',
                mVec = accToUVec ms',
                vVec = accToUVec vs',
                step = st0.step + numDocs
              }
       in (lossesVec, finalState)
  | shouldUseAccelerate cfg =
      let (final, losses) =
            Traversable.mapAccumL
              ( \st doc ->
                  let (loss, st') = trainStep cfg hp st doc
                   in st' `deepseq` (st', loss)
              )
              st0
              docs
       in (losses, final)
  | otherwise =
      let numDocs = Vector.length docs
          nParams = UVector.length st0.valuesVec
          validShape = isValidModelShape cfg nParams
          (!lossesVec, !vals', !ms', !vs') = runST $ do
            valsM <- UVector.thaw st0.valuesVec
            mM <- UVector.thaw st0.mVec
            vM <- UVector.thaw st0.vVec
            ws <- newSTWorkspace cfg nParams
            lossesM <- MVector.new numDocs
            let stepDoc idx
                  | idx >= numDocs = pure ()
                  | otherwise = do
                      let doc = Vector.unsafeIndex docs idx
                          nPairs = min cfg.blockSize (max 0 (Vector.length doc - 1))
                          !curStep = st0.step + idx
                          !lr = hp.learningRate * (1.0 - fromIntegral curStep / fromIntegral hp.numSteps)
                          !t = curStep + 1
                      docLoss <-
                        if nPairs <= 0 || not validShape
                          then do
                            stZeroSlice ws.gradParamsBuf 0 nParams
                            stAdamUpdate hp lr t nParams ws.gradParamsBuf valsM mM vM
                            pure 0.0
                          else do
                            l <- stForwardBackward cfg ws valsM doc nPairs
                            stAdamUpdate hp lr t nParams ws.gradParamsBuf valsM mM vM
                            pure l
                      MVector.unsafeWrite lossesM idx docLoss
                      stepDoc (idx + 1)
            stepDoc 0
            lossesFrozen <- Vector.unsafeFreeze lossesM
            valsFrozen <- UVector.unsafeFreeze valsM
            mFrozen <- UVector.unsafeFreeze mM
            vFrozen <- UVector.unsafeFreeze vM
            pure (lossesFrozen, valsFrozen, mFrozen, vFrozen)
       in ( lossesVec,
            st0
              { valuesVec = vals',
                mVec = ms',
                vVec = vs',
                step = st0.step + numDocs
              }
          )

-- 6. Inference & Sampling

-- | Single-token autoregressive inference step with KV-cache update inside Accelerate.
accInferStep ::
  A.Acc (A.Vector Int) ->
  A.Acc (A.Vector Int) ->
  A.Acc (A.Scalar Double) ->
  A.Acc (A.Vector Double) ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Array A.DIM3 Double) ->
  A.Acc (A.Vector Double, A.Array A.DIM3 Double, A.Array A.DIM3 Double)
accInferStep cfgVec stepVec invTemp flatW kCache0 vCache0 =
  let (vSize, lCount, dDim, bSize, nHead, hd) = unpackAccConfig cfgVec
      invScale = sqrt (A.fromIntegral hd) ** (-1.0)
      tok = stepVec A.! A.I1 0
      pos = stepVec A.! A.I1 1
      validLen = pos + 1
      wte = sliceMatrix flatW 0 vSize dDim
      wpe = sliceMatrix flatW (vSize * dDim) bSize dDim
      lmHead = sliceMatrix flatW ((vSize + bSize) * dDim) vSize dDim
      safeTok = A.max 0 (A.min (vSize - 1) tok)
      safePos = A.max 0 (A.min (bSize - 1) pos)
      x0 =
        A.generate (A.I1 dDim) $ \(A.I1 d) ->
          let tv = (0 A.<= tok A.&& tok A.< vSize) A.? (wte A.! A.I2 safeTok d, 0.0)
              pv = (0 A.<= pos A.&& pos A.< bSize) A.? (wpe A.! A.I2 safePos d, 0.0)
           in tv + pv
      x1 = accRmsNorm1D x0
      A.T4 _ xFinal kCacheFinal vCacheFinal =
        A.awhile
          (\(A.T4 l _ _ _) -> A.map (A.< lCount) l)
          ( \(A.T4 l x kCache vCache) ->
              let lIdx = A.the l
                  (wQKV, wo, fc1, fc2) = sliceLayerMatrices vSize bSize dDim flatW lIdx
                  xNorm1 = accRmsNorm1D x
                  qkv = accMatVec wQKV xNorm1
                  kCache' =
                    A.generate (A.I3 lCount bSize dDim) $ \(A.I3 li pi' d) ->
                      (li A.== lIdx A.&& pi' A.== pos) A.? (qkv A.! A.I1 (dDim + d), kCache A.! A.I3 li pi' d)
                  vCache' =
                    A.generate (A.I3 lCount bSize dDim) $ \(A.I3 li pi' d) ->
                      (li A.== lIdx A.&& pi' A.== pos) A.? (qkv A.! A.I1 (2 * dDim + d), vCache A.! A.I3 li pi' d)
                  scoreTerms =
                    A.generate (A.I3 nHead validLen hd) $ \(A.I3 h s c) ->
                      (qkv A.! A.I1 (h * hd + c)) * (kCache' A.! A.I3 lIdx s (h * hd + c))
                  scores = A.map (* invScale) (A.sum scoreTerms)
                  lseScores = A.fold1 combineMaxSumExp (A.map (`A.T2` 1.0) scores)
                  attnW =
                    A.generate (A.I2 nHead validLen) $ \(A.I2 h s) ->
                      let A.T2 m sumE = lseScores A.! A.I1 h
                       in exp (scores A.! A.I2 h s - m) * (sumE ** (-1.0))
                  valTerms =
                    A.generate (A.I3 nHead hd validLen) $ \(A.I3 h c s) ->
                      (attnW A.! A.I2 h s) * (vCache' A.! A.I3 lIdx s (h * hd + c))
                  headOut = A.sum valTerms
                  attnOut = accMatVec wo (A.flatten headOut)
                  xMid = A.zipWith (+) attnOut x
                  xNorm2 = accRmsNorm1D xMid
                  hidden = A.map (\z -> (z A.> 0.0) A.? (z, 0.0)) (accMatVec fc1 xNorm2)
                  mlpOut = accMatVec fc2 hidden
                  xOut = A.zipWith (+) mlpOut xMid
               in A.T4 (A.map (+ 1) l) xOut kCache' vCache'
          )
          (A.T4 (A.unit 0) x1 kCache0 vCache0)
      logits = accMatVec lmHead xFinal
      scaled = A.map (* A.the invTemp) logits
      lseL = A.fold1 combineMaxSumExp (A.map (`A.T2` 1.0) scaled)
      A.T2 maxL sumEL = A.the lseL
      probs = A.map (\z -> exp (z - maxL) * (sumEL ** (-1.0))) scaled
   in A.T3 probs kCacheFinal vCacheFinal

{-# NOINLINE compiledInferStep #-}
compiledInferStep ::
  A.Vector Int ->
  A.Vector Int ->
  A.Scalar Double ->
  A.Vector Double ->
  A.Array A.DIM3 Double ->
  A.Array A.DIM3 Double ->
  (A.Vector Double, A.Array A.DIM3 Double, A.Array A.DIM3 Double)
compiledInferStep = CPU.runN accInferStep

-- | random.choices: samples an index proportionally to the weights without intermediate allocations.
sampleCategorical :: Vector Double -> StdGen -> (Int, StdGen)
sampleCategorical weights gen =
  let (u, gen') = Random.randomR (0.0, 1.0) gen
      n = Vector.length weights
      total = Vector.foldl' (+) 0.0 weights
      !threshold = u * total
      scan idx !cum
        | idx >= n - 1 = max 0 (n - 1)
        | otherwise =
            let !cum' = cum + Vector.unsafeIndex weights idx
             in if cum' <= threshold
                  then scan (idx + 1) cum'
                  else idx
   in (scan 0 0.0, gen')

{-# INLINE sampleCategoricalM #-}
sampleCategoricalM :: Int -> UMVector.MVector s Double -> StdGen -> ST s (Int, StdGen)
sampleCategoricalM n weightsM gen = do
  let (u, gen') = Random.randomR (0.0, 1.0) gen
      sumLoop i !acc
        | i >= n = pure acc
        | otherwise = do
            w <- UMVector.unsafeRead weightsM i
            sumLoop (i + 1) (acc + w)
  total <- sumLoop 0 0.0
  let !threshold = u * total
      scan idx !cum
        | idx >= n - 1 = pure (max 0 (n - 1))
        | otherwise = do
            w <- UMVector.unsafeRead weightsM idx
            let !cum' = cum + w
            if cum' <= threshold
              then scan (idx + 1) cum'
              else pure idx
  chosen <- scan 0 0.0
  pure (chosen, gen')

-- | Autoregressive generation mutating the 3D KV-cache directly in-place in 'ST'.
stGenerate :: Config -> Int -> Double -> UVector.Vector Double -> StdGen -> (Vector Int, StdGen)
stGenerate cfg bosId temperature flatW gen0
  | cfg.blockSize <= 0 || not (isValidModelShape cfg (UVector.length flatW)) = (Vector.empty, gen0)
  | otherwise = runST $ do
      let vSize = cfg.vocabSize
          lCount = cfg.nLayer
          dDim = cfg.nEmbd
          bSize = cfg.blockSize
          nHead = cfg.nHead
          hd = headDim cfg
          mlpDim = 4 * dDim
          dd = dDim * dDim
          offWte = 0
          offWpe = vSize * dDim
          offLmHead = (vSize + bSize) * dDim
          baseLayers = (2 * vSize + bSize) * dDim
          !invScale = 1.0 / sqrt (fromIntegral hd)
          !invTemp = temperature ** (-1.0)
          kvLen = max 1 (lCount * bSize * dDim)
          !oKCache = 0
          !oVCache = oKCache + kvLen
          !oXCur = oVCache + kvLen
          !oXNorm = oXCur + dDim
          !oQVec = oXNorm + dDim
          !oAttnConcat = oQVec + dDim
          !oAttnOut = oAttnConcat + dDim
          !oHidden = oAttnOut + dDim
          !oScores = oHidden + mlpDim
          !oProbs = oScores + bSize
          !totalArena = oProbs + vSize
      arena <- UMVector.new totalArena
      let kCache = UMVector.unsafeSlice oKCache kvLen arena
          vCache = UMVector.unsafeSlice oVCache kvLen arena
          xCur = UMVector.unsafeSlice oXCur dDim arena
          xNorm = UMVector.unsafeSlice oXNorm dDim arena
          qVec = UMVector.unsafeSlice oQVec dDim arena
          attnConcat = UMVector.unsafeSlice oAttnConcat dDim arena
          attnOut = UMVector.unsafeSlice oAttnOut dDim arena
          hidden = UMVector.unsafeSlice oHidden mlpDim arena
          scoresRow = UMVector.unsafeSlice oScores bSize arena
          probs = UMVector.unsafeSlice oProbs vSize arena
      outTokens <- MVector.new bSize
      let stepPos posId tokenId !gen
            | posId >= bSize = do
                frozen <- Vector.unsafeFreeze outTokens
                pure (Vector.take posId frozen, gen)
            | otherwise = do
                let initX d
                      | d >= dDim = pure ()
                      | otherwise = do
                          let !tv =
                                if tokenId >= 0 && tokenId < vSize
                                  then UVector.unsafeIndex flatW (offWte + tokenId * dDim + d)
                                  else 0.0
                              !pv = UVector.unsafeIndex flatW (offWpe + posId * dDim + d)
                          UMVector.unsafeWrite xCur d (tv + pv)
                          initX (d + 1)
                initX 0
                stRmsNorm1D dDim xCur xCur
                let layerLoop l
                      | l >= lCount = pure ()
                      | otherwise = do
                          let !lBase = baseLayers + l * 12 * dd
                              !offWq = lBase
                              !offWk = lBase + dd
                              !offWv = lBase + 2 * dd
                              !offWo = lBase + 3 * dd
                              !offFc1 = lBase + 4 * dd
                              !offFc2 = lBase + 8 * dd
                              !kvSlot = (l * bSize + posId) * dDim
                          stRmsNorm1D dDim xCur xNorm
                          stMatVecImmW dDim dDim flatW offWq xNorm qVec 0
                          stMatVecImmW dDim dDim flatW offWk xNorm kCache kvSlot
                          stMatVecImmW dDim dDim flatW offWv xNorm vCache kvSlot
                          stZeroSlice attnConcat 0 dDim
                          let headLoop h
                                | h >= nHead = pure ()
                                | otherwise = do
                                    let !hOff = h * hd
                                        scoreS s !maxSc
                                          | s > posId = pure maxSc
                                          | otherwise = do
                                              let !kOff = (l * bSize + s) * dDim + hOff
                                                  dotC c !acc
                                                    | c >= hd = pure acc
                                                    | otherwise = do
                                                        qv <- UMVector.unsafeRead qVec (hOff + c)
                                                        kv <- UMVector.unsafeRead kCache (kOff + c)
                                                        dotC (c + 1) (acc + qv * kv)
                                              raw <- dotC 0 0.0
                                              let !sc = raw * invScale
                                              UMVector.unsafeWrite scoresRow s sc
                                              scoreS (s + 1) (max maxSc sc)
                                    maxSc <- scoreS 0 (-1.0e308)
                                    let expS s !sumE
                                          | s > posId = pure sumE
                                          | otherwise = do
                                              sc <- UMVector.unsafeRead scoresRow s
                                              let !ev = exp (sc - maxSc)
                                              UMVector.unsafeWrite scoresRow s ev
                                              expS (s + 1) (sumE + ev)
                                    sumE <- expS 0 0.0
                                    let !invSum = 1.0 / sumE
                                        valS s
                                          | s > posId = pure ()
                                          | otherwise = do
                                              ev <- UMVector.unsafeRead scoresRow s
                                              let !wS = ev * invSum
                                                  !vOff = (l * bSize + s) * dDim + hOff
                                                  addC c
                                                    | c >= hd = pure ()
                                                    | otherwise = do
                                                        vv <- UMVector.unsafeRead vCache (vOff + c)
                                                        stAddAt attnConcat (hOff + c) (wS * vv)
                                                        addC (c + 1)
                                              addC 0
                                              valS (s + 1)
                                    valS 0
                                    headLoop (h + 1)
                          headLoop 0
                          stMatVecImmW dDim dDim flatW offWo attnConcat attnOut 0
                          let addRes1 d
                                | d >= dDim = pure ()
                                | otherwise = do
                                    ao <- UMVector.unsafeRead attnOut d
                                    stAddAt xCur d ao
                                    addRes1 (d + 1)
                          addRes1 0
                          stRmsNorm1D dDim xCur xNorm
                          stMatVecImmW mlpDim dDim flatW offFc1 xNorm hidden 0
                          let reluLoop m
                                | m >= mlpDim = pure ()
                                | otherwise = do
                                    hv <- UMVector.unsafeRead hidden m
                                    when (hv < 0.0) $ UMVector.unsafeWrite hidden m 0.0
                                    reluLoop (m + 1)
                          reluLoop 0
                          stMatVecImmW dDim mlpDim flatW offFc2 hidden attnOut 0
                          let addRes2 d
                                | d >= dDim = pure ()
                                | otherwise = do
                                    mo <- UMVector.unsafeRead attnOut d
                                    stAddAt xCur d mo
                                    addRes2 (d + 1)
                          addRes2 0
                          layerLoop (l + 1)
                layerLoop 0
                stMatVecImmW vSize dDim flatW offLmHead xCur probs 0
                let scaleMax v !maxL
                      | v >= vSize = pure maxL
                      | otherwise = do
                          lv <- UMVector.unsafeRead probs v
                          let !sv = lv * invTemp
                          UMVector.unsafeWrite probs v sv
                          scaleMax (v + 1) (max maxL sv)
                maxL <- scaleMax 0 (-1.0e308)
                let expV v !sumE
                      | v >= vSize = pure sumE
                      | otherwise = do
                          sv <- UMVector.unsafeRead probs v
                          let !ev = exp (sv - maxL)
                          UMVector.unsafeWrite probs v ev
                          expV (v + 1) (sumE + ev)
                sumE <- expV 0 0.0
                let !invSum = 1.0 / sumE
                    normV v
                      | v >= vSize = pure ()
                      | otherwise = do
                          ev <- UMVector.unsafeRead probs v
                          UMVector.unsafeWrite probs v (ev * invSum)
                          normV (v + 1)
                normV 0
                (nextToken, gen') <- sampleCategoricalM vSize probs gen
                if nextToken == bosId
                  then do
                    frozen <- Vector.unsafeFreeze outTokens
                    pure (Vector.take posId frozen, gen')
                  else do
                    MVector.unsafeWrite outTokens posId nextToken
                    stepPos (posId + 1) nextToken gen'
      stepPos 0 bosId gen0

-- | Samples a document, starting from BOS, until BOS is predicted or block_size tokens are generated
generate :: Config -> Int -> Double -> Weights Double -> StdGen -> (Vector Int, StdGen)
generate cfg bosId temperature pureWeights
  | shouldUseAccelerate cfg =
      let !cfgVec = packConfig cfg
          !flatW = packWeights pureWeights
          !invTemp = scalarOf (temperature ** (-1.0))
          !emptyKV = A.fromFunction (A.Z :. cfg.nLayer :. cfg.blockSize :. cfg.nEmbd) (const 0.0)
          go posId tokenId kCache vCache acc gen
            | posId >= cfg.blockSize || cfg.vocabSize <= 0 = (acc, gen)
            | otherwise =
                let !stepVec =
                      A.fromFunction
                        (A.Z :. 2)
                        ( \case
                            A.Z :. 0 -> tokenId
                            _ -> posId
                        )
                    (probsAcc, kCache', vCache') =
                      compiledInferStep cfgVec stepVec invTemp flatW kCache vCache
                    probsVec = Vector.generate cfg.vocabSize (\i -> A.indexArray probsAcc (A.Z :. i))
                    (nextToken, gen') = sampleCategorical probsVec gen
                 in if nextToken == bosId
                      then (acc, gen')
                      else go (posId + 1) nextToken kCache' vCache' (Vector.snoc acc nextToken) gen'
       in go 0 bosId emptyKV emptyKV Vector.empty
  | otherwise =
      let !flatU = flattenWeightsUVec pureWeights
       in stGenerate cfg bosId temperature flatU

-- 7. The whole microgpt.py script

-- | In-place Fisher-Yates shuffle on a mutable vector in 'ST'.
shuffle :: Vector a -> StdGen -> (Vector a, StdGen)
shuffle xs gen0 = runST $ do
  let n = Vector.length xs
  mvec <- Vector.thaw xs
  let loop i !gen
        | i <= 0 = pure gen
        | otherwise = do
            let (j, gen') = Random.randomR (0, i) gen
            MVector.unsafeSwap mvec i j
            loop (i - 1) gen'
  genFinal <- loop (n - 1) gen0
  frozen <- Vector.unsafeFreeze mvec
  pure (frozen, genFinal)

microGPT :: IO ()
microGPT = do
  allDocs <- loadDocs datasetPath
  let (docs, gen1) = shuffle allDocs (Random.mkStdGen 42)
  putStrLn $ "num docs: " <> show (Vector.length docs)

  let tok = mkTokenizer docs
  putStrLn $ "vocab size: " <> show tok.vocabSize

  let cfg = microConfig tok.vocabSize
      hp = microHyperparams
      (w, gen2) = randomWeights cfg gen1
      padded n x = let s = show x in replicate (n - length s) ' ' <> s
  putStrLn $ "num params: " <> show (length w)

  let stepDocs
        | Vector.null docs = Vector.empty
        | otherwise =
            Vector.generate hp.numSteps $ \i ->
              Maybe.fromMaybe "" (docs Vector.!? (i `mod` Vector.length docs))
      st0 = initTrainState w

  final <-
    if shouldUseAccelerate cfg
      then
        Vector.ifoldM
          ( \st i doc -> do
              let (loss, st') = trainStep cfg hp st (encode tok doc)
              st' `deepseq` putStr ("step " <> padded 4 (i + 1) <> " / " <> padded 4 hp.numSteps <> " | loss " <> Numeric.showFFloat (Just 4) loss "\r")
              hFlush stdout
              pure st'
          )
          st0
          stepDocs
      else do
        let nParams = UVector.length st0.valuesVec
            validShape = isValidModelShape cfg nParams
        (valsM, mM, vM, ws) <- stToIO $ do
          valsBuf <- UVector.thaw st0.valuesVec
          mBuf <- UVector.thaw st0.mVec
          vBuf <- UVector.thaw st0.vVec
          wsBuf <- newSTWorkspace cfg nParams
          pure (valsBuf, mBuf, vBuf, wsBuf)
        Vector.iforM_ stepDocs $ \i doc -> do
          let tokens = encode tok doc
              nPairs = min cfg.blockSize (max 0 (Vector.length tokens - 1))
              !lr = hp.learningRate * (1.0 - fromIntegral i / fromIntegral hp.numSteps)
              !t = i + 1
          loss <- stToIO $ do
            if nPairs <= 0 || not validShape
              then do
                stZeroSlice ws.gradParamsBuf 0 nParams
                stAdamUpdate hp lr t nParams ws.gradParamsBuf valsM mM vM
                pure 0.0
              else do
                l <- stForwardBackward cfg ws valsM tokens nPairs
                stAdamUpdate hp lr t nParams ws.gradParamsBuf valsM mM vM
                pure l
          putStr ("step " <> padded 4 (i + 1) <> " / " <> padded 4 hp.numSteps <> " | loss " <> Numeric.showFFloat (Just 4) loss "\r")
          hFlush stdout
        (vals', ms', vs') <- stToIO $ do
          valsFrozen <- UVector.unsafeFreeze valsM
          mFrozen <- UVector.unsafeFreeze mM
          vFrozen <- UVector.unsafeFreeze vM
          pure (valsFrozen, mFrozen, vFrozen)
        pure
          st0
            { valuesVec = vals',
              mVec = ms',
              vVec = vs',
              step = Vector.length stepDocs
            }

  putStrLn "\n--- inference (new, hallucinated names) ---"
  let weights = fmap (.value) final.params
      sampleOne = generate cfg tok.bos 0.5 weights
  Vector.foldM_
    ( \gen i -> do
        let (sample, gen') = sampleOne gen
        putStrLn $ "sample " <> padded 2 i <> ": " <> decode tok sample
        pure gen'
    )
    gen2
    (Vector.enumFromN (1 :: Int) 20)