///|
/// Demo: batched GQA causal attention — naive reference vs flash kernel
/// on random 4D inputs.
fn main {
  let b = 2 // batch
  let h = 4 // query heads
  let hkv = 2 // key/value heads (GQA)
  let sq = 64 // query length
  let skv = 64 // key/value length
  let d = 32 // head dim
  let rand = fn(i : Int, j : Int, k : Int) -> Float {
    (Float::from_int((i * 37 + j * 91 + k * 13 + 11) % 29) - 14.0F) / 14.0F
  }
  let mk = fn(bh : Int, rows : Int, cols : Int, k : Int) -> @nparray.NpArray {
    let data = FixedArray::make(bh * rows * cols, 0.0F)
    for i in 0..<(bh * rows) {
      for j in 0.. max_diff {
      max_diff = diff
    }
  }
  println(
    "attention: b=\{b} h=\{h} hkv=\{hkv} sq=\{sq} skv=\{skv} d=\{d} causal",
  )
  println("naive out[0] = \{da[0]}")
  println("flash out[0] = \{db[0]}")
  println("max |naive - flash| = \{max_diff}")
}