#⚡ flashmoon

#Why flashmoon

• 📐 flash + flash/gpu: standalone attention library — batched 4D MHA/GQA/MQA over [B,H,S,D] tensors, bottom-right causal masks, cross-attention with cached KV, arbitrary head dims; f32x4-SIMD wasm/native kernels plus a WebGPU backend. One moon add, no framework attached.
• 🎛️ gpu: WebGPU compute runtime — device/queue management, pipeline cache, single-pass batched submit, zero-copy bf16 weight upload, and a verified WGSL kernel library (flash prefill / flash-decoding / bf16 GEMV·GEMM / fused SiLU·residual·RoPE / GPU argmax).
• 🧮 Numerics you can trust — every kernel ships with a CPU/f64 oracle check; the attention library carries randomized property tests against its naive reference. No silent drift.
• 📦 Zero-dependency core — the entire GPU path is MoonBit code + WebGPU. No native libs, no Python, no ONNX.

#Use it as a library

moon add starAndHonor/flashmoon

#1. Batched attention — flash

q   [B, H,   Sq, D ]      out [B, H, Sq, DV]
k   [B, HKV, Skv, D ]
v   [B, HKV, Skv, DV]     out[b,h,i,:] = softmax(q[b,h,i,:] · k[b,kvh,:,:]ᵀ · scale) · v[b,kvh,:,:]

import { "starAndHonor/flashmoon/flash" }

let cfg = @flash.AttnConfig::new(
  heads=8,          // query heads H
  kv_heads=2,       // key/value heads HKV: H -> MHA, H/n -> GQA, 1 -> MQA
  causal=true,      // bottom-right-aligned causal mask
  scale=None,       // None = 1/sqrt(D); Some(x) to override
)
let out = @flash.flash_attention(q, k, v, cfg, block_rows=64, block_cols=64)

• Causal is bottom-right aligned: row i attends keys 0 ..= i + Skv - Sq. So Sq == Skv is ordinary self-attention, Sq < Skv is attention over a cached KV prefix (incremental decode), and Sq > Skv with causal=true is rejected (abort).
• GQA/MQA repeat each KV head across H / HKV query heads; kv_heads here is the shape of k/v, and must match cfg.
• Head dims are arbitrary (D and DV may differ, no alignment required) — the SIMD kernels have a scalar tail. block_rows/block_cols are tuning knobs only; results are identical up to float summation order (measured ≤ 6e-6 at 4096², see Performance).
• Validation is fail-fast: dims, batch, head counts and Skv >= Sq are checked on entry and reported via abort with the offending numbers.
• Scale defaults to 1/sqrt(D); the config is shared by the WebGPU backend, so the same cfg gives the same math on both.

#2. WebGPU attention — flash/gpu

import {
  "starAndHonor/flashmoon/flash",
  "starAndHonor/flashmoon/flash/gpu" @flashgpu,
  "starAndHonor/flashmoon/gpu",
}

@gpu.Gpu::init(fn(g) {
  @flashgpu.flash_attention(g, q, k, v, cfg, fn(out) {
    // out : NpArray [B, H, Sq, DV]
  })
}, log=println)

#3. WebGPU runtime — gpu

let qb = g.upload(q_data)              // FixedArray[Float] -> GPUBuffer
let outb = g.alloc(b * h * sq * dv)    // f32 storage buffer
let scale = 1.0 / d.to_double().sqrt()
g.attn_4d(qb, kb, vb, outb, b, h, hkv, sq, skv, d, dv, /* causal */ true, scale)
g.sync(fn(_) { ... })                  // wait for the queue
g.readback(outb, n, fn(data) { ... })  // GPUBuffer -> FixedArray[Float]

#Runnable examples

#Performance

#GPU: naive vs flash (WebGPU, RTX 4060 Laptop, 64×64 tiles)

#CPU: naive vs flash — moon run bench --target wasm --release

#Kernels (RTX 4060 Laptop, release, dispatch-only unless noted)

#Quick start

moon build --target js
deno run --allow-read scripts/webgpu_host.js      # kernel checks + throughput
deno run --allow-read scripts/bench_gpu_host.js   # naive vs flash scenario bench

moon test                                  # full suite, incl. property tests
moon run examples/attn_cpu --target wasm   # minimal 4D GQA + causal example

#Architecture

flowchart LR
    subgraph host["Hosts (demo)"]
        B["🌐 Browser page<br/>cmd/webchat"]
        D["🖥️ Deno REPL<br/>cmd/qwengpu"]
        K["🔬 Bench host<br/>cmd/gpubench"]
    end
    subgraph lib["flashmoon library (pure MoonBit)"]
        G["gpu<br/>WebGPU runtime + WGSL kernels<br/>pipelines · buffers · submit"]
        F["flash + flash/gpu<br/>4D batched attention (MHA/GQA/MQA)<br/>wasm SIMD · native · WebGPU"]
    end
    subgraph demo["demo/ — model components"]
        Q["demo/qwen<br/>safetensors + BPE tokenizer"]
        R["demo/qwenrun<br/>host-agnostic runner<br/>weights · KV cache · decode loop"]
    end
    B & D --> R
    K --> G & F
    R --> G & Q
    F --> G

On the name. The flash library and the WebGPU inference kernels implement the FlashAttention-family algorithm — online softmax with running max/sum and deferred 1/l normalization (prefill: tiled K/V staging; decode: flash-decoding split-KV + combine). FA2's headline contributions are CUDA grid/warp scheduling (Q-parallel thread blocks, per-warp Q partitioning), which don't translate to WGSL — so we claim the family, not the version.

#Verification

#Repository layout

flash/                         4D batched flash attention library (wasm/native, f32x4 SIMD)
flash/gpu/                     WebGPU backend for the flash library (js target)
gpu/                           WebGPU runtime + compute kernels (js target)
demo/                          downstream model components + LLM demo (see demo/README.md)
demo/qwen/                     safetensors parser + Qwen2 byte-level BPE tokenizer
demo/qwenrun/                  Qwen3-0.6B runner core (host-agnostic: read/log injected)
examples/                      standalone library usage examples (CPU + WebGPU)
bench/                         naive-vs-flash scenario suite + sweeps (wasm/native)
bench/gpu/                     naive-vs-flash GPU bench (js/WebGPU, Deno host)
cmd/fa/                        naive-vs-flash attention demo (wasm/native)
cmd/gpubench/                  WebGPU kernel checks + benchmarks (Deno host)
cmd/qwencpu/                   Qwen3 CPU reference runner + tokenizer oracle test
cmd/qwengpu/                   Deno REPL chat (MATCH gate + slash commands)
cmd/webchat/                   browser chat page (chat.html + DOM frontend)
test/                          blackbox tests (flash attention, benchmarks, tokenizer oracle)
refs/                          model + HF reference data (gitignored, ~1.5 GB)
scripts/                       Deno host shims (webgpu_host.js, bench_gpu_host.js, attn_gpu_host.js, qwengpu_host.js)

#References & licenses

#License