{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE UnboxedTuples #-}
{-# 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); @cabal run microgpt -- --fp32@
--   trains with the single-precision kernels instead
module Main (main, benchMain, microGPT, parameter, pow, exp', log', relu) where

import Control.Concurrent (forkOn, getNumCapabilities)
import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Control.DeepSeq (NFData (..), deepseq, force)
import Control.Exception (SomeException, evaluate, throwIO, try)
import Control.Monad (forM, forM_, when)
import Control.Monad.ST (RealWorld, ST, runST, stToIO)
import Control.Monad.State.Strict (State, get, put, runState, state)
import qualified Criterion.Main as Criterion
import Data.Bits (unsafeShiftL)
import Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as Char8
import qualified Data.ByteString.Unsafe as ByteStringUnsafe
import qualified Data.Char as Char
import qualified Data.Foldable as Foldable
import qualified Data.Functor as Functor
import qualified Data.IORef as IORef
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Map.Strict as Map
import qualified Data.Maybe as Maybe
import qualified Data.Primitive.ByteArray as ByteArray
import qualified Data.STRef as STRef
import qualified Data.Set as Set
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TextEncoding
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.Primitive as PVector
import qualified Data.Vector.Primitive.Mutable as PMVector
import qualified Data.Vector.Unboxed as UVector
import qualified Data.Vector.Unboxed.Base as UBase
import qualified Data.Vector.Unboxed.Mutable as UMVector
import Data.Word (Word32, Word8)
import GHC.Exts
  ( Double (D#),
    Float (F#),
    Int (I#),
    indexDoubleArray#,
    isTrue#,
    readDoubleArray#,
    writeDoubleArray#,
    (*#),
    (*##),
    (+#),
    (+##),
    (-#),
    (>=#),
  )
import qualified GHC.Exts as Exts
import GHC.Float (castWord32ToFloat, float2Int, int2Float)
import GHC.Generics (Generic)
import GHC.IO (ioToST, unsafePerformIO)
import GHC.Records (HasField (..))
import GHC.ST (ST (..))
import qualified Numeric
import System.Environment (getArgs)
import System.IO (hFlush, hIsTerminalDevice, 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

      nimLookup 0 grads `shouldBe` Just 4.0 -- df/da
      nimLookup 1 grads `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
      nimSize grads `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 (UVector.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)
      UVector.length sample `shouldSatisfy` (<= cfg.blockSize)
      sample `shouldSatisfy` UVector.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"

isAsciiSpace :: Word8 -> Bool
isAsciiSpace b = b == 32 || (b >= 9 && b <= 13)

-- | Loads documents like microgpt.py: one per non-blank line, stripped, as zero-copy 'ByteString' slices.
loadDocs :: FilePath -> IO (Vector ByteString)
loadDocs path = do
  contents <- ByteString.readFile path
  let totalLen = ByteString.length contents
      findLineEnd !i
        | i >= totalLen = totalLen
        | ByteStringUnsafe.unsafeIndex contents i == 10 = i
        | otherwise = findLineEnd (i + 1)
      trimStart !s !lineEnd
        | s < lineEnd && isAsciiSpace (ByteStringUnsafe.unsafeIndex contents s) = trimStart (s + 1) lineEnd
        | otherwise = s
      trimEnd !s !e
        | e > s && isAsciiSpace (ByteStringUnsafe.unsafeIndex contents (e - 1)) = trimEnd s (e - 1)
        | otherwise = e
      countDocs !pos !acc
        | pos >= totalLen = acc
        | otherwise =
            let !lineEnd = findLineEnd pos
                !s = trimStart pos lineEnd
                !e = trimEnd s lineEnd
                !acc' = if e > s then acc + 1 else acc
             in countDocs (lineEnd + 1) acc'
      !numDocs = countDocs 0 0
  pure $ runST $ do
    out <- MVector.new numDocs
    let fillDocs !pos !wIdx
          | pos >= totalLen = pure ()
          | otherwise = do
              let !lineEnd = findLineEnd pos
                  !s = trimStart pos lineEnd
                  !e = trimEnd s lineEnd
              if e > s
                then do
                  let !slice = ByteStringUnsafe.unsafeTake (e - s) (ByteStringUnsafe.unsafeDrop s contents)
                  MVector.unsafeWrite out wIdx slice
                  fillDocs (lineEnd + 1) (wIdx + 1)
                else fillDocs (lineEnd + 1) wIdx
    fillDocs 0 0
    Vector.unsafeFreeze out

-- | 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 = 1000
  w <- evaluate (force (deterministicWeights cfg))
  let flatW = flattenWeightsUVec w
      flatF32 = toFloatVec flatW
  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 ("forward/multicore (" <> show n <> " docs)") $
            Criterion.nfAppIO (forwardLossBatch cfg w) batch,
          Criterion.bench ("forward+backward/multicore (" <> show n <> " docs)") $
            Criterion.nfAppIO (trainingStepBatch cfg w) batch,
          Criterion.bench ("forward/fp32 (" <> show n <> " docs)") $
            Criterion.nfAppIO (forwardLossBatchF32 cfg flatF32) batch,
          Criterion.bench ("forward+backward/fp32 (" <> show n <> " docs)") $
            Criterion.nfAppIO (trainingStepBatchF32 cfg flatF32) 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)

data NodeIdxMap a where
  SparseNodeIdxMap :: !(IntMap.IntMap a) -> NodeIdxMap a
  DenseDoubleMap :: !(UVector.Vector Double) -> NodeIdxMap Double

instance (NFData a) => NFData (NodeIdxMap a) where
  rnf = \case
    SparseNodeIdxMap m -> rnf m
    DenseDoubleMap v -> rnf v

nimEmpty :: NodeIdxMap a
nimEmpty = SparseNodeIdxMap IntMap.empty

nimSize :: NodeIdxMap a -> Int
nimSize = \case
  SparseNodeIdxMap m -> IntMap.size m
  DenseDoubleMap v -> UVector.length v

nimInsert :: NodeId -> a -> NodeIdxMap a -> NodeIdxMap a
nimInsert (NodeId idx) x = \case
  SparseNodeIdxMap m -> SparseNodeIdxMap (IntMap.insert idx x m)
  DenseDoubleMap v ->
    let len = UVector.length v
        toPairs i
          | i >= len = []
          | otherwise = (i, UVector.unsafeIndex v i) : toPairs (i + 1)
     in SparseNodeIdxMap (IntMap.insert idx x (IntMap.fromDistinctAscList (toPairs 0)))

{-# INLINE nimLookup #-}
nimLookup :: NodeId -> NodeIdxMap a -> Maybe a
nimLookup (NodeId k) = \case
  SparseNodeIdxMap m -> IntMap.lookup k m
  DenseDoubleMap v
    | k >= 0 && k < UVector.length v -> Just (UVector.unsafeIndex v k)
    | otherwise -> Nothing

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 (SparseNodeIdxMap (IntMap.fromDistinctAscList (collect 0)))

-- 2. Tensor Primitives (Mutable Unboxed ST)

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

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

-- ============================================================================
-- AVX2/FMA SIMD PrimOp Micro-Kernels (Optimization #2)
-- ============================================================================

{-# INLINE simdDotMImm #-}
simdDotMImm :: ByteArray.ByteArray# -> Exts.Int# -> ByteArray.MutableByteArray# s -> Exts.Int# -> Exts.Int# -> Exts.State# s -> (# Exts.State# s, Double #)
simdDotMImm ba# oW# mba# oX# n# s0 =
  let n4# = n# -# Exts.remInt# n# 4#
      goV i# accV s
        | isTrue# (i# >=# n4#) = (# s, accV #)
        | otherwise =
            let !wv = Exts.indexDoubleArrayAsDoubleX4# ba# (oW# +# i#)
             in case Exts.readDoubleArrayAsDoubleX4# mba# (oX# +# i#) s of
                  (# s', !xv #) ->
                    goV (i# +# 4#) (Exts.plusDoubleX4# accV (Exts.timesDoubleX4# wv xv)) s'
      goS i# !accD s
        | isTrue# (i# >=# n#) = (# s, accD #)
        | otherwise =
            let !wv = Exts.indexDoubleArray# ba# (oW# +# i#)
             in case Exts.readDoubleArray# mba# (oX# +# i#) s of
                  (# s', !xv #) -> goS (i# +# 1#) (accD +## (wv *## xv)) s'
   in case goV 0# (Exts.broadcastDoubleX4# 0.0##) s0 of
        (# s1, accV #) ->
          case Exts.unpackDoubleX4# accV of
            (# a, b, c, d #) ->
              case goS n4# ((a +## b) +## (c +## d)) s1 of
                (# s2, res #) -> (# s2, D# res #)

{-# INLINE simdDotMM #-}
simdDotMM :: ByteArray.MutableByteArray# s -> Exts.Int# -> ByteArray.MutableByteArray# s -> Exts.Int# -> Exts.Int# -> Exts.State# s -> (# Exts.State# s, Double #)
simdDotMM mbaW# oW# mbaX# oX# n# s0 =
  let n4# = n# -# Exts.remInt# n# 4#
      goV i# accV s
        | isTrue# (i# >=# n4#) = (# s, accV #)
        | otherwise =
            case Exts.readDoubleArrayAsDoubleX4# mbaW# (oW# +# i#) s of
              (# s', !wv #) ->
                case Exts.readDoubleArrayAsDoubleX4# mbaX# (oX# +# i#) s' of
                  (# s'', !xv #) ->
                    goV (i# +# 4#) (Exts.plusDoubleX4# accV (Exts.timesDoubleX4# wv xv)) s''
      goS i# !accD s
        | isTrue# (i# >=# n#) = (# s, accD #)
        | otherwise =
            case Exts.readDoubleArray# mbaW# (oW# +# i#) s of
              (# s', !wv #) ->
                case Exts.readDoubleArray# mbaX# (oX# +# i#) s' of
                  (# s'', !xv #) ->
                    goS (i# +# 1#) (accD +## (wv *## xv)) s''
   in case goV 0# (Exts.broadcastDoubleX4# 0.0##) s0 of
        (# s1, accV #) ->
          case Exts.unpackDoubleX4# accV of
            (# a, b, c, d #) ->
              case goS n4# ((a +## b) +## (c +## d)) s1 of
                (# s2, res #) -> (# s2, D# res #)

{-# INLINE simdAxpyM #-}
simdAxpyM :: ByteArray.MutableByteArray# s -> Exts.Int# -> ByteArray.MutableByteArray# s -> Exts.Int# -> Exts.Double# -> Exts.Int# -> Exts.State# s -> Exts.State# s
simdAxpyM mbaY# oY# mbaX# oX# a# n# s0 =
  let n4# = n# -# Exts.remInt# n# 4#
      !a4 = Exts.broadcastDoubleX4# a#
      goV i# s
        | isTrue# (i# >=# n4#) = s
        | otherwise =
            case Exts.readDoubleArrayAsDoubleX4# mbaX# (oX# +# i#) s of
              (# s', !xv #) ->
                case Exts.readDoubleArrayAsDoubleX4# mbaY# (oY# +# i#) s' of
                  (# s'', !yv #) ->
                    let !yv' = Exts.plusDoubleX4# yv (Exts.timesDoubleX4# a4 xv)
                        s''' = Exts.writeDoubleArrayAsDoubleX4# mbaY# (oY# +# i#) yv' s''
                     in goV (i# +# 4#) s'''
      goS i# s
        | isTrue# (i# >=# n#) = s
        | otherwise =
            case Exts.readDoubleArray# mbaX# (oX# +# i#) s of
              (# s', !xv #) ->
                case Exts.readDoubleArray# mbaY# (oY# +# i#) s' of
                  (# s'', !yv #) ->
                    let !yv' = yv +## (a# *## xv)
                        s''' = Exts.writeDoubleArray# mbaY# (oY# +# i#) yv' s''
                     in goS (i# +# 1#) s'''
   in goS n4# (goV 0# s0)

{-# INLINE simdSumSqM #-}
simdSumSqM :: ByteArray.MutableByteArray# s -> Exts.Int# -> Exts.Int# -> Exts.State# s -> (# Exts.State# s, Double #)
simdSumSqM mba# off# n# s0 =
  let n4# = n# -# Exts.remInt# n# 4#
      goV i# accV s
        | isTrue# (i# >=# n4#) = (# s, accV #)
        | otherwise =
            case Exts.readDoubleArrayAsDoubleX4# mba# (off# +# i#) s of
              (# s', !xv #) ->
                goV (i# +# 4#) (Exts.plusDoubleX4# accV (Exts.timesDoubleX4# xv xv)) s'
      goS i# !accD s
        | isTrue# (i# >=# n#) = (# s, accD #)
        | otherwise =
            case Exts.readDoubleArray# mba# (off# +# i#) s of
              (# s', !xv #) -> goS (i# +# 1#) (accD +## (xv *## xv)) s'
   in case goV 0# (Exts.broadcastDoubleX4# 0.0##) s0 of
        (# s1, accV #) ->
          case Exts.unpackDoubleX4# accV of
            (# a, b, c, d #) ->
              case goS n4# ((a +## b) +## (c +## d)) s1 of
                (# s2, res #) -> (# s2, D# res #)

{-# INLINE simdScaleM #-}
simdScaleM :: ByteArray.MutableByteArray# s -> Exts.Int# -> ByteArray.MutableByteArray# s -> Exts.Int# -> Exts.Double# -> Exts.Int# -> Exts.State# s -> Exts.State# s
simdScaleM dst# oDst# src# oSrc# scale# n# s0 =
  let n4# = n# -# Exts.remInt# n# 4#
      !s4 = Exts.broadcastDoubleX4# scale#
      goV i# s
        | isTrue# (i# >=# n4#) = s
        | otherwise =
            case Exts.readDoubleArrayAsDoubleX4# src# (oSrc# +# i#) s of
              (# s', !xv #) ->
                let !yv = Exts.timesDoubleX4# xv s4
                    s'' = Exts.writeDoubleArrayAsDoubleX4# dst# (oDst# +# i#) yv s'
                 in goV (i# +# 4#) s''
      goS i# s
        | isTrue# (i# >=# n#) = s
        | otherwise =
            case Exts.readDoubleArray# src# (oSrc# +# i#) s of
              (# s', !xv #) ->
                let !yv = xv *## scale#
                    s'' = Exts.writeDoubleArray# dst# (oDst# +# i#) yv s'
                 in goS (i# +# 1#) s''
   in goS n4# (goV 0# s0)

{-# INLINE simdNormBwdAddM #-}
simdNormBwdAddM :: ByteArray.MutableByteArray# s -> Exts.Int# -> ByteArray.MutableByteArray# s -> Exts.Int# -> ByteArray.MutableByteArray# s -> Exts.Int# -> Exts.Double# -> Exts.Double# -> Exts.Int# -> Exts.State# s -> Exts.State# s
simdNormBwdAddM dx# oDx# dy# oDy# x# oX# scale# xCoeff# n# s0 =
  let n4# = n# -# Exts.remInt# n# 4#
      !scale4 = Exts.broadcastDoubleX4# scale#
      !xCoeff4 = Exts.broadcastDoubleX4# xCoeff#
      goV i# s
        | isTrue# (i# >=# n4#) = s
        | otherwise =
            case Exts.readDoubleArrayAsDoubleX4# dx# (oDx# +# i#) s of
              (# s1, !dxv #) ->
                case Exts.readDoubleArrayAsDoubleX4# dy# (oDy# +# i#) s1 of
                  (# s2, !dyv #) ->
                    case Exts.readDoubleArrayAsDoubleX4# x# (oX# +# i#) s2 of
                      (# s3, !xv #) ->
                        let !term = Exts.plusDoubleX4# (Exts.timesDoubleX4# scale4 dyv) (Exts.timesDoubleX4# xCoeff4 xv)
                            !res = Exts.plusDoubleX4# dxv term
                            s4 = Exts.writeDoubleArrayAsDoubleX4# dx# (oDx# +# i#) res s3
                         in goV (i# +# 4#) s4
      goS i# s
        | isTrue# (i# >=# n#) = s
        | otherwise =
            case Exts.readDoubleArray# dx# (oDx# +# i#) s of
              (# s1, !dxv #) ->
                case readDoubleArray# dy# (oDy# +# i#) s1 of
                  (# s2, !dyv #) ->
                    case readDoubleArray# x# (oX# +# i#) s2 of
                      (# s3, !xv #) ->
                        let !res = dxv +## ((scale# *## dyv) +## (xCoeff# *## xv))
                            s4 = writeDoubleArray# dx# (oDx# +# i#) res s3
                         in goS (i# +# 1#) s4
   in goS n4# (goV 0# s0)

-- | 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 (I# dDim#) (I# tLen#) (UBase.MV_Double (PMVector.MVector (I# oSrc) _ (ByteArray.MutableByteArray srcMba#))) (I# srcOff#) (UBase.MV_Double (PMVector.MVector (I# oDst) _ (ByteArray.MutableByteArray dstMba#))) (I# dstOff#) (UBase.MV_Double (PMVector.MVector (I# oSc) _ (ByteArray.MutableByteArray scMba#))) (I# scOff#) =
  let !invD = 1.0 / fromIntegral (I# dDim#)
      !baseSrc# = oSrc +# srcOff#
      !baseDst# = oDst +# dstOff#
      !baseSc# = oSc +# scOff#
      loopP p#
        | isTrue# (p# >=# tLen#) = pure ()
        | otherwise = do
            let !rowSrc# = baseSrc# +# (p# *# dDim#)
                !rowDst# = baseDst# +# (p# *# dDim#)
            sq <- ST $ \s -> simdSumSqM srcMba# rowSrc# dDim# s
            let !(D# scale#) = 1.0 / sqrt (sq * invD + 1.0e-5)
            ST $ \s ->
              let s1 = writeDoubleArray# scMba# (baseSc# +# p#) scale# s
                  s2 = simdScaleM dstMba# rowDst# srcMba# rowSrc# scale# dDim# s1
               in (# s2, () #)
            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 (I# dDim#) (UBase.MV_Double (PMVector.MVector (I# oSrc) _ (ByteArray.MutableByteArray srcMba#))) (UBase.MV_Double (PMVector.MVector (I# oDst) _ (ByteArray.MutableByteArray dstMba#))) = do
  let !invD = 1.0 / fromIntegral (I# dDim#)
  sq <- ST $ \s -> simdSumSqM srcMba# oSrc dDim# s
  let !(D# scale#) = 1.0 / sqrt (sq * invD + 1.0e-5)
  ST $ \s ->
    let s' = simdScaleM dstMba# oDst srcMba# oSrc scale# dDim# s
     in (# s', () #)

-- | 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 (I# dDim#) (I# tLen#) (UBase.MV_Double (PMVector.MVector (I# oX) _ (ByteArray.MutableByteArray xMba#))) (I# xOff#) (UBase.MV_Double (PMVector.MVector (I# oSc) _ (ByteArray.MutableByteArray scMba#))) (I# scOff#) (UBase.MV_Double (PMVector.MVector (I# oDy) _ (ByteArray.MutableByteArray dyMba#))) (I# dyOff#) (UBase.MV_Double (PMVector.MVector (I# oDx) _ (ByteArray.MutableByteArray dxMba#))) (I# dxOff#) =
  let !invD = 1.0 / fromIntegral (I# dDim#)
      !baseX# = oX +# xOff#
      !baseSc# = oSc +# scOff#
      !baseDy# = oDy +# dyOff#
      !baseDx# = oDx +# dxOff#
      loopP p#
        | isTrue# (p# >=# tLen#) = pure ()
        | otherwise = do
            let !rowX# = baseX# +# (p# *# dDim#)
                !rowDy# = baseDy# +# (p# *# dDim#)
                !rowDx# = baseDx# +# (p# *# dDim#)
            D# scale# <- ST $ \s ->
              case readDoubleArray# scMba# (baseSc# +# p#) s of
                (# s', v# #) -> (# s', D# v# #)
            dotVal <- ST $ \s -> simdDotMM dyMba# rowDy# xMba# rowX# dDim# s
            let !(D# xCoeff#) = -(dotVal * (D# scale# * D# scale# * D# scale#) * invD)
            ST $ \s ->
              let s' = simdNormBwdAddM dxMba# rowDx# dyMba# rowDy# xMba# rowX# scale# xCoeff# dDim# s
               in (# s', () #)
            loopP (p# +# 1#)
   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 (I# tLen#) (I# nin#) (I# nout#) (UBase.MV_Double (PMVector.MVector (I# oX) _ (ByteArray.MutableByteArray xMba#))) (I# xOff#) (UBase.MV_Double (PMVector.MVector (I# oW) _ (ByteArray.MutableByteArray wMba#))) (I# wOff#) (UBase.MV_Double (PMVector.MVector (I# oY) _ (ByteArray.MutableByteArray yMba#))) (I# yOff#) =
  let !baseX# = oX +# xOff#
      !baseW# = oW +# wOff#
      !baseY# = oY +# yOff#
      loopP p#
        | isTrue# (p# >=# tLen#) = pure ()
        | otherwise = do
            let !rowX# = baseX# +# (p# *# nin#)
                !rowY# = baseY# +# (p# *# nout#)
                loopR r#
                  | isTrue# (r# >=# nout#) = pure ()
                  | otherwise = do
                      let !rowW# = baseW# +# (r# *# nin#)
                      D# val# <- ST $ \s -> simdDotMM wMba# rowW# xMba# rowX# nin# s
                      ST $ \s ->
                        let s' = writeDoubleArray# yMba# (rowY# +# r#) val# s
                         in (# s', () #)
                      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 (I# tLen#) (I# nin#) (I# nout#) (UBase.MV_Double (PMVector.MVector (I# oX) _ (ByteArray.MutableByteArray xMba#))) (I# xOff#) (UBase.MV_Double (PMVector.MVector (I# oW) _ (ByteArray.MutableByteArray wMba#))) (I# wOff#) (UBase.MV_Double (PMVector.MVector (I# oY) _ (ByteArray.MutableByteArray yMba#))) (I# yOff#) =
  let !baseX# = oX +# xOff#
      !baseW# = oW +# wOff#
      !baseY# = oY +# yOff#
      loopP p#
        | isTrue# (p# >=# tLen#) = pure ()
        | otherwise = do
            let !rowX# = baseX# +# (p# *# nin#)
                !rowY# = baseY# +# (p# *# nout#)
                loopR r#
                  | isTrue# (r# >=# nout#) = pure ()
                  | otherwise = do
                      let !rowW# = baseW# +# (r# *# nin#)
                      val <- ST $ \s -> simdDotMM wMba# rowW# xMba# rowX# nin# s
                      let !(D# reluVal#) = max 0.0 val
                      ST $ \s ->
                        let s' = writeDoubleArray# yMba# (rowY# +# r#) reluVal# s
                         in (# s', () #)
                      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 (I# tLen#) (I# nin#) (I# nout#) (UBase.MV_Double (PMVector.MVector (I# oX) _ (ByteArray.MutableByteArray xMba#))) (I# xOff#) (UBase.V_Double (PVector.Vector (I# oW) _ (ByteArray.ByteArray wBa#))) (I# wOff#) (UBase.MV_Double (PMVector.MVector (I# oY) _ (ByteArray.MutableByteArray yMba#))) (I# yOff#) =
  let !baseX# = oX +# xOff#
      !baseW# = oW +# wOff#
      !baseY# = oY +# yOff#
      loopP p#
        | isTrue# (p# >=# tLen#) = pure ()
        | otherwise = do
            let !rowX# = baseX# +# (p# *# nin#)
                !rowY# = baseY# +# (p# *# nout#)
                loopR r#
                  | isTrue# (r# >=# nout#) = pure ()
                  | otherwise = do
                      let !rowW# = baseW# +# (r# *# nin#)
                      D# val# <- ST $ \s -> simdDotMImm wBa# rowW# xMba# rowX# nin# s
                      ST $ \s ->
                        let s' = writeDoubleArray# yMba# (rowY# +# r#) val# s
                         in (# s', () #)
                      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 (I# tLen#) (I# nin#) (I# nout#) (UBase.MV_Double (PMVector.MVector (I# oX) _ (ByteArray.MutableByteArray xMba#))) (I# xOff#) (UBase.V_Double (PVector.Vector (I# oW) _ (ByteArray.ByteArray wBa#))) (I# wOff#) (UBase.MV_Double (PMVector.MVector (I# oY) _ (ByteArray.MutableByteArray yMba#))) (I# yOff#) =
  let !baseX# = oX +# xOff#
      !baseW# = oW +# wOff#
      !baseY# = oY +# yOff#
      loopP p#
        | isTrue# (p# >=# tLen#) = pure ()
        | otherwise = do
            let !rowX# = baseX# +# (p# *# nin#)
                !rowY# = baseY# +# (p# *# nout#)
                loopR r#
                  | isTrue# (r# >=# nout#) = pure ()
                  | otherwise = do
                      let !rowW# = baseW# +# (r# *# nin#)
                      val <- ST $ \s -> simdDotMImm wBa# rowW# xMba# rowX# nin# s
                      let !(D# reluVal#) = max 0.0 val
                      ST $ \s ->
                        let s' = writeDoubleArray# yMba# (rowY# +# r#) reluVal# s
                         in (# s', () #)
                      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 (I# tLen#) (I# nout#) (I# nin#) (UBase.MV_Double (PMVector.MVector (I# oDy) _ (ByteArray.MutableByteArray dyMba#))) (I# dyOff#) (UBase.MV_Double (PMVector.MVector (I# oX) _ (ByteArray.MutableByteArray xMba#))) (I# xOff#) (UBase.MV_Double (PMVector.MVector (I# oW) _ (ByteArray.MutableByteArray wMba#))) (I# wOff#) (UBase.MV_Double (PMVector.MVector (I# oDw) _ (ByteArray.MutableByteArray dwMba#))) (UBase.MV_Double (PMVector.MVector (I# oDx) _ (ByteArray.MutableByteArray dxMba#))) (I# dxOff#) =
  let !baseDy# = oDy +# dyOff#
      !baseX# = oX +# xOff#
      !baseW# = oW +# wOff#
      !baseDw# = oDw +# wOff#
      !baseDx# = oDx +# dxOff#
      loopP p#
        | isTrue# (p# >=# tLen#) = pure ()
        | otherwise = do
            let !rowDy# = baseDy# +# (p# *# nout#)
                !rowX# = baseX# +# (p# *# nin#)
                !rowDx# = baseDx# +# (p# *# nin#)
                loopR r#
                  | isTrue# (r# >=# nout#) = pure ()
                  | otherwise = do
                      D# g# <- ST $ \s ->
                        case readDoubleArray# dyMba# (rowDy# +# r#) s of
                          (# s', v# #) -> (# s', D# v# #)
                      when (D# g# /= 0.0) $ do
                        let !rowW# = baseW# +# (r# *# nin#)
                            !rowDw# = baseDw# +# (r# *# nin#)
                        ST $ \s ->
                          let s1 = simdAxpyM dwMba# rowDw# xMba# rowX# g# nin# s
                              s2 = simdAxpyM dxMba# rowDx# wMba# rowW# g# nin# s1
                           in (# s2, () #)
                      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 (I# nout#) (I# nin#) (UBase.V_Double (PVector.Vector (I# oW) _ (ByteArray.ByteArray wBa#))) (I# wOff#) (UBase.MV_Double (PMVector.MVector (I# oX) _ (ByteArray.MutableByteArray xMba#))) (UBase.MV_Double (PMVector.MVector (I# oY) _ (ByteArray.MutableByteArray yMba#))) (I# yOff#) =
  let !baseW# = oW +# wOff#
      !baseY# = oY +# yOff#
      loopR r#
        | isTrue# (r# >=# nout#) = pure ()
        | otherwise = do
            let !rowW# = baseW# +# (r# *# nin#)
            D# val# <- ST $ \s -> simdDotMImm wBa# rowW# xMba# oX nin# s
            ST $ \s ->
              let s' = writeDoubleArray# yMba# (baseY# +# r#) val# s
               in (# s', () #)
            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

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

-- | 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 -> UVector.Vector Int -> Double
stForwardLoss cfg flatW tokens =
  let nPairs = min cfg.blockSize (max 0 (UVector.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 = UVector.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 = UVector.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 ->
  UVector.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 = UVector.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 = UVector.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 = UVector.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

-- | Forward-only evaluation reusing a preallocated 'STWorkspace'.
{-# INLINE stForwardLossWs #-}
stForwardLossWs :: Config -> STWorkspace s -> UMVector.MVector s Double -> UVector.Vector Int -> Int -> ST s Double
stForwardLossWs 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
      !invScale = 1.0 / sqrt (fromIntegral hd)
      !invT = 1.0 / fromIntegral tLen
  let embedLoop p
        | p >= tLen = pure ()
        | otherwise = do
            let !tok = UVector.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
  let !finalOff = lCount * bd
      lossLoop p !accLoss
        | p >= tLen = pure (accLoss * invT)
        | otherwise = do
            let !tgt = UVector.unsafeIndex tokens (p + 1)
            if tgt < 0 || tgt >= vSize
              then lossLoop (p + 1) accLoss
              else do
                let !rowX = finalOff + 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
                lossLoop (p + 1) (accLoss - log probTgt)
  lossLoop 0 0.0

-- ============================================================================
-- FP32 Single-Precision Specialization & Fast Transcendentals (Optimization #7)
-- ============================================================================

-- | Schraudolph / Cody-Waite single-precision fast exponential function (~2 cycles, no libm call).
{-# INLINE fastExpF32 #-}
fastExpF32 :: Float -> Float
fastExpF32 !x
  | x < (-87.0) = 0.0
  | x > 88.0 = 1.6516363e38
  | otherwise =
      let !a = 12102203.0 :: Float
          !b = 1064866805.0 :: Float
          !val = a * x + b
          !i = float2Int val
       in castWord32ToFloat (fromIntegral i)

{-# INLINE toFloatVec #-}
toFloatVec :: UVector.Vector Double -> UVector.Vector Float
toFloatVec = UVector.map (\(D# d#) -> F# (Exts.double2Float# d#))

-- | Reusable FP32 mutable workspace in 'ST'.
data STWorkspaceF32 s = STWorkspaceF32
  { x0Buf :: UMVector.MVector s Float,
    scale0Buf :: UMVector.MVector s Float,
    xInBuf :: UMVector.MVector s Float,
    scale1Buf :: UMVector.MVector s Float,
    xNorm1Buf :: UMVector.MVector s Float,
    qBuf :: UMVector.MVector s Float,
    kBuf :: UMVector.MVector s Float,
    vBuf :: UMVector.MVector s Float,
    attnWBuf :: UMVector.MVector s Float,
    attnConcatBuf :: UMVector.MVector s Float,
    xMidBuf :: UMVector.MVector s Float,
    scale2Buf :: UMVector.MVector s Float,
    xNorm2Buf :: UMVector.MVector s Float,
    hiddenBuf :: UMVector.MVector s Float,
    probsBuf :: UMVector.MVector s Float,
    gradParamsBuf :: UMVector.MVector s Float,
    dXCurBuf :: UMVector.MVector s Float,
    dXNormBuf :: UMVector.MVector s Float,
    dHiddenBuf :: UMVector.MVector s Float,
    dAttnConcatBuf :: UMVector.MVector s Float,
    dQBuf :: UMVector.MVector s Float,
    dKBuf :: UMVector.MVector s Float,
    dVBuf :: UMVector.MVector s Float,
    dAttnWBuf :: UMVector.MVector s Float
  }

{-# INLINE newSTWorkspaceF32 #-}
newSTWorkspaceF32 :: Config -> Int -> ST s (STWorkspaceF32 s)
newSTWorkspaceF32 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 STWorkspaceF32 {..}

{-# INLINE stZeroSliceF32 #-}
stZeroSliceF32 :: UMVector.MVector s Float -> Int -> Int -> ST s ()
stZeroSliceF32 buf off len = UMVector.set (UMVector.unsafeSlice off len buf) 0.0

{-# INLINE stAddAtF32 #-}
stAddAtF32 :: UMVector.MVector s Float -> Int -> Float -> ST s ()
stAddAtF32 buf idx !delta = do
  old <- UMVector.unsafeRead buf idx
  let !new = old + delta
  UMVector.unsafeWrite buf idx new

{-# INLINE stDotF32 #-}
stDotF32 :: UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> Int -> ST s Float
stDotF32 bufA offA bufB offB n =
  let loop i !acc
        | i >= n = pure acc
        | otherwise = do
            va <- UMVector.unsafeRead bufA (offA + i)
            vb <- UMVector.unsafeRead bufB (offB + i)
            loop (i + 1) (acc + va * vb)
   in loop 0 0.0

{-# INLINE stDotImmwF32 #-}
stDotImmwF32 :: UVector.Vector Float -> Int -> UMVector.MVector s Float -> Int -> Int -> ST s Float
stDotImmwF32 vecA offA bufB offB n =
  let loop i !acc
        | i >= n = pure acc
        | otherwise = do
            let !va = UVector.unsafeIndex vecA (offA + i)
            vb <- UMVector.unsafeRead bufB (offB + i)
            loop (i + 1) (acc + va * vb)
   in loop 0 0.0

{-# INLINE stAxpyF32 #-}
stAxpyF32 :: UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> Float -> Int -> ST s ()
stAxpyF32 dst oDst src oSrc !a n =
  let loop i
        | i >= n = pure ()
        | otherwise = do
            s <- UMVector.unsafeRead src (oSrc + i)
            stAddAtF32 dst (oDst + i) (a * s)
            loop (i + 1)
   in loop 0

{-# INLINE stRmsNormFwdF32 #-}
stRmsNormFwdF32 :: Int -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> ST s ()
stRmsNormFwdF32 dDim tLen xBuf xOff dstBuf dstOff scaleBuf scaleOff =
  let !invD = 1.0 / fromIntegral dDim
      loopP p
        | p >= tLen = pure ()
        | otherwise = do
            let !rowX = xOff + p * dDim
                !rowDst = dstOff + p * dDim
                sumSq d !acc
                  | d >= dDim = pure acc
                  | otherwise = do
                      v <- UMVector.unsafeRead xBuf (rowX + 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 scaleRow d
                  | d >= dDim = pure ()
                  | otherwise = do
                      v <- UMVector.unsafeRead xBuf (rowX + d)
                      UMVector.unsafeWrite dstBuf (rowDst + d) (v * scale)
                      scaleRow (d + 1)
            scaleRow 0
            loopP (p + 1)
   in loopP 0

{-# INLINE stRmsNorm1DF32 #-}
stRmsNorm1DF32 :: Int -> UMVector.MVector s Float -> UMVector.MVector s Float -> ST s ()
stRmsNorm1DF32 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)
  let scaleRow d
        | d >= dDim = pure ()
        | otherwise = do
            v <- UMVector.unsafeRead src d
            UMVector.unsafeWrite dst d (v * scale)
            scaleRow (d + 1)
  scaleRow 0

{-# INLINE stRmsNormBwdAddF32 #-}
stRmsNormBwdAddF32 :: Int -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> ST s ()
stRmsNormBwdAddF32 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 = xOff + p * dDim
                !rowDy = dyOff + p * dDim
                !rowDx = dxOff + p * dDim
            scale <- UMVector.unsafeRead scaleBuf (scaleOff + p)
            dotVal <- stDotF32 dyBuf rowDy xBuf rowX dDim
            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)
                      stAddAtF32 dxDst (rowDx + d) (scale * dy + xCoeff * xv)
                      accumRow (d + 1)
            accumRow 0
            loopP (p + 1)
   in loopP 0

{-# INLINE stLinearFwdF32 #-}
stLinearFwdF32 :: Int -> Int -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> ST s ()
stLinearFwdF32 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
                      val <- stDotF32 wBuf (wOff + r * nin) xBuf rowX nin
                      UMVector.unsafeWrite yBuf (rowY + r) val
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

{-# INLINE stLinearFwdReluF32 #-}
stLinearFwdReluF32 :: Int -> Int -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> ST s ()
stLinearFwdReluF32 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
                      val <- stDotF32 wBuf (wOff + r * nin) xBuf rowX nin
                      UMVector.unsafeWrite yBuf (rowY + r) (max 0.0 val)
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

{-# INLINE stLinearFwdImmwF32 #-}
stLinearFwdImmwF32 :: Int -> Int -> Int -> UMVector.MVector s Float -> Int -> UVector.Vector Float -> Int -> UMVector.MVector s Float -> Int -> ST s ()
stLinearFwdImmwF32 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
                      val <- stDotImmwF32 wVec (wOff + r * nin) xBuf rowX nin
                      UMVector.unsafeWrite yBuf (rowY + r) val
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

{-# INLINE stLinearFwdReluImmwF32 #-}
stLinearFwdReluImmwF32 :: Int -> Int -> Int -> UMVector.MVector s Float -> Int -> UVector.Vector Float -> Int -> UMVector.MVector s Float -> Int -> ST s ()
stLinearFwdReluImmwF32 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
                      val <- stDotImmwF32 wVec (wOff + r * nin) xBuf rowX nin
                      UMVector.unsafeWrite yBuf (rowY + r) (max 0.0 val)
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

{-# INLINE stLinearBwdAccumF32 #-}
stLinearBwdAccumF32 :: Int -> Int -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> UMVector.MVector s Float -> Int -> ST s ()
stLinearBwdAccumF32 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
                        stAxpyF32 dwBuf rowW xBuf rowX g nin
                        stAxpyF32 dxBuf rowDx wBuf rowW g nin
                      loopR (r + 1)
            loopR 0
            loopP (p + 1)
   in loopP 0

{-# INLINE stLinearBwdF32 #-}
stLinearBwdF32 :: Int -> Int -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> Int -> UMVector.MVector s Float -> UMVector.MVector s Float -> Int -> ST s ()
stLinearBwdF32 tLen nout nin dyBuf dyOff xBuf xOff wBuf wOff dwBuf dxBuf dxOff = do
  stZeroSliceF32 dxBuf dxOff (tLen * nin)
  stLinearBwdAccumF32 tLen nout nin dyBuf dyOff xBuf xOff wBuf wOff dwBuf dxBuf dxOff

{-# INLINE stMatVecImmwF32 #-}
stMatVecImmwF32 :: Int -> Int -> UVector.Vector Float -> Int -> UMVector.MVector s Float -> UMVector.MVector s Float -> Int -> ST s ()
stMatVecImmwF32 nout nin wVec wOff xBuf yBuf yOff =
  let loopR r
        | r >= nout = pure ()
        | otherwise = do
            val <- stDotImmwF32 wVec (wOff + r * nin) xBuf 0 nin
            UMVector.unsafeWrite yBuf (yOff + r) val
            loopR (r + 1)
   in loopR 0

{-# INLINE stVectorMaxF32 #-}
stVectorMaxF32 :: UMVector.MVector s Float -> Int -> Int -> ST s Float
stVectorMaxF32 buf off n =
  let loop i !maxV
        | i >= n = pure maxV
        | otherwise = do
            v <- UMVector.unsafeRead buf (off + i)
            loop (i + 1) (max maxV v)
   in loop 0 (-1.0e38)

-- | Forward-only sequence loss reusing a preallocated 'STWorkspaceF32'.
{-# INLINE stForwardLossWsF32 #-}
stForwardLossWsF32 :: Config -> STWorkspaceF32 s -> UMVector.MVector s Float -> UVector.Vector Int -> Int -> ST s Float
stForwardLossWsF32 cfg ws valsM tokens tLen = do
  let STWorkspaceF32 {..} = 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
      !invScale = 1.0 / sqrt (fromIntegral hd)
      !invT = 1.0 / fromIntegral tLen
  let embedLoop p
        | p >= tLen = pure ()
        | otherwise = do
            let !tok = UVector.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
  stRmsNormFwdF32 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
            stRmsNormFwdF32 dDim tLen xInBuf lOff xNorm1Buf lOff scale1Buf lScaleOff
            stLinearFwdF32 tLen dDim dDim xNorm1Buf lOff valsM offWq qBuf lOff
            stLinearFwdF32 tLen dDim dDim xNorm1Buf lOff valsM offWk kBuf lOff
            stLinearFwdF32 tLen dDim dDim xNorm1Buf lOff valsM offWv vBuf lOff
            stZeroSliceF32 attnConcatBuf lOff (tLen * dDim)
            let headLoop h
                  | h >= nHead = pure ()
                  | otherwise = do
                      let !hOff = h * hd
                          posLoop qPos
                            | qPos >= tLen = pure ()
                            | otherwise = do
                                let !qRow = lOff + qPos * dDim + hOff
                                    scoreS s !maxSc
                                      | s > qPos = pure maxSc
                                      | otherwise = do
                                          let !kRow = lOff + s * dDim + hOff
                                          raw <- stDotF32 qBuf qRow kBuf kRow hd
                                          let !sc = raw * invScale
                                              !wIdx = (l * nHead + h) * bSize * bSize + qPos * bSize + s
                                          UMVector.unsafeWrite attnWBuf wIdx sc
                                          scoreS (s + 1) (max maxSc sc)
                                maxSc <- scoreS 0 (-1.0e38)
                                let expS s !sumE
                                      | s > qPos = pure sumE
                                      | otherwise = do
                                          let !wIdx = (l * nHead + h) * bSize * bSize + qPos * bSize + s
                                          sc <- UMVector.unsafeRead attnWBuf wIdx
                                          let !ev = fastExpF32 (sc - maxSc)
                                          UMVector.unsafeWrite attnWBuf wIdx ev
                                          expS (s + 1) (sumE + ev)
                                sumE <- expS 0 0.0
                                let !invSum = 1.0 / sumE
                                    normS s
                                      | s > qPos = pure ()
                                      | otherwise = do
                                          let !wIdx = (l * nHead + h) * bSize * bSize + qPos * bSize + s
                                          ev <- UMVector.unsafeRead attnWBuf wIdx
                                          let !w = ev * invSum
                                              !vRow = lOff + s * dDim + hOff
                                              !outRow = lOff + qPos * dDim + hOff
                                          UMVector.unsafeWrite attnWBuf wIdx w
                                          stAxpyF32 attnConcatBuf outRow vBuf vRow w hd
                                          normS (s + 1)
                                normS 0
                                posLoop (qPos + 1)
                      posLoop 0
                      headLoop (h + 1)
            headLoop 0
            stLinearFwdF32 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
            stRmsNormFwdF32 dDim tLen xMidBuf lOff xNorm2Buf lOff scale2Buf lScaleOff
            stLinearFwdReluF32 tLen dDim mlpDim xNorm2Buf lOff valsM offFc1 hiddenBuf lMlpOff
            stLinearFwdF32 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
  let !finalOff = lCount * bd
      lossLoop p !accLoss
        | p >= tLen = pure (accLoss * invT)
        | otherwise = do
            let !tgt = UVector.unsafeIndex tokens (p + 1)
            if tgt < 0 || tgt >= vSize
              then lossLoop (p + 1) accLoss
              else do
                let !rowX = finalOff + p * dDim
                stLinearFwdF32 1 dDim vSize xInBuf rowX valsM offLmHead probsBuf 0
                maxL <- stVectorMaxF32 probsBuf 0 vSize
                let expV v !sumE
                      | v >= vSize = pure sumE
                      | otherwise = do
                          lv <- UMVector.unsafeRead probsBuf v
                          let !ev = fastExpF32 (lv - maxL)
                          UMVector.unsafeWrite probsBuf v ev
                          expV (v + 1) (sumE + ev)
                sumE <- expV 0 0.0
                tgtExp <- UMVector.unsafeRead probsBuf tgt
                let !prob = tgtExp / sumE
                lossLoop (p + 1) (accLoss - log prob)
  lossLoop 0 0.0

-- | Full forward and backward pass in FP32 mutating @ws.gradParamsBuf@ in-place.
{-# INLINE stForwardBackwardF32 #-}
stForwardBackwardF32 ::
  Config ->
  STWorkspaceF32 s ->
  UMVector.MVector s Float ->
  UVector.Vector Int ->
  Int ->
  ST s Float
stForwardBackwardF32 cfg ws valsM tokens tLen = do
  let STWorkspaceF32 {..} = 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 = UVector.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
  stRmsNormFwdF32 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
            stRmsNormFwdF32 dDim tLen xInBuf lOff xNorm1Buf lOff scale1Buf lScaleOff
            stLinearFwdF32 tLen dDim dDim xNorm1Buf lOff valsM offWq qBuf lOff
            stLinearFwdF32 tLen dDim dDim xNorm1Buf lOff valsM offWk kBuf lOff
            stLinearFwdF32 tLen dDim dDim xNorm1Buf lOff valsM offWv vBuf lOff
            stZeroSliceF32 attnConcatBuf lOff (tLen * dDim)
            let headLoop h
                  | h >= nHead = pure ()
                  | otherwise = do
                      let !hOff = h * hd
                          posLoop qPos
                            | qPos >= tLen = pure ()
                            | otherwise = do
                                let !qRow = lOff + qPos * dDim + hOff
                                    scoreS s !maxSc
                                      | s > qPos = pure maxSc
                                      | otherwise = do
                                          let !kRow = lOff + s * dDim + hOff
                                          raw <- stDotF32 qBuf qRow kBuf kRow hd
                                          let !sc = raw * invScale
                                              !wIdx = (l * nHead + h) * bSize * bSize + qPos * bSize + s
                                          UMVector.unsafeWrite attnWBuf wIdx sc
                                          scoreS (s + 1) (max maxSc sc)
                                maxSc <- scoreS 0 (-1.0e38)
                                let expS s !sumE
                                      | s > qPos = pure sumE
                                      | otherwise = do
                                          let !wIdx = (l * nHead + h) * bSize * bSize + qPos * bSize + s
                                          sc <- UMVector.unsafeRead attnWBuf wIdx
                                          let !ev = fastExpF32 (sc - maxSc)
                                          UMVector.unsafeWrite attnWBuf wIdx ev
                                          expS (s + 1) (sumE + ev)
                                sumE <- expS 0 0.0
                                let !invSum = 1.0 / sumE
                                    normS s
                                      | s > qPos = pure ()
                                      | otherwise = do
                                          let !wIdx = (l * nHead + h) * bSize * bSize + qPos * bSize + s
                                          ev <- UMVector.unsafeRead attnWBuf wIdx
                                          let !w = ev * invSum
                                              !vRow = lOff + s * dDim + hOff
                                              !outRow = lOff + qPos * dDim + hOff
                                          UMVector.unsafeWrite attnWBuf wIdx w
                                          stAxpyF32 attnConcatBuf outRow vBuf vRow w hd
                                          normS (s + 1)
                                normS 0
                                posLoop (qPos + 1)
                      posLoop 0
                      headLoop (h + 1)
            headLoop 0
            stLinearFwdF32 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
            stRmsNormFwdF32 dDim tLen xMidBuf lOff xNorm2Buf lOff scale2Buf lScaleOff
            stLinearFwdReluF32 tLen dDim mlpDim xNorm2Buf lOff valsM offFc1 hiddenBuf lMlpOff
            stLinearFwdF32 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
  stZeroSliceF32 gradParamsBuf 0 totalParams
  stZeroSliceF32 dXCurBuf 0 (tLen * dDim)
  let !finalOff = lCount * bd
      lossAndHeadBwd p !accLoss
        | p >= tLen = pure (accLoss * invT)
        | otherwise = do
            let !tgt = UVector.unsafeIndex tokens (p + 1)
            if tgt < 0 || tgt >= vSize
              then lossAndHeadBwd (p + 1) accLoss
              else do
                let !rowX = finalOff + p * dDim
                stLinearFwdF32 1 dDim vSize xInBuf rowX valsM offLmHead probsBuf 0
                maxL <- stVectorMaxF32 probsBuf 0 vSize
                let expV v !sumE
                      | v >= vSize = pure sumE
                      | otherwise = do
                          lv <- UMVector.unsafeRead probsBuf v
                          let !ev = fastExpF32 (lv - maxL)
                          UMVector.unsafeWrite probsBuf 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 probsBuf v
                          let !prob = ev * invSum
                              !indicator = if v == tgt then 1.0 else 0.0
                              !dLogit = invT * (prob - indicator)
                          UMVector.unsafeWrite probsBuf v dLogit
                          normV (v + 1)
                normV 0
                tgtEv <- UMVector.unsafeRead probsBuf tgt
                let !tgtProb = (tgtEv * fromIntegral tLen) + 1.0
                    !tgtProbClamped = max 1e-12 tgtProb
                    !lossVal = -log tgtProbClamped
                let bwdHeadCol v
                      | v >= vSize = pure ()
                      | otherwise = do
                          dL <- UMVector.unsafeRead probsBuf v
                          when (dL /= 0.0) $ do
                            let !rowW = offLmHead + v * dDim
                            stAxpyF32 gradParamsBuf rowW xInBuf rowX dL dDim
                            stAxpyF32 dXCurBuf (p * dDim) valsM rowW dL dDim
                          bwdHeadCol (v + 1)
                bwdHeadCol 0
                lossAndHeadBwd (p + 1) (accLoss + lossVal)
  lossVal <- 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
            stLinearBwdF32 tLen dDim mlpDim dXCurBuf 0 hiddenBuf lMlpOff valsM offFc2 gradParamsBuf dHiddenBuf 0
            let bwdRelu i
                  | i >= tLen * mlpDim = pure ()
                  | otherwise = do
                      h <- UMVector.unsafeRead hiddenBuf (lMlpOff + i)
                      when (h <= 0.0) $ UMVector.unsafeWrite dHiddenBuf i 0.0
                      bwdRelu (i + 1)
            bwdRelu 0
            stLinearBwdF32 tLen mlpDim dDim dHiddenBuf 0 xNorm2Buf lOff valsM offFc1 gradParamsBuf dXNormBuf 0
            stRmsNormBwdAddF32 dDim tLen xMidBuf lOff scale2Buf lScaleOff dXNormBuf 0 dXCurBuf 0
            stLinearBwdF32 tLen dDim dDim dXCurBuf 0 attnConcatBuf lOff valsM offWo gradParamsBuf dAttnConcatBuf 0
            stZeroSliceF32 dQBuf 0 (tLen * dDim)
            stZeroSliceF32 dKBuf 0 (tLen * dDim)
            stZeroSliceF32 dVBuf 0 (tLen * dDim)
            let headLoop h
                  | h >= nHead = pure ()
                  | otherwise = do
                      let !hOff = h * hd
                          posLoop qPos
                            | qPos >= tLen = pure ()
                            | otherwise = do
                                let !qOff = qPos * dDim + hOff
                                    !dAttnOff = qOff
                                    dotV s !accDot
                                      | s > qPos = pure accDot
                                      | otherwise = do
                                          let !wIdx = (l * nHead + h) * bSize * bSize + qPos * bSize + s
                                          w <- UMVector.unsafeRead attnWBuf wIdx
                                          let !vRow = lOff + s * dDim + hOff
                                          dotProd <- stDotF32 dAttnConcatBuf dAttnOff vBuf vRow hd
                                          stAxpyF32 dVBuf (s * dDim + hOff) dAttnConcatBuf dAttnOff w hd
                                          UMVector.unsafeWrite dAttnWBuf s dotProd
                                          dotV (s + 1) (accDot + dotProd * w)
                                sumDot <- dotV 0 0.0
                                let bwdSoftmax s
                                      | s > qPos = pure ()
                                      | otherwise = do
                                          let !wIdx = (l * nHead + h) * bSize * bSize + qPos * bSize + s
                                          w <- UMVector.unsafeRead attnWBuf wIdx
                                          dp <- UMVector.unsafeRead dAttnWBuf s
                                          let !dScore = (w * (dp - sumDot)) * invScale
                                              !kRow = lOff + s * dDim + hOff
                                          stAxpyF32 dQBuf qOff kBuf kRow dScore hd
                                          stAxpyF32 dKBuf (s * dDim + hOff) qBuf (lOff + qOff) dScore hd
                                          bwdSoftmax (s + 1)
                                bwdSoftmax 0
                                posLoop (qPos + 1)
                      posLoop 0
                      headLoop (h + 1)
            headLoop 0
            stZeroSliceF32 dXNormBuf 0 (tLen * dDim)
            stLinearBwdAccumF32 tLen dDim dDim dQBuf 0 xNorm1Buf lOff valsM offWq gradParamsBuf dXNormBuf 0
            stLinearBwdAccumF32 tLen dDim dDim dKBuf 0 xNorm1Buf lOff valsM offWk gradParamsBuf dXNormBuf 0
            stLinearBwdAccumF32 tLen dDim dDim dVBuf 0 xNorm1Buf lOff valsM offWv gradParamsBuf dXNormBuf 0
            stRmsNormBwdAddF32 dDim tLen xInBuf lOff scale1Buf lScaleOff dXNormBuf 0 dXCurBuf 0
            bwdLayers (l - 1)
  bwdLayers (lCount - 1)
  stZeroSliceF32 dXNormBuf 0 (tLen * dDim)
  stRmsNormBwdAddF32 dDim tLen x0Buf 0 scale0Buf 0 dXCurBuf 0 dXNormBuf 0
  let bwdEmbed p
        | p >= tLen = pure ()
        | otherwise = do
            let !tok = UVector.unsafeIndex tokens p
                !rowDx = p * dDim
                loopD d
                  | d >= dDim = pure ()
                  | otherwise = do
                      dx <- UMVector.unsafeRead dXNormBuf (rowDx + d)
                      when (tok >= 0 && tok < vSize) $
                        stAddAtF32 gradParamsBuf (offWte + tok * dDim + d) dx
                      stAddAtF32 gradParamsBuf (offWpe + p * dDim + d) dx
                      loopD (d + 1)
            loopD 0
            bwdEmbed (p + 1)
  bwdEmbed 0
  pure lossVal

-- | Adam update in FP32 single precision.
{-# INLINE stAdamUpdateF32 #-}
stAdamUpdateF32 :: Hyperparams -> Float -> Int -> Int -> UMVector.MVector s Float -> UMVector.MVector s Float -> UMVector.MVector s Float -> UMVector.MVector s Float -> ST s ()
stAdamUpdateF32 hp lr t nParams gradsM valsM mM vM =
  let !b1 = realToFrac hp.beta1 :: Float
      !b2 = realToFrac hp.beta2 :: Float
      !oneMinusB1 = 1.0 - b1
      !oneMinusB2 = 1.0 - b2
      !tF = fromIntegral t :: Float
      !b1Corr = 1.0 - (b1 ** tF)
      !b2Corr = 1.0 - (b2 ** tF)
      !epsilon = realToFrac hp.epsilon :: Float
      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 + epsilon)
            UMVector.unsafeWrite mM i mNew
            UMVector.unsafeWrite vM i vNew
            UMVector.unsafeWrite valsM i valNew
            loop (i + 1)
   in loop 0

-- | Autoregressive generation in FP32 mutating the 3D KV-cache in 'ST'.
stGenerateF32 :: Config -> Int -> Float -> UVector.Vector Float -> StdGen -> (UVector.Vector Int, StdGen)
stGenerateF32 cfg bosId temperature flatW gen0
  | cfg.blockSize <= 0 || not (isValidModelShape cfg (UVector.length flatW)) = (UVector.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 <- UMVector.new bSize
      let stepPos posId tokenId !gen
            | posId >= bSize = do
                frozen <- UVector.unsafeFreeze outTokens
                pure (UVector.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
                stRmsNorm1DF32 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
                          stRmsNorm1DF32 dDim xCur xNorm
                          stMatVecImmwF32 dDim dDim flatW offWq xNorm qVec 0
                          stMatVecImmwF32 dDim dDim flatW offWk xNorm kCache kvSlot
                          stMatVecImmwF32 dDim dDim flatW offWv xNorm vCache kvSlot
                          stZeroSliceF32 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.0e38)
                                    let expS s !sumE
                                          | s > posId = pure sumE
                                          | otherwise = do
                                              sc <- UMVector.unsafeRead scoresRow s
                                              let !ev = fastExpF32 (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)
                                                        stAddAtF32 attnConcat (hOff + c) (wS * vv)
                                                        addC (c + 1)
                                              addC 0
                                              valS (s + 1)
                                    valS 0
                                    headLoop (h + 1)
                          headLoop 0
                          stMatVecImmwF32 dDim dDim flatW offWo attnConcat attnOut 0
                          let addRes1 d
                                | d >= dDim = pure ()
                                | otherwise = do
                                    ao <- UMVector.unsafeRead attnOut d
                                    stAddAtF32 xCur d ao
                                    addRes1 (d + 1)
                          addRes1 0
                          stRmsNorm1DF32 dDim xCur xNorm
                          stMatVecImmwF32 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
                          stMatVecImmwF32 dDim mlpDim flatW offFc2 hidden attnOut 0
                          let addRes2 d
                                | d >= dDim = pure ()
                                | otherwise = do
                                    mo <- UMVector.unsafeRead attnOut d
                                    stAddAtF32 xCur d mo
                                    addRes2 (d + 1)
                          addRes2 0
                          layerLoop (l + 1)
                layerLoop 0
                stMatVecImmwF32 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.0e38)
                let expV v !sumE
                      | v >= vSize = pure sumE
                      | otherwise = do
                          sv <- UMVector.unsafeRead probs v
                          let !ev = fastExpF32 (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
                let (r, gen') = Random.randomR (0.0 :: Float, 1.0) gen
                    scan idx !cum
                      | idx >= vSize - 1 = pure idx
                      | otherwise = do
                          p <- UMVector.unsafeRead probs idx
                          let !cum' = cum + p
                          if r <= cum'
                            then pure idx
                            else scan (idx + 1) cum'
                chosen <- scan 0 0.0
                UMVector.unsafeWrite outTokens posId chosen
                if chosen == bosId
                  then do
                    frozen <- UVector.unsafeFreeze outTokens
                    pure (UVector.take (posId + 1) frozen, gen')
                  else stepPos (posId + 1) chosen gen'
      stepPos 0 bosId gen0

-- 4. Tokenizer

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

-- | Builds the character vocabulary directly from 'ByteString' documents using an in-place ASCII bitset in 'ST'
-- (with a UTF-8 fallback for non-ASCII code points).
mkTokenizer :: Vector ByteString -> Tokenizer
mkTokenizer docs =
  let uchars = runST $ do
        asciiSeen <- UMVector.replicate 128 False
        extraRef <- STRef.newSTRef Set.empty
        Vector.forM_ docs $ \doc -> do
          let bsLen = ByteString.length doc
              scanBytes !i
                | i >= bsLen = pure ()
                | otherwise = do
                    let !b = ByteStringUnsafe.unsafeIndex doc i
                    if b < 128
                      then do
                        UMVector.unsafeWrite asciiSeen (fromIntegral b) True
                        scanBytes (i + 1)
                      else do
                        let restTxt = TextEncoding.decodeUtf8Lenient (ByteStringUnsafe.unsafeDrop i doc)
                        Text.foldl'
                          ( \act c -> do
                              act
                              let code = Char.ord c
                              if code < 128
                                then UMVector.unsafeWrite asciiSeen code True
                                else STRef.modifySTRef' extraRef (Set.insert c)
                          )
                          (pure ())
                          restTxt
          scanBytes 0
        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 (UVector.fromList (collectAscii 0))
      tokenMap = UVector.ifoldl' (\m i c -> Map.insert c i m) Map.empty uchars
      asciiTable = runST $ do
        tbl <- UMVector.replicate 256 (-1)
        UVector.iforM_ uchars $ \i c -> do
          let code = Char.ord c
          when (code < 128) $ UMVector.unsafeWrite tbl code i
        UVector.unsafeFreeze tbl
      nChars = UVector.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 unboxed token vector in 'ST'.
encode :: Tokenizer -> ByteString -> UVector.Vector Int
encode tok doc = runST $ do
  let bsLen = ByteString.length doc
  buf <- UMVector.new (bsLen + 2)
  UMVector.unsafeWrite buf 0 tok.bos
  let finish !wIdx = do
        UMVector.unsafeWrite buf wIdx tok.bos
        frozen <- UVector.unsafeFreeze buf
        pure (UVector.take (wIdx + 1) frozen)
      fillChars chars !wIdx = case chars of
        [] -> finish wIdx
        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 -> fillChars cs wIdx
            Just tid -> do
              UMVector.unsafeWrite buf wIdx tid
              fillChars cs (wIdx + 1)
      fillBytes !i !wIdx
        | i >= bsLen = finish wIdx
        | otherwise = do
            let !b = ByteStringUnsafe.unsafeIndex doc i
            if b < 128
              then do
                let !tid = UVector.unsafeIndex tok.asciiIds (fromIntegral b)
                if tid >= 0
                  then do
                    UMVector.unsafeWrite buf wIdx tid
                    fillBytes (i + 1) (wIdx + 1)
                  else fillBytes (i + 1) wIdx
              else do
                let restChars = Text.unpack (TextEncoding.decodeUtf8Lenient (ByteStringUnsafe.unsafeDrop i doc))
                fillChars restChars wIdx
  fillBytes 0 1

decode :: Tokenizer -> UVector.Vector Int -> String
decode tok tokens =
  let n = UVector.length tokens
      nChars = UVector.length tok.uchars
      go i
        | i >= n = []
        | otherwise =
            let !tid = UVector.unsafeIndex tokens i
             in if tid >= 0 && tid < nChars
                  then UVector.unsafeIndex tok.uchars tid : go (i + 1)
                  else go (i + 1)
   in go 0

-- 5. Training

{-# INLINE uvecToGradients #-}
uvecToGradients :: UVector.Vector Double -> Gradients
uvecToGradients = DenseDoubleMap

-- | Loss of a document, forward pass only
{-# INLINE forwardLoss #-}
forwardLoss :: Config -> Weights Double -> UVector.Vector Int -> Double
forwardLoss cfg pureWeights =
  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 -> UVector.Vector Int -> (Double, Gradients)
trainingStep cfg pureWeights =
  let flatU = flattenWeightsUVec pureWeights
   in \tokens ->
        let nPairs = min cfg.blockSize (max 0 (UVector.length tokens - 1))
            nParams = UVector.length flatU
         in if nPairs <= 0 || not (isValidModelShape cfg nParams)
              then (0.0, DenseDoubleMap UVector.empty)
              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)

-- ============================================================================
-- Multi-Core Document Batching & Capability Re-Use (Optimization #4)
-- ============================================================================

-- | Splits a vector into at most @k@ balanced contiguous chunks.
splitIntoChunks :: Int -> Vector a -> [Vector a]
splitIntoChunks k vec
  | k <= 1 || Vector.null vec = [vec]
  | otherwise =
      let !n = Vector.length vec
          !chunkSz = (n + k - 1) `div` k
          go i
            | i >= n = []
            | otherwise = Vector.slice i (min chunkSz (n - i)) vec : go (i + chunkSz)
       in go 0

-- | Evaluates forward loss in parallel across CPU capabilities reusing 1 'STWorkspace' per worker.
forwardLossBatch :: Config -> Weights Double -> Vector (UVector.Vector Int) -> IO (Vector Double)
forwardLossBatch cfg pureWeights batch
  | Vector.null batch = pure Vector.empty
  | otherwise = do
      numCaps <- getNumCapabilities
      let flatU = flattenWeightsUVec pureWeights
          n = Vector.length batch
          nParams = UVector.length flatU
      outM <- MVector.new n
      let chunks = splitIntoChunks numCaps (Vector.enumFromN (0 :: Int) n)
      doneMVars <- forM (zip [0 ..] chunks) $ \(cap, idxSlice) -> do
        doneVar <- newEmptyMVar
        _ <- forkOn cap $ do
          stToIO $ do
            ws <- newSTWorkspace cfg nParams
            valsM <- UVector.thaw flatU
            Vector.forM_ idxSlice $ \i -> do
              let tokens = Vector.unsafeIndex batch i
                  nPairs = min cfg.blockSize (max 0 (UVector.length tokens - 1))
              if nPairs <= 0 || not (isValidModelShape cfg nParams)
                then ioToST (MVector.unsafeWrite outM i 0.0)
                else do
                  l <- stForwardLossWs cfg ws valsM tokens nPairs
                  ioToST (MVector.unsafeWrite outM i l)
          putMVar doneVar ()
        pure doneVar
      forM_ doneMVars takeMVar
      Vector.unsafeFreeze outM

-- | Runs training steps (forward + backward) in parallel across CPU capabilities reusing 1 'STWorkspace' per worker.
trainingStepBatch :: Config -> Weights Double -> Vector (UVector.Vector Int) -> IO (Vector (Double, Gradients))
trainingStepBatch cfg pureWeights batch
  | Vector.null batch = pure Vector.empty
  | otherwise = do
      numCaps <- getNumCapabilities
      let flatU = flattenWeightsUVec pureWeights
          n = Vector.length batch
          nParams = UVector.length flatU
      outM <- MVector.new n
      let chunks = splitIntoChunks numCaps (Vector.enumFromN (0 :: Int) n)
      doneMVars <- forM (zip [0 ..] chunks) $ \(cap, idxSlice) -> do
        doneVar <- newEmptyMVar
        _ <- forkOn cap $ do
          stToIO $ do
            ws <- newSTWorkspace cfg nParams
            valsM <- UVector.thaw flatU
            Vector.forM_ idxSlice $ \i -> do
              let tokens = Vector.unsafeIndex batch i
                  nPairs = min cfg.blockSize (max 0 (UVector.length tokens - 1))
              if nPairs <= 0 || not (isValidModelShape cfg nParams)
                then ioToST (MVector.unsafeWrite outM i (0.0, DenseDoubleMap UVector.empty))
                else do
                  l <- stForwardBackward cfg ws valsM tokens nPairs
                  gFrozen <- UVector.freeze ws.gradParamsBuf
                  ioToST (MVector.unsafeWrite outM i (l, uvecToGradients gFrozen))
          putMVar doneVar ()
        pure doneVar
      forM_ doneMVars takeMVar
      Vector.unsafeFreeze outM

-- | Evaluates forward loss in FP32 in parallel across CPU capabilities reusing 1 'STWorkspaceF32' per worker.
forwardLossBatchF32 :: Config -> UVector.Vector Float -> Vector (UVector.Vector Int) -> IO (Vector Float)
forwardLossBatchF32 cfg flatF32 batch
  | Vector.null batch = pure Vector.empty
  | otherwise = do
      numCaps <- getNumCapabilities
      let n = Vector.length batch
          nParams = UVector.length flatF32
      outM <- MVector.new n
      let chunks = splitIntoChunks numCaps (Vector.enumFromN (0 :: Int) n)
      doneMVars <- forM (zip [0 ..] chunks) $ \(cap, idxSlice) -> do
        doneVar <- newEmptyMVar
        _ <- forkOn cap $ do
          stToIO $ do
            ws <- newSTWorkspaceF32 cfg nParams
            valsM <- UVector.thaw flatF32
            Vector.forM_ idxSlice $ \i -> do
              let tokens = Vector.unsafeIndex batch i
                  nPairs = min cfg.blockSize (max 0 (UVector.length tokens - 1))
              if nPairs <= 0 || not (isValidModelShape cfg nParams)
                then ioToST (MVector.unsafeWrite outM i 0.0)
                else do
                  l <- stForwardLossWsF32 cfg ws valsM tokens nPairs
                  ioToST (MVector.unsafeWrite outM i l)
          putMVar doneVar ()
        pure doneVar
      forM_ doneMVars takeMVar
      Vector.unsafeFreeze outM

-- | Runs training steps (forward + backward) in FP32 in parallel across CPU capabilities.
trainingStepBatchF32 :: Config -> UVector.Vector Float -> Vector (UVector.Vector Int) -> IO (Vector (Float, UVector.Vector Float))
trainingStepBatchF32 cfg flatF32 batch
  | Vector.null batch = pure Vector.empty
  | otherwise = do
      numCaps <- getNumCapabilities
      let n = Vector.length batch
          nParams = UVector.length flatF32
      outM <- MVector.new n
      let chunks = splitIntoChunks numCaps (Vector.enumFromN (0 :: Int) n)
      doneMVars <- forM (zip [0 ..] chunks) $ \(cap, idxSlice) -> do
        doneVar <- newEmptyMVar
        _ <- forkOn cap $ do
          stToIO $ do
            ws <- newSTWorkspaceF32 cfg nParams
            valsM <- UVector.thaw flatF32
            Vector.forM_ idxSlice $ \i -> do
              let tokens = Vector.unsafeIndex batch i
                  nPairs = min cfg.blockSize (max 0 (UVector.length tokens - 1))
              if nPairs <= 0 || not (isValidModelShape cfg nParams)
                then ioToST (MVector.unsafeWrite outM i (0.0, UVector.empty))
                else do
                  l <- stForwardBackwardF32 cfg ws valsM tokens nPairs
                  gFrozen <- UVector.freeze ws.gradParamsBuf
                  ioToST (MVector.unsafeWrite outM i (l, gFrozen))
          putMVar doneVar ()
        pure doneVar
      forM_ doneMVars takeMVar
      Vector.unsafeFreeze outM

-- | 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
        }

-- | Forward, backward and Adam update on a document, returns the loss before the update
trainStep :: Config -> Hyperparams -> TrainState -> UVector.Vector Int -> (Double, TrainState)
trainStep cfg hp st tokens =
  let nPairs = min cfg.blockSize (max 0 (UVector.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
train :: Config -> Hyperparams -> TrainState -> Vector (UVector.Vector Int) -> (Vector Double, TrainState)
train cfg hp st0 docs
  | Vector.null docs = (Vector.empty, st0)
  | 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 (UVector.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

-- | 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 -> (UVector.Vector Int, StdGen)
stGenerate cfg bosId temperature flatW gen0
  | cfg.blockSize <= 0 || not (isValidModelShape cfg (UVector.length flatW)) = (UVector.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 <- UMVector.new bSize
      let stepPos posId tokenId !gen
            | posId >= bSize = do
                frozen <- UVector.unsafeFreeze outTokens
                pure (UVector.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 <- UVector.unsafeFreeze outTokens
                    pure (UVector.take posId frozen, gen')
                  else do
                    UMVector.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 -> (UVector.Vector Int, StdGen)
generate cfg bosId temperature pureWeights =
  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
  args <- getArgs
  let useFp32 = "--fp32" `elem` args
  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 (max 0 (n - length s)) ' ' <> s
  putStrLn $ "num params: " <> show (length w)
  when useFp32 $
    putStrLn "[mode] Single-precision FP32 with fastExpF32"

  let stepTokens
        | Vector.null docs = Vector.empty
        | otherwise =
            Vector.generate hp.numSteps $ \i ->
              encode tok (Vector.unsafeIndex docs (i `mod` Vector.length docs))
      st0 = initTrainState w

  isTty <- hIsTerminalDevice stdout
  let !totalStepsStr = padded 4 hp.numSteps
      reportStep i loss = do
        Char8.hPutStr
          stdout
          ( Char8.pack
              ( "step "
                  <> padded 4 (i + 1)
                  <> " / "
                  <> totalStepsStr
                  <> " | loss "
                  <> Numeric.showFFloat (Just 4) loss "\r"
              )
          )
        when isTty (hFlush stdout)

  final <-
    if useFp32
      then do
        let nParams = UVector.length st0.valuesVec
            validShape = isValidModelShape cfg nParams
            flatF32 = toFloatVec st0.valuesVec
        (valsM, mM, vM, ws) <- stToIO $ do
          valsBuf <- UVector.thaw flatF32
          mBuf <- UMVector.replicate nParams (0.0 :: Float)
          vBuf <- UMVector.replicate nParams (0.0 :: Float)
          wsBuf <- newSTWorkspaceF32 cfg nParams
          pure (valsBuf, mBuf, vBuf, wsBuf)
        Vector.iforM_ stepTokens $ \i tokens -> do
          let nPairs = min cfg.blockSize (max 0 (UVector.length tokens - 1))
              !lr = realToFrac (hp.learningRate * (1.0 - fromIntegral i / fromIntegral hp.numSteps)) :: Float
              !t = i + 1
          loss <- stToIO $ do
            if nPairs <= 0 || not validShape
              then do
                stZeroSliceF32 ws.gradParamsBuf 0 nParams
                stAdamUpdateF32 hp lr t nParams ws.gradParamsBuf valsM mM vM
                pure 0.0
              else do
                l <- stForwardBackwardF32 cfg ws valsM tokens nPairs
                stAdamUpdateF32 hp lr t nParams ws.gradParamsBuf valsM mM vM
                pure l
          reportStep i (realToFrac loss)
        vals' <- stToIO (UVector.unsafeFreeze valsM)
        pure
          st0
            { valuesVec = UVector.map realToFrac vals',
              step = Vector.length stepTokens
            }
      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_ stepTokens $ \i tokens -> do
          let nPairs = min cfg.blockSize (max 0 (UVector.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
          reportStep i loss
        (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 stepTokens
            }

  putStrLn "\n--- inference (new, hallucinated names) ---"
  let sampleOne = stGenerate cfg tok.bos 0.5 final.valuesVec
  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)