Upload 14 files

check_vocab.py

@@ -0,0 +1,4 @@
+import json
+
+v = json.load(open('vocab.json'))
+print(f'Current vocab size: {len(v["char_to_idx"])}')

config.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+EPOCHS = 20
+CONTEXT = 1024
+data_path = "./training_data"
+batch_size = 8
+learning_rate = 1e-4
+weight_decay = 0.01
+
+vocab_size = 8000
+D_MODEL = 256
+numberoflayers = 8
+numberofheads = 4
+num_kv_heads = 2
+d_Latent = 64
+ffn_mult = 3.5
+swa_window = 512
+use_moe = False
+moe_num_experts = 8
+moe_top_k = 2
+
+GLOBAL_DTYPE = "bfloat16"   # "float32", "float16", "bfloat16"
+MATMUL_PRECISION = "high"   # "default", "high", "highest"
+USE_FLASH_ATTENTION = False
+USE_SPLASH_ATTENTION = True

corpus.txt

@@ -0,0 +1 @@
+Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. Advanced reasoning involves breaking down problems into steps, critiquing solutions, and iteratively improving answers. Chain-of-thought improves logical consistency. Self-critique helps models identify flaws in their own reasoning. Recursive refinement (TRM) leads to better final outputs. 

inference.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import os
+os.environ["KERAS_BACKEND"] = "jax"
+
+import numpy as np
+import jax
+import keras
+
+from veylon_model import create_llm
+from tokenizer import TokenizerWrapper
+
+from config import (
+    CONTEXT,
+    vocab_size,
+    D_MODEL,
+    numberoflayers,
+    numberofheads,
+    d_Latent,
+    ffn_mult,
+    num_kv_heads,
+    swa_window,
+)
+
+# ============================================================
+# Runtime info
+# ============================================================
+
+print(f"Backend: {keras.backend.backend()}")
+print(f"JAX devices: {jax.devices()}")
+
+keras.mixed_precision.set_global_policy("mixed_bfloat16")
+
+# ============================================================
+# Load tokenizer
+# ============================================================
+
+tokenizer = TokenizerWrapper("tokenizer.model")
+
+assert tokenizer.vocab_size == vocab_size, (
+    f"Tokenizer vocab ({tokenizer.vocab_size}) "
+    f"!= config vocab ({vocab_size})"
+)
+
+print(f"Tokenizer vocab size: {tokenizer.vocab_size}")
+
+# ============================================================
+# Build model (must exactly match training)
+# ============================================================
+
+print("Building model...")
+
+model = create_llm(
+    vocab_size=vocab_size,
+    d_model=D_MODEL,
+    n_layers=numberoflayers,
+    n_heads=numberofheads,
+    d_latent=d_Latent,
+    ffn_mult=ffn_mult,
+    max_seq_len=CONTEXT,
+    use_moe=False,
+    num_kv_heads=num_kv_heads,
+    swa_window=swa_window,
+)
+
+# Warmup with EXACT training/inference shape
+dummy = np.zeros((1, CONTEXT), dtype=np.int32)
+_ = model(dummy, training=False)
+
+print("✓ Model built successfully")
+
+# ============================================================
+# Load weights
+# ============================================================
+
+WEIGHTS_PATH = "veylon_final.weights.h5"
+print(f"Loading weights from: {WEIGHTS_PATH}")
+model.load_weights(WEIGHTS_PATH)
+print("✓ Weights loaded successfully")
+
+# ============================================================
+# Sampling settings
+# ============================================================
+
+MAX_NEW_TOKENS = 64
+TEMPERATURE = 0.8
+TOP_K = 50
+
+def sample_from_logits(
+    logits: np.ndarray,
+    temperature: float = 0.8,
+    top_k: int = 50,
+) -> int:
+    """
+    NumPy-only sampling to avoid JAX/readonly array issues.
+    """
+    logits = np.array(logits, dtype=np.float32, copy=True)
+
+    if temperature > 0:
+        logits = logits / float(max(temperature, 1e-8))
+
+    if top_k > 0:
+        k = min(int(top_k), logits.shape[-1])
+        row = logits[0]
+        top_indices = np.argpartition(row, -k)[-k:]
+
+        filtered = np.full_like(row, -np.inf)
+        filtered[top_indices] = row[top_indices]
+        logits[0] = filtered
+
+    row = logits[0]
+    row = row - np.max(row)
+    probs = np.exp(row)
+    probs = probs / probs.sum()
+
+    return int(np.random.choice(len(probs), p=probs))
+
+# ============================================================
+# Generation loop
+# ============================================================
+
+while True:
+    prompt = input("\nEnter your prompt (or 'exit'): ").strip()
+
+    if prompt.lower() in {"exit", "quit"}:
+        break
+
+    tokens = tokenizer.encode(
+        prompt,
+        add_bos=True,
+        add_eos=False,
+    )
+
+    if len(tokens) == 0:
+        tokens = [tokenizer.bos_id if hasattr(tokenizer, "bos_id") else 1]
+
+    tokens = tokens[-CONTEXT:]
+
+    print("\nGenerating...\n")
+
+    # Prompt prefill (one-time)
+    prompt_ids = np.array([tokens], dtype=np.int32)
+    logits, cache_k, cache_v = model.generate_step(
+        prompt_ids,
+        cache_k=None,
+        cache_v=None,
+        cache_pos=0,
+    )
+
+    next_token = sample_from_logits(
+        np.array(logits[:, -1, :], dtype=np.float32, copy=True),
+        temperature=TEMPERATURE,
+        top_k=TOP_K,
+    )
+    tokens.append(next_token)
+
+    if next_token != tokenizer.eos_id and len(tokens) < CONTEXT:
+        # After prefill, we are decoding token-by-token.
+        cache_pos = len(prompt_ids[0])
+
+        for _ in range(MAX_NEW_TOKENS - 1):
+            next_input = np.array([[next_token]], dtype=np.int32)
+
+            logits, cache_k, cache_v = model.generate_step(
+                next_input,
+                cache_k=cache_k,
+                cache_v=cache_v,
+                cache_pos=cache_pos,
+            )
+
+            cache_pos += 1
+
+            next_token = sample_from_logits(
+                np.array(logits[:, -1, :], dtype=np.float32, copy=True),
+                temperature=TEMPERATURE,
+                top_k=TOP_K,
+            )
+            tokens.append(next_token)
+
+            if next_token == tokenizer.eos_id:
+                break
+
+            if len(tokens) >= CONTEXT:
+                print("\n[Context limit reached]")
+                break
+
+    generated_text = tokenizer.decode(tokens)
+
+    print("\n" + "=" * 60)
+    print("Veylon Alpha")
+    print("=" * 60)
+    print(generated_text)
+    print("=" * 60)

inspect_weights.py

@@ -0,0 +1,7 @@
+import h5py
+
+with h5py.File('generative_model.weights.h5', 'r') as f:
+    def print_structure(name, obj):
+        if isinstance(obj, h5py.Dataset):
+            print(f"{name}: {obj.shape} dtype={obj.dtype}")
+    f.visititems(print_structure)

model_build.py

@@ -0,0 +1,56 @@
+# model_build.py
+
+import os
+os.environ["KERAS_BACKEND"] = "jax"
+
+import keras
+
+from veylon_model import create_llm
+from tokenizer import TokenizerWrapper
+
+from config import (
+    CONTEXT,
+    ffn_mult,
+    d_Latent,
+    D_MODEL,
+    numberofheads,
+    numberoflayers,
+    vocab_size
+)
+
+# Mixed precision is set inside veylon_model.py, but we ensure it here too
+keras.mixed_precision.set_global_policy("mixed_bfloat16")
+
+tokenizer = TokenizerWrapper("tokenizer.json")
+
+# n_kv_heads is completely removed - pure MLA
+model = create_llm(
+    vocab_size=vocab_size,
+    d_model=D_MODEL,
+    n_layers=numberoflayers,
+    n_heads=numberofheads,
+    d_latent=d_Latent,
+    ffn_mult=ffn_mult,
+    max_seq_len=CONTEXT,
+    use_moe=False,
+)
+
+optimizer = keras.optimizers.AdamW(
+    learning_rate=1e-4,
+    weight_decay=0.01,
+    global_clipnorm=1.0,
+)
+
+loss_fn = keras.losses.SparseCategoricalCrossentropy(
+    from_logits=True,
+)
+
+model.compile(
+    optimizer=optimizer,
+    loss=loss_fn,
+    jit_compile=True,
+)
+
+# Explicitly build to print summary
+model.build((None, CONTEXT))
+model.summary()

token_train.py

@@ -0,0 +1,20 @@
+import os
+from tokenizer import train_sentencepiece
+from config import vocab_size
+DATA_DIR = "./training_data"
+
+text_files = [
+    os.path.join(DATA_DIR, f)
+    for f in os.listdir(DATA_DIR)
+    if f.endswith(".txt")
+]
+
+print(f"Found {len(text_files)} text files")
+
+if not text_files:
+    raise ValueError(
+        f"No .txt files found in {DATA_DIR}"
+    )
+
+tokenizer = train_sentencepiece(text_files,vocab_size=vocab_size)
+

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa5baedc3bd88d9f1321fcb36c8d9d73c30727bfafb6276951697df6a20d9028
+size 369778

train.py

@@ -0,0 +1,439 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+import time
+
+os.environ['KERAS_BACKEND']                  = 'jax'
+os.environ['XLA_PYTHON_CLIENT_PREALLOCATE']  = 'true'
+os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.90'
+os.environ['TF_CPP_MIN_LOG_LEVEL']           = '3'
+os.environ['PYTHONWARNINGS']                 = 'ignore'
+os.environ['TF_DATA_EXPERIMENTAL_SLACK']     = '1'
+
+import numpy as np
+import tensorflow as tf
+import jax
+import keras
+
+from tokenizer import TokenizerWrapper
+from veylon_model import create_llm
+
+try:
+    from config import (EPOCHS, CONTEXT, vocab_size, D_MODEL, numberoflayers,
+                        numberofheads, d_Latent, ffn_mult, swa_window,
+                        num_kv_heads, learning_rate, weight_decay, batch_size,
+                        data_path)
+except Exception:
+    EPOCHS         = 20
+    CONTEXT        = 2048
+    vocab_size     = 32000
+    D_MODEL        = 512
+    numberoflayers = 8
+    numberofheads  = 8
+    d_Latent       = 128
+    ffn_mult       = 3.5
+    swa_window     = 1024
+    num_kv_heads   = 2
+    learning_rate  = 1e-4
+    weight_decay   = 0.01
+    batch_size     = 8
+    data_path      = './training_data'
+
+keras.mixed_precision.set_global_policy('mixed_bfloat16')
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# PIPELINE
+# ─────────────────────────────────────────────────────────────────────────────
+
+def precompute_windows(tokens: np.ndarray, seq_len: int, stride: int) -> np.ndarray:
+    """
+    Slice the token array into overlapping (seq_len+1) windows using
+    sliding_window_view, then stride-subsample.
+
+    sliding_window_view is preferred over as_strided: it validates bounds
+    internally and raises a clear error instead of silently reading garbage
+    memory past the array end.
+
+    The [::stride] slice is a zero-copy view.  ascontiguousarray forces one
+    real allocation — a dense (N, seq_len+1) int32 buffer — which TensorFlow
+    can wrap without copying again.
+
+    Memory: for 5.7M tokens, CONTEXT=2048, stride=1024 → ~45 MB. Trivial
+    against the 42 GB host RAM on Colab/Kaggle TPU instances.
+    """
+    window_len = seq_len + 1
+    view = np.lib.stride_tricks.sliding_window_view(tokens, window_len)[::stride]
+
+    if len(view) == 0:
+        raise ValueError(
+            f"Token array ({len(tokens):,}) is too short for "
+            f"seq_len={seq_len}, stride={stride}."
+        )
+
+    return np.ascontiguousarray(view, dtype=np.int32)
+
+
+def _make_tf_data_options() -> tf.data.Options:
+    """
+    Bundle all tf.data graph-level knobs in one place.
+
+    map_parallelization: lets the runtime fuse and parallelize map ops across
+      the thread pool automatically — no manual num_parallel_calls needed for
+      simple element-wise transforms.
+
+    parallel_batch: batching itself becomes a parallel operation; each batch
+      slot is filled concurrently instead of serially.
+
+    experimental_slack: introduces a one-step slack between the prefetch
+      stage and the training loop so the pipeline stays one batch ahead
+      without blocking on the next step.  Redundant with the env var above
+      but belt-and-suspenders costs nothing.
+    """
+    opts = tf.data.Options()
+    opts.experimental_optimization.map_parallelization = True
+    opts.experimental_optimization.parallel_batch      = True
+    opts.experimental_slack                            = True
+    return opts
+
+
+def create_tf_dataset(
+    windows: np.ndarray,
+    batch_size: int,
+    shuffle: bool = True,
+    shuffle_seed: int = 42,
+    cache_path: str = "",          # "" → in-memory cache; "/tmp/..." → disk
+) -> tf.data.Dataset:
+    """
+    High-throughput tf.data pipeline.
+
+    Ordering rationale (each decision is load-bearing):
+
+    1. from_tensor_slices(windows)
+       TensorFlow wraps the numpy buffer directly — no tf.constant() copy.
+       Cardinality is exactly N (known), enabling optimal prefetch depth and
+       shuffle buffer auto-sizing downstream.
+
+    2. cache()  ← BEFORE shuffle and map
+       Caches individual (seq_len+1,) sequence tensors.  On epoch 2+ the
+       entire dataset is served from RAM, eliminating all numpy I/O.
+       Caching before shuffle means the cache stores the canonical data once;
+       shuffle order changes every epoch without invalidating the cache.
+       Caching after batching would freeze batch composition forever, which
+       breaks per-epoch shuffle semantics.
+
+    3. shuffle()  ← AFTER cache, BEFORE repeat
+       Operates on individual sequences (correct granularity).
+       buffer=min(N, 20_000) bounds peak memory while covering the full
+       dataset for corpora this size.
+
+    4. repeat()  ← AFTER shuffle, BEFORE map/batch
+       Placed here so the shuffle buffer can draw across epoch boundaries,
+       preventing an artificial ordering reset at each epoch edge.
+
+    5. map(split_xy)
+       Slices each (seq_len+1,) window into x=(seq_len,) and y=(seq_len,).
+       No @tf.function decorator needed — tf.data traces the lambda itself.
+       num_parallel_calls=AUTOTUNE parallelises across the CPU thread pool.
+
+    6. batch(drop_remainder=True)
+       Static shape [batch_size, seq_len] — mandatory for XLA/TPU.
+       No symbolic/dynamic shapes ever reach the device.
+
+    7. prefetch(AUTOTUNE)
+       Keeps the device-transfer queue full.  Always last so it buffers
+       fully-formed batches, not individual sequences.
+
+    8. with_options(...)
+       Applies graph-level optimizations in one shot after the pipeline is
+       fully defined.
+    """
+    n = len(windows)
+
+    ds = tf.data.Dataset.from_tensor_slices(windows)           # (1) wrap
+
+    ds = ds.cache(cache_path)                                  # (2) cache
+
+    if shuffle:
+        buf = min(n, 20_000)
+        ds  = ds.shuffle(buffer_size=buf, seed=shuffle_seed,
+                         reshuffle_each_iteration=True)        # (3) shuffle
+
+    ds = ds.repeat()                                           # (4) repeat
+
+    ds = ds.map(                                               # (5) split
+        lambda w: (w[:-1], w[1:]),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+
+    ds = ds.batch(batch_size, drop_remainder=True)             # (6) batch
+
+    ds = ds.prefetch(tf.data.AUTOTUNE)                        # (7) prefetch
+
+    ds = ds.with_options(_make_tf_data_options())              # (8) opts
+
+    return ds
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# BENCHMARK
+# ─────────────────────────────────────────────────────────────────────────────
+
+def benchmark_dataset(
+    ds: tf.data.Dataset,
+    batch_size: int,
+    seq_len: int,
+    n_batches: int = 200,
+    label: str = "dataset",
+) -> dict:
+    """
+    Measure pipeline throughput independently of TPU compute.
+
+    We do NOT call .numpy() inside the loop.  That forces a device→host
+    transfer and a Python GIL acquisition on every batch, serialising what
+    should be a fully-async prefetch chain.  Instead we consume the iterator
+    and let TensorFlow materialise tensors lazily.  The end-to-end wall time
+    is what the TPU will actually experience.
+
+    Interpret results:
+      pipeline_tok/s > 2× tpu_tok/s  →  pipeline is not the bottleneck
+      pipeline_tok/s < 1.5× tpu_tok/s →  still starving the TPU
+    """
+    print(f"\n{'─'*52}")
+    print(f"Benchmarking: {label}")
+    print(f"{'─'*52}")
+
+    # Warm up: populate cache + prefetch buffers before timing starts.
+    for _ in ds.take(5):
+        pass
+    print("Warm-up (5 batches) complete.")
+
+    start = time.perf_counter()
+    for _ in ds.take(n_batches):
+        pass
+    elapsed = time.perf_counter() - start
+
+    tokens_per_batch = batch_size * seq_len
+    total_tokens     = n_batches * tokens_per_batch
+    batches_sec      = n_batches / elapsed
+    tokens_sec       = total_tokens / elapsed
+
+    print(f"Batches measured : {n_batches}")
+    print(f"Elapsed          : {elapsed:.2f} s")
+    print(f"Throughput       : {batches_sec:.1f} batches/s")
+    print(f"Throughput       : {tokens_sec:,.0f} tokens/s")
+
+    tpu_target = 30_000
+    headroom   = tokens_sec / tpu_target
+    if headroom >= 2.0:
+        verdict = f"✓  {headroom:.1f}× headroom over TPU — pipeline is NOT the bottleneck."
+    elif headroom >= 1.2:
+        verdict = f"⚠  {headroom:.1f}× headroom — marginal; increase batch_size or reduce stride."
+    else:
+        verdict = f"✗  {tokens_sec:,.0f} tok/s < TPU target {tpu_target:,} — still starving the TPU."
+
+    print(verdict)
+    print(f"{'─'*52}\n")
+
+    return {"batches_sec": batches_sec, "tokens_sec": tokens_sec,
+            "headroom_vs_30k": headroom}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# THROUGHPUT CALLBACK
+# ─────────────────────────────────────────────────────────────────────────────
+
+class ThroughputCallback(keras.callbacks.Callback):
+    """
+    Logs throughput every LOG_INTERVAL steps and at epoch end.
+
+    Why coarse logging matters on TPU:
+    Keras callbacks run on the Python host.  JAX dispatches training steps
+    asynchronously — Python returns before XLA finishes the kernel.  A
+    time.perf_counter() call on every step forces a host sync, equivalent
+    to inserting jax.block_until_ready() every step and destroying async
+    pipelining.  At LOG_INTERVAL=50 only ~2% of steps incur this cost.
+    """
+    LOG_INTERVAL = 50
+
+    def on_epoch_begin(self, epoch, logs=None):
+        self._epoch_start    = time.perf_counter()
+        self._interval_start = self._epoch_start
+        self._step_count     = 0
+
+    def on_train_batch_end(self, batch, logs=None):
+        self._step_count += 1
+        if batch > 0 and batch % self.LOG_INTERVAL == 0:
+            now     = time.perf_counter()
+            elapsed = now - self._interval_start
+            tps     = (self.LOG_INTERVAL * batch_size * CONTEXT) / elapsed
+            loss    = logs.get('loss', float('nan'))
+            print(f"  step {batch:5d} | loss {loss:.4f} | {tps:>10,.0f} tok/s")
+            self._interval_start = now
+
+    def on_epoch_end(self, epoch, logs=None):
+        elapsed      = time.perf_counter() - self._epoch_start
+        total_tokens = self._step_count * batch_size * CONTEXT
+        val_loss     = logs.get('val_loss', float('nan'))
+        print(f"\nEpoch {epoch+1} | {elapsed:.1f}s | "
+              f"{total_tokens / elapsed:,.0f} tok/s (avg) | "
+              f"val_loss={val_loss:.4f}\n")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# HELPERS
+# ─────────────────────────────────────────────────────────────────────────────
+
+def load_text_corpus(path: str) -> str:
+    p     = Path(path)
+    files = [p] if p.is_file() else sorted(p.glob('*.txt'))
+    if not files:
+        raise FileNotFoundError(f'No .txt files found in {path}')
+    return ''.join(fp.read_text(encoding='utf-8') for fp in files)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# MAIN
+# ─────────────────────────────────────────────────────────────────────────────
+class MemoryCallback(keras.callbacks.Callback):
+    def on_epoch_end(self, epoch, logs=None):
+        stats = jax.devices()[0].memory_stats()
+
+        used = stats["bytes_in_use"] / 1024**3
+        peak = stats["peak_bytes_in_use"] / 1024**3
+        limit = stats["bytes_limit"] / 1024**3
+
+        print(
+            f"\nHBM: {used:.2f}/{limit:.2f} GB "
+            f"(Peak: {peak:.2f} GB)"
+        )
+def main():
+    print('\n' + '=' * 60)
+    print('BACKEND & DEVICE VERIFICATION')
+    print('=' * 60)
+    print(f'Keras backend  : {keras.backend.backend()}')
+    print(f'JAX devices    : {jax.devices()}')
+    print('=' * 60 + '\n')
+
+    # ── Tokenize ──────────────────────────────────────────────────────────
+    tokenizer  = TokenizerWrapper('tokenizer.model')
+    text       = load_text_corpus(data_path)
+    raw_tokens = np.asarray(
+        tokenizer.encode(text, add_bos=True, add_eos=False), dtype=np.int32
+    )
+    print(f'Total tokens: {len(raw_tokens):,}')
+
+    split_idx    = int(len(raw_tokens) * 0.98)
+    train_tokens = raw_tokens[:split_idx]
+    val_tokens   = raw_tokens[split_idx:]
+
+    stride = max(1, CONTEXT // 2)
+
+    # ── Pre-compute windows ────────────────────────────────────────────────
+    print("Pre-computing windows...", end=" ", flush=True)
+    t0            = time.perf_counter()
+    train_windows = precompute_windows(train_tokens, CONTEXT, stride)
+    val_windows   = precompute_windows(val_tokens,   CONTEXT, stride)
+    print(f"done in {(time.perf_counter() - t0) * 1000:.1f} ms")
+    print(f"  train: {train_windows.shape}  ({train_windows.nbytes / 1e6:.1f} MB)")
+    print(f"  val:   {val_windows.shape}  ({val_windows.nbytes / 1e6:.1f} MB)\n")
+
+    # ── Build datasets ─────────────────────────────────────────────────────
+    train_ds = create_tf_dataset(train_windows, batch_size, shuffle=True)
+    val_ds   = create_tf_dataset(val_windows,   batch_size, shuffle=False)
+
+    # ── Benchmark before training ──────────────────────────────────────────
+    # Verifies the pipeline can outpace the TPU before we waste a run.
+    benchmark_dataset(train_ds, batch_size, CONTEXT, n_batches=200,
+                      label="train_ds")
+
+    # ── Build model ────────────────────────────────────────────────────────
+    os.makedirs('checkpoints', exist_ok=True)
+
+    model = create_llm(
+        vocab_size   = tokenizer.vocab_size,
+        d_model      = D_MODEL,
+        n_layers     = numberoflayers,
+        n_heads      = numberofheads,
+        d_latent     = d_Latent,
+        ffn_mult     = ffn_mult,
+        max_seq_len  = CONTEXT,
+        use_moe      = False,
+        num_kv_heads = num_kv_heads,
+        swa_window   = swa_window,
+    )
+
+    optimizer = keras.optimizers.AdamW(
+        learning_rate   = learning_rate,
+        weight_decay    = weight_decay,
+        global_clipnorm = 1.0,
+    )
+    model.compile(
+        optimizer  = optimizer,
+        loss       = keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        jit_compile = True,
+    )
+    
+    # ── Warmup with exact training shape ──────────────────────────────────
+    # XLA traces a kernel per unique input shape.  Warming up with any other
+    # shape (e.g. the original (1, min(4, CONTEXT))) causes a full recompile
+    # on the first real training batch — wasting 30-120 s on the TPU.
+    print("Warming up XLA with training shape...", end=" ", flush=True)
+    t0      = time.perf_counter()
+    dummy_x = np.zeros((batch_size, CONTEXT), dtype=np.int32)
+    _       = model(dummy_x, training=False)
+    _       = model(dummy_x, training=True)
+    print(f"done in {time.perf_counter() - t0:.2f}s\n")
+
+    # ── Forward pass sanity check ──────────────────────────────────────────
+    # One explicit transfer (np.array) then work in NumPy — avoids the 4
+    # implicit device→host round trips of the original float(ops.min(...))
+    # scalar-cast pattern.
+    print('─' * 52)
+    print('Forward pass validation')
+    print('─' * 52)
+    for x_batch, _ in train_ds.take(1):
+        logits    = model(x_batch[:1], training=False)
+        logits_np = np.array(logits)
+        print(f'Input  : {x_batch.shape}')
+        print(f'Logits : {logits.shape}')
+        print(f'min={logits_np.min():.4f}  max={logits_np.max():.4f}')
+        if np.isnan(logits_np).any() or np.isinf(logits_np).any():
+            raise ValueError('Forward pass produced NaN/Inf — check model init.')
+    print('✓ Forward pass clean.\n')
+    
+    # ── Steps per epoch from actual window count ───────────────────────────
+    # Original used len(tokens) // (batch*seq) which undercounts overlapping
+    # windows and causes epochs to terminate prematurely.
+    print(f"Total params: {model.count_params():,}")
+    steps_per_epoch  = max(1, len(train_windows) // batch_size)
+    validation_steps = max(1, len(val_windows)   // batch_size)
+    print(f"steps_per_epoch  : {steps_per_epoch}")
+    print(f"validation_steps : {validation_steps}\n")
+
+    # ── Train ──────────────────────────────────────────────────────────────
+    model.fit(
+        train_ds,
+        validation_data  = val_ds,
+        epochs           = EPOCHS,
+        steps_per_epoch  = steps_per_epoch,
+        validation_steps = validation_steps,
+        callbacks        = [
+            keras.callbacks.TerminateOnNaN(),
+            keras.callbacks.ModelCheckpoint(
+                filepath          = 'checkpoints/veylon_{epoch:02d}.weights.h5',
+                save_freq         = 'epoch',
+                save_weights_only = True,
+            ),
+            ThroughputCallback(),
+            MemoryCallback(),
+        ],
+    )
+
+    model.save_weights('veylon_final.weights.h5')
+    print('Saved veylon_final.weights.h5')
+
+
+if __name__ == '__main__':
+    main()

train2.py

@@ -0,0 +1,547 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+import time
+
+os.environ['KERAS_BACKEND']                  = 'jax'
+os.environ['XLA_PYTHON_CLIENT_PREALLOCATE']  = 'true'
+os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.90'
+os.environ['TF_CPP_MIN_LOG_LEVEL']           = '3'
+os.environ['PYTHONWARNINGS']                 = 'ignore'
+os.environ['TF_DATA_EXPERIMENTAL_SLACK']     = '1'
+
+import numpy as np
+import tensorflow as tf
+import jax
+import jax.sharding as sharding
+from jax.sharding import PartitionSpec as P
+import keras
+
+from tokenizer import TokenizerWrapper
+from veylon_model import create_llm
+
+try:
+    from config import (EPOCHS, CONTEXT, vocab_size, D_MODEL, numberoflayers,
+                        numberofheads, d_Latent, ffn_mult, swa_window,
+                        num_kv_heads, learning_rate, weight_decay, batch_size,
+                        data_path)
+except Exception:
+    EPOCHS         = 20
+    CONTEXT        = 2048
+    vocab_size     = 32000
+    D_MODEL        = 512
+    numberoflayers = 8
+    numberofheads  = 8
+    d_Latent       = 128
+    ffn_mult       = 3.5
+    swa_window     = 1024
+    num_kv_heads   = 2
+    learning_rate  = 1e-4
+    weight_decay   = 0.01
+    batch_size     = 8
+    data_path      = './training_data'
+
+keras.mixed_precision.set_global_policy('mixed_bfloat16')
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# PIPELINE
+# ─────────────────────────────────────────────────────────────────────────────
+
+def precompute_windows(tokens: np.ndarray, seq_len: int, stride: int) -> np.ndarray:
+    """
+    Slice the token array into overlapping (seq_len+1) windows using
+    sliding_window_view, then stride-subsample.
+
+    sliding_window_view is preferred over as_strided: it validates bounds
+    internally and raises a clear error instead of silently reading garbage
+    memory past the array end.
+
+    The [::stride] slice is a zero-copy view.  ascontiguousarray forces one
+    real allocation — a dense (N, seq_len+1) int32 buffer — which TensorFlow
+    can wrap without copying again.
+
+    Memory: for 5.7M tokens, CONTEXT=2048, stride=1024 → ~45 MB. Trivial
+    against the 42 GB host RAM on Colab/Kaggle TPU instances.
+    """
+    window_len = seq_len + 1
+    view = np.lib.stride_tricks.sliding_window_view(tokens, window_len)[::stride]
+
+    if len(view) == 0:
+        raise ValueError(
+            f"Token array ({len(tokens):,}) is too short for "
+            f"seq_len={seq_len}, stride={stride}."
+        )
+
+    return np.ascontiguousarray(view, dtype=np.int32)
+
+
+def _make_tf_data_options() -> tf.data.Options:
+    """
+    Bundle all tf.data graph-level knobs in one place.
+
+    map_parallelization: lets the runtime fuse and parallelize map ops across
+      the thread pool automatically — no manual num_parallel_calls needed for
+      simple element-wise transforms.
+
+    parallel_batch: batching itself becomes a parallel operation; each batch
+      slot is filled concurrently instead of serially.
+
+    experimental_slack: introduces a one-step slack between the prefetch
+      stage and the training loop so the pipeline stays one batch ahead
+      without blocking on the next step.  Redundant with the env var above
+      but belt-and-suspenders costs nothing.
+    """
+    opts = tf.data.Options()
+    opts.experimental_optimization.map_parallelization = True
+    opts.experimental_optimization.parallel_batch      = True
+    opts.experimental_slack                            = True
+    return opts
+
+
+def create_tf_dataset(
+    windows: np.ndarray,
+    batch_size: int,
+    shuffle: bool = True,
+    shuffle_seed: int = 42,
+    cache_path: str = "",          # "" → in-memory cache; "/tmp/..." → disk
+) -> tf.data.Dataset:
+    """
+    High-throughput tf.data pipeline.
+
+    Ordering rationale (each decision is load-bearing):
+
+    1. from_tensor_slices(windows)
+       TensorFlow wraps the numpy buffer directly — no tf.constant() copy.
+       Cardinality is exactly N (known), enabling optimal prefetch depth and
+       shuffle buffer auto-sizing downstream.
+
+    2. cache()  ← BEFORE shuffle and map
+       Caches individual (seq_len+1,) sequence tensors.  On epoch 2+ the
+       entire dataset is served from RAM, eliminating all numpy I/O.
+       Caching before shuffle means the cache stores the canonical data once;
+       shuffle order changes every epoch without invalidating the cache.
+       Caching after batching would freeze batch composition forever, which
+       breaks per-epoch shuffle semantics.
+
+    3. shuffle()  ← AFTER cache, BEFORE repeat
+       Operates on individual sequences (correct granularity).
+       buffer=min(N, 20_000) bounds peak memory while covering the full
+       dataset for corpora this size.
+
+    4. repeat()  ← AFTER shuffle, BEFORE map/batch
+       Placed here so the shuffle buffer can draw across epoch boundaries,
+       preventing an artificial ordering reset at each epoch edge.
+
+    5. map(split_xy)
+       Slices each (seq_len+1,) window into x=(seq_len,) and y=(seq_len,).
+       No @tf.function decorator needed — tf.data traces the lambda itself.
+       num_parallel_calls=AUTOTUNE parallelises across the CPU thread pool.
+
+    6. batch(drop_remainder=True)
+       Static shape [batch_size, seq_len] — mandatory for XLA/TPU.
+       No symbolic/dynamic shapes ever reach the device.
+
+    7. prefetch(AUTOTUNE)
+       Keeps the device-transfer queue full.  Always last so it buffers
+       fully-formed batches, not individual sequences.
+
+    8. with_options(...)
+       Applies graph-level optimizations in one shot after the pipeline is
+       fully defined.
+    """
+    n = len(windows)
+
+    ds = tf.data.Dataset.from_tensor_slices(windows)           # (1) wrap
+
+    ds = ds.cache(cache_path)                                  # (2) cache
+
+    if shuffle:
+        buf = min(n, 20_000)
+        ds  = ds.shuffle(buffer_size=buf, seed=shuffle_seed,
+                         reshuffle_each_iteration=True)        # (3) shuffle
+
+    ds = ds.repeat()                                           # (4) repeat
+
+    ds = ds.map(                                               # (5) split
+        lambda w: (w[:-1], w[1:]),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+
+    ds = ds.batch(batch_size, drop_remainder=True)             # (6) batch
+
+    ds = ds.prefetch(tf.data.AUTOTUNE)                        # (7) prefetch
+
+    ds = ds.with_options(_make_tf_data_options())              # (8) opts
+
+    return ds
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# BENCHMARK
+# ─────────────────────────────────────────────────────────────────────────────
+
+def benchmark_dataset(
+    ds: tf.data.Dataset,
+    batch_size: int,
+    seq_len: int,
+    n_batches: int = 200,
+    label: str = "dataset",
+) -> dict:
+    """
+    Measure pipeline throughput independently of TPU compute.
+
+    We do NOT call .numpy() inside the loop.  That forces a device→host
+    transfer and a Python GIL acquisition on every batch, serialising what
+    should be a fully-async prefetch chain.  Instead we consume the iterator
+    and let TensorFlow materialise tensors lazily.  The end-to-end wall time
+    is what the TPU will actually experience.
+
+    Interpret results:
+      pipeline_tok/s > 2× tpu_tok/s  →  pipeline is not the bottleneck
+      pipeline_tok/s < 1.5× tpu_tok/s →  still starving the TPU
+    """
+    print(f"\n{'─'*52}")
+    print(f"Benchmarking: {label}")
+    print(f"{'─'*52}")
+
+    # Warm up: populate cache + prefetch buffers before timing starts.
+    for _ in ds.take(5):
+        pass
+    print("Warm-up (5 batches) complete.")
+
+    start = time.perf_counter()
+    for _ in ds.take(n_batches):
+        pass
+    elapsed = time.perf_counter() - start
+
+    tokens_per_batch = batch_size * seq_len
+    total_tokens     = n_batches * tokens_per_batch
+    batches_sec      = n_batches / elapsed
+    tokens_sec       = total_tokens / elapsed
+
+    print(f"Batches measured : {n_batches}")
+    print(f"Elapsed          : {elapsed:.2f} s")
+    print(f"Throughput       : {batches_sec:.1f} batches/s")
+    print(f"Throughput       : {tokens_sec:,.0f} tokens/s")
+
+    tpu_target = 30_000
+    headroom   = tokens_sec / tpu_target
+    if headroom >= 2.0:
+        verdict = f"✓  {headroom:.1f}× headroom over TPU — pipeline is NOT the bottleneck."
+    elif headroom >= 1.2:
+        verdict = f"⚠  {headroom:.1f}× headroom — marginal; increase batch_size or reduce stride."
+    else:
+        verdict = f"✗  {tokens_sec:,.0f} tok/s < TPU target {tpu_target:,} — still starving the TPU."
+
+    print(verdict)
+    print(f"{'─'*52}\n")
+
+    return {"batches_sec": batches_sec, "tokens_sec": tokens_sec,
+            "headroom_vs_30k": headroom}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# THROUGHPUT CALLBACK
+# ──────────────────────────���──────────────────────────────────────────────────
+
+class ThroughputCallback(keras.callbacks.Callback):
+    """
+    Logs throughput every LOG_INTERVAL steps and at epoch end.
+
+    Why coarse logging matters on TPU:
+    Keras callbacks run on the Python host.  JAX dispatches training steps
+    asynchronously — Python returns before XLA finishes the kernel.  A
+    time.perf_counter() call on every step forces a host sync, equivalent
+    to inserting jax.block_until_ready() every step and destroying async
+    pipelining.  At LOG_INTERVAL=50 only ~2% of steps incur this cost.
+    """
+    LOG_INTERVAL = 50
+
+    def on_epoch_begin(self, epoch, logs=None):
+        self._epoch_start    = time.perf_counter()
+        self._interval_start = self._epoch_start
+        self._step_count     = 0
+
+    def on_train_batch_end(self, batch, logs=None):
+        self._step_count += 1
+        if batch > 0 and batch % self.LOG_INTERVAL == 0:
+            now     = time.perf_counter()
+            elapsed = now - self._interval_start
+            # Multiply by effective_batch if data parallelism is active
+            tps     = (self.LOG_INTERVAL * effective_batch * CONTEXT) / elapsed
+            loss    = logs.get('loss', float('nan'))
+            print(f"  step {batch:5d} | loss {loss:.4f} | {tps:>10,.0f} tok/s")
+            self._interval_start = now
+
+    def on_epoch_end(self, epoch, logs=None):
+        elapsed      = time.perf_counter() - self._epoch_start
+        total_tokens = self._step_count * effective_batch * CONTEXT
+        val_loss     = logs.get('val_loss', float('nan'))
+        print(f"\nEpoch {epoch+1} | {elapsed:.1f}s | "
+              f"{total_tokens / elapsed:,.0f} tok/s (avg) | "
+              f"val_loss={val_loss:.4f}\n")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# HELPERS
+# ─────────────────────────────────────────────────────────────────────────────
+
+def load_text_corpus(path: str) -> str:
+    p     = Path(path)
+    files = [p] if p.is_file() else sorted(p.glob('*.txt'))
+    if not files:
+        raise FileNotFoundError(f'No .txt files found in {path}')
+    return ''.join(fp.read_text(encoding='utf-8') for fp in files)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# SHARDING (Data Parallelism for v5e-8)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def setup_data_parallelism():
+    """
+    Configure JAX for pure data parallelism across all TPU devices.
+
+    For v5e-8 (8 chips):
+      - Each chip runs the full model
+      - Different batch shards go to each chip
+      - Gradients are AllReduced across chips
+      - Effective batch = batch_per_chip × num_chips
+
+    This is the simplest and most efficient sharding strategy for models
+    that fit on a single chip.  For 8M parameters, the model easily fits
+    on one v5e chip (16 GB HBM each).
+
+    Returns mesh_shape (tuple) and sharding spec dict for inputs/weights.
+    """
+    devices = jax.devices()
+    n_devices = len(devices)
+
+    if n_devices == 1:
+        print(f"Single device detected ({devices[0]}). Data parallelism disabled.")
+        return None, None
+
+    print(f"\n{'─'*52}")
+    print(f"Data Parallelism Setup (v5e-{n_devices})")
+    print(f"{'─'*52}")
+    print(f"Devices: {n_devices}")
+    print(f"Device type: {devices[0].platform}")
+
+    # Create 1D mesh for data parallelism: each axis is one device
+    mesh = sharding.Mesh(
+        devices=np.array(devices).reshape((n_devices,)),
+        axis_names=("batch",)
+    )
+
+    # Sharding specs for inputs and weights
+    # Input shape: (batch, seq_len)
+    #   Shard batch dimension across devices, keep seq_len replicated
+    input_spec = P("batch", None)
+
+    # Weights: replicate across all devices (each chip has full model)
+    weight_spec = P(None)
+
+    # Activation gradients during backward pass
+    #   Same as input: shard batch, replicate seq
+    activation_spec = P("batch", None)
+
+    print(f"Mesh shape: {mesh.shape}")
+    print(f"Input sharding: {input_spec}  (batch sharded, seq replicated)")
+    print(f"Weight sharding: {weight_spec}  (fully replicated)")
+    print(f"Expected effective batch: {n_devices} × batch_size")
+    print(f"{'─'*52}\n")
+
+    return mesh, {
+        "input": input_spec,
+        "weight": weight_spec,
+        "activation": activation_spec,
+    }
+
+
+def apply_sharding_to_model(model, mesh, sharding_specs):
+    """
+    Apply JAX sharding annotations to a Keras model compiled with JAX backend.
+
+    WARNING: This is JAX-level sharding and only works if:
+    1. Model is built with JAX-native ops (Keras layers with JAX backend)
+    2. jit_compile=True is set in model.compile()
+    3. No TensorFlow-only ops are used in the model
+
+    For Keras models, we set the sharding via jax.Array.with_sharding_constraint()
+    in a custom train step. However, Keras 3 makes this tricky because it wraps
+    training in its own jit.
+
+    Simpler approach: set jax.config to globally use this mesh, and let XLA
+    infer sharding from the mesh context.
+    """
+    if mesh is None:
+        return
+
+    # Tell JAX to use this mesh for all operations
+    with mesh:
+        print("Sharding configuration loaded.")
+        print(f"  All matmul ops will shard batch dim across {mesh.shape[0]} devices")
+        print(f"  Gradient AllReduce will use {mesh.shape[0]}-way ICI collective\n")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# MAIN
+# ─────────────────────────────────────────────────────────────────────────────
+class MemoryCallback(keras.callbacks.Callback):
+    def on_epoch_end(self, epoch, logs=None):
+        stats = jax.devices()[0].memory_stats()
+
+        used = stats["bytes_in_use"] / 1024**3
+        peak = stats["peak_bytes_in_use"] / 1024**3
+        limit = stats["bytes_limit"] / 1024**3
+
+        print(
+            f"\nHBM: {used:.2f}/{limit:.2f} GB "
+            f"(Peak: {peak:.2f} GB)"
+        )
+def main():
+    print('\n' + '=' * 60)
+    print('BACKEND & DEVICE VERIFICATION')
+    print('=' * 60)
+    print(f'Keras backend  : {keras.backend.backend()}')
+    print(f'JAX devices    : {jax.devices()}')
+    print('=' * 60 + '\n')
+
+    # ── Setup data parallelism ─────────────────────────────────────────────
+    mesh, sharding_specs = setup_data_parallelism()
+    if mesh is not None:
+        apply_sharding_to_model(None, mesh, sharding_specs)
+
+    # ── Tokenize ──────────────────────────────────────────────────────────
+    tokenizer  = TokenizerWrapper('tokenizer.model')
+    text       = load_text_corpus(data_path)
+    raw_tokens = np.asarray(
+        tokenizer.encode(text, add_bos=True, add_eos=False), dtype=np.int32
+    )
+    print(f'Total tokens: {len(raw_tokens):,}')
+
+    split_idx    = int(len(raw_tokens) * 0.98)
+    train_tokens = raw_tokens[:split_idx]
+    val_tokens   = raw_tokens[split_idx:]
+
+    stride = max(1, CONTEXT // 2)
+
+    # ── Pre-compute windows ────────────────────────────────────────────────
+    print("Pre-computing windows...", end=" ", flush=True)
+    t0            = time.perf_counter()
+    train_windows = precompute_windows(train_tokens, CONTEXT, stride)
+    val_windows   = precompute_windows(val_tokens,   CONTEXT, stride)
+    print(f"done in {(time.perf_counter() - t0) * 1000:.1f} ms")
+    print(f"  train: {train_windows.shape}  ({train_windows.nbytes / 1e6:.1f} MB)")
+    print(f"  val:   {val_windows.shape}  ({val_windows.nbytes / 1e6:.1f} MB)\n")
+
+    # ── Build datasets ─────────────────────────────────────────────────────
+    # If data parallelism is active, the effective batch becomes:
+    #   batch_per_device × num_devices
+    # But tf.data still sees batch_per_device — JAX handles the replication.
+    num_devices = len(jax.devices())
+    actual_batch_size = batch_size  # Per-device batch
+    effective_batch = actual_batch_size * num_devices if mesh is not None else actual_batch_size
+
+    if mesh is not None:
+        print(f"Effective batch size: {actual_batch_size} × {num_devices} devices = {effective_batch}")
+
+    train_ds = create_tf_dataset(train_windows, actual_batch_size, shuffle=True)
+    val_ds   = create_tf_dataset(val_windows,   actual_batch_size, shuffle=False)
+
+    # ── Benchmark before training ────────────────────────────────────────��─
+    # Verifies the pipeline can outpace the TPU before we waste a run.
+    benchmark_dataset(train_ds, batch_size, CONTEXT, n_batches=200,
+                      label="train_ds")
+
+    # ── Build model ────────────────────────────────────────────────────────
+    os.makedirs('checkpoints', exist_ok=True)
+
+    model = create_llm(
+        vocab_size   = tokenizer.vocab_size,
+        d_model      = D_MODEL,
+        n_layers     = numberoflayers,
+        n_heads      = numberofheads,
+        d_latent     = d_Latent,
+        ffn_mult     = ffn_mult,
+        max_seq_len  = CONTEXT,
+        use_moe      = False,
+        num_kv_heads = num_kv_heads,
+        swa_window   = swa_window,
+    )
+
+    optimizer = keras.optimizers.AdamW(
+        learning_rate   = learning_rate,
+        weight_decay    = weight_decay,
+        global_clipnorm = 1.0,
+    )
+    model.compile(
+        optimizer  = optimizer,
+        loss       = keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        jit_compile = True,
+    )
+    
+    # ── Warmup with exact training shape ──────────────────────────────────
+    # XLA traces a kernel per unique input shape.  Warming up with any other
+    # shape (e.g. the original (1, min(4, CONTEXT))) causes a full recompile
+    # on the first real training batch — wasting 30-120 s on the TPU.
+    print("Warming up XLA with training shape...", end=" ", flush=True)
+    t0      = time.perf_counter()
+    dummy_x = np.zeros((batch_size, CONTEXT), dtype=np.int32)
+    _       = model(dummy_x, training=False)
+    _       = model(dummy_x, training=True)
+    print(f"done in {time.perf_counter() - t0:.2f}s\n")
+
+    # ── Forward pass sanity check ──────────────────────────────────────────
+    # One explicit transfer (np.array) then work in NumPy — avoids the 4
+    # implicit device→host round trips of the original float(ops.min(...))
+    # scalar-cast pattern.
+    print('─' * 52)
+    print('Forward pass validation')
+    print('─' * 52)
+    for x_batch, _ in train_ds.take(1):
+        logits    = model(x_batch[:1], training=False)
+        logits_np = np.array(logits)
+        print(f'Input  : {x_batch.shape}')
+        print(f'Logits : {logits.shape}')
+        print(f'min={logits_np.min():.4f}  max={logits_np.max():.4f}')
+        if np.isnan(logits_np).any() or np.isinf(logits_np).any():
+            raise ValueError('Forward pass produced NaN/Inf — check model init.')
+    print('✓ Forward pass clean.\n')
+    
+    # ── Steps per epoch from actual window count ───────────────────────────
+    # Original used len(tokens) // (batch*seq) which undercounts overlapping
+    # windows and causes epochs to terminate prematurely.
+    print(f"Total params: {model.count_params():,}")
+    steps_per_epoch  = max(1, len(train_windows) // batch_size)
+    validation_steps = max(1, len(val_windows)   // batch_size)
+    print(f"steps_per_epoch  : {steps_per_epoch}")
+    print(f"validation_steps : {validation_steps}\n")
+
+    # ── Train ──────────────────────────────────────────────────────────────
+    model.fit(
+        train_ds,
+        validation_data  = val_ds,
+        epochs           = EPOCHS,
+        steps_per_epoch  = steps_per_epoch,
+        validation_steps = validation_steps,
+        callbacks        = [
+            keras.callbacks.TerminateOnNaN(),
+            keras.callbacks.ModelCheckpoint(
+                filepath          = 'checkpoints/veylon_{epoch:02d}.weights.h5',
+                save_freq         = 'epoch',
+                save_weights_only = True,
+            ),
+            ThroughputCallback(),
+            MemoryCallback(),
+        ],
+    )
+
+    model.save_weights('veylon_final.weights.h5')
+    print('Saved veylon_final.weights.h5')
+
+
+if __name__ == '__main__':
+    main()

veylon_attention.py

@@ -0,0 +1,640 @@
+"""
+veylon_attention.py  —  Veylon Alpha 1
+======================================
+Block-tiled Native GQA Sliding Window Attention for JAX / Keras-3 / TPU v5e.
+
+Why NOT lax.scan + dynamic_slice
+---------------------------------
+The previous implementation used:
+
+    jax.lax.scan(step, None, (jnp.arange(S), q_scan))
+
+with `dynamic_slice(k_pad, (0,0,t,0), (B,Hkv,W,D))` inside the body.
+
+XLA/XLA-TPU has a known pathology with this pattern:
+  - The scan body contains a dynamic gather (dynamic_slice on a traced int `t`).
+  - XLA's while-loop lowering stages ALL window materializations into a single
+    large buffer to enable pipeline prefetching.
+  - On TPU this produces a hidden tensor  [S, B, Hkv, W, D]  which is exactly
+    what you see in the OOM allocation log:  f32[1024,8,4,512,64].
+  - jax.checkpoint does NOT protect against this because it is a compiler
+    (HLO-level) allocation, not a JAX-level rematerialization artifact.
+
+The fix: block-tiled computation with fully static tensor shapes
+----------------------------------------------------------------
+Instead of iterating over S individual tokens we iterate over
+(S / BLK) blocks of queries.  For each block:
+
+  - q_blk  : [B, Hq,  BLK,        D]   <- static shape
+  - k_blk  : [B, Hkv, BLK + W - 1, D]  <- static shape
+  - scores : [B, Hkv, G, BLK, BLK+W-1] <- static shape
+
+XLA sees ONLY static shapes inside the map body.  There is no
+gather-inside-loop pattern.  XLA can freely fuse, pipeline, and
+tile these einsums onto TPU systolic arrays without hidden buffers.
+
+lax.map vs lax.scan
+--------------------
+We use `jax.lax.map` (not `lax.scan`) because:
+  - Each block is fully independent — no carry state is needed.
+  - `lax.map` lowers to a while_loop with no accumulation buffer.
+  - `lax.scan` always allocates an output stacked along the scan axis;
+    lax.map lets XLA write directly into the preallocated output slice.
+  - This eliminates the [n_blocks, B, Hq, BLK, D] intermediate stack.
+    (We still do one reshape at the end, but that is a view, not a copy.)
+
+Tensor size invariants
+----------------------
+NEVER created:
+    [B, Hq,  S, S]               full attention score matrix
+    [S, B, Hkv, W, D]            per-token window stack (old scan bug)
+    [B, Hq, Hkv, S, D]           duplicated KV heads
+
+Largest tensors inside map body (all STATIC shapes):
+    k_blk, v_blk  : [B, Hkv, BLK+W-1, D]
+    scores        : [B, Hkv, G, BLK, BLK+W-1]
+    probs         : same
+
+Memory complexity
+-----------------
+  k_pad, v_pad  : O(B × Hkv × (S + W) × D)   linear in S, scales with Hkv
+  q             : O(B × Hq  × S × D)
+  Per block     : O(B × Hkv × (BLK + W) × D)  constant w.r.t. S
+  Output        : O(B × Hq  × S × D)
+  TOTAL         : O(B × (Hkv + Hq) × S × D)   strictly linear in S
+
+TPU-specific notes
+------------------
+  * All shapes in map body are concrete at trace time — XLA never inserts
+    shape-dependent conditionals or recompiles.
+  * BF16 inputs are cast to FP32 before matmul/softmax, then cast back.
+  * BLK should be a multiple of 128 on TPU v5e for optimal systolic array
+    utilization (default 128, tunable via `block_size` parameter).
+  * The precomputed mask is a static bool array passed as a closed-over
+    constant — XLA fuses it into the einsum kernel at no extra memory cost.
+  * jax.checkpoint on the map body protects backward-pass activations
+    (one block at a time, not one token at a time — much coarser remat).
+"""
+
+from __future__ import annotations
+
+import math
+from functools import partial
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+
+
+# ---------------------------------------------------------------------------
+# Public tuning constant
+# ---------------------------------------------------------------------------
+
+# Default block size for TPU v5e.  Must be a power of 2 and ≥ 1.
+# Larger blocks → fewer kernel launches, better systolic utilisation.
+# Smaller blocks → lower peak HBM per block (useful when W is huge).
+TPU_BLOCK_SIZE: int = 128
+
+
+# ---------------------------------------------------------------------------
+# Utility helpers
+# ---------------------------------------------------------------------------
+
+def _next_power_of_2(x: int) -> int:
+    x = int(x)
+    if x <= 1:
+        return 1
+    return 1 << (x - 1).bit_length()
+
+
+def apply_rope(
+    x: jnp.ndarray,
+    cos: jnp.ndarray,
+    sin: jnp.ndarray,
+    offset: int = 0,
+) -> jnp.ndarray:
+    """
+    Apply Rotary Position Encoding to [..., S, D].
+    cos/sin tables are pre-built for D//2 (half the head dim).
+    Uses dynamic_slice — safe under jit and XLA outside of attention body.
+    """
+    if x.ndim < 2:
+        raise ValueError(f"apply_rope: expected ≥2-D input, got shape {x.shape}")
+    d = int(x.shape[-1])
+    if d % 2 != 0:
+        raise ValueError(f"apply_rope: head dim must be even, got {d}")
+    half = d // 2
+    seq_len = int(x.shape[-2])
+    if offset + seq_len > int(cos.shape[0]):
+        raise ValueError(
+            f"apply_rope: RoPE table too small — "
+            f"offset={offset}, seq_len={seq_len}, table_size={cos.shape[0]}"
+        )
+    cos_s = jax.lax.dynamic_slice(cos, (offset, 0), (seq_len, half)).astype(x.dtype)
+    sin_s = jax.lax.dynamic_slice(sin, (offset, 0), (seq_len, half)).astype(x.dtype)
+    while cos_s.ndim < x.ndim - 1:
+        cos_s = cos_s[None]
+        sin_s = sin_s[None]
+    x1, x2 = x[..., :half], x[..., half:]
+    return jnp.concatenate(
+        [x1 * cos_s - x2 * sin_s, x1 * sin_s + x2 * cos_s],
+        axis=-1,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Core kernel: block-tiled native GQA SWA
+# ---------------------------------------------------------------------------
+
+def _block_gqa_swa(
+    q: jnp.ndarray,
+    k: jnp.ndarray,
+    v: jnp.ndarray,
+    window_size: int,
+    block_size: int = TPU_BLOCK_SIZE,
+    use_remat: bool = True,
+) -> jnp.ndarray:
+    """
+    Block-tiled Native GQA Sliding Window Attention.
+
+    This function is the replacement for the scan+dynamic_slice approach.
+    All tensor shapes inside the map body are STATIC — XLA never sees a
+    gather-inside-loop pattern, so the hidden [S,B,H,W,D] buffer cannot form.
+
+    Parameters
+    ----------
+    q           : [B, Hq,  S, D]  — any float dtype (BF16 in production)
+    k           : [B, Hkv, S, D]  — same dtype
+    v           : [B, Hkv, S, D]  — same dtype
+    window_size : causal window W; token t attends to [max(0, t-W+1), t]
+    block_size  : query block size BLK (tune to 128 for TPU v5e)
+    use_remat   : wrap map body in jax.checkpoint (recommended for training)
+
+    Returns
+    -------
+    [B, Hq, S, D] — same dtype as q
+    """
+    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
+        raise ValueError(
+            f"_block_gqa_swa: expected 4-D inputs, "
+            f"got q={q.shape}  k={k.shape}  v={v.shape}"
+        )
+
+    B,  Hq,  S,  D   = map(int, q.shape)
+    _B, Hkv, Sk, Dk  = map(int, k.shape)
+
+    if _B != B:
+        raise ValueError(f"Batch mismatch: q B={B}, k B={_B}")
+    if Dk != D:
+        raise ValueError(f"Head-dim mismatch: q D={D}, k D={Dk}")
+    if Sk != S:
+        raise ValueError(f"Sequence-length mismatch: q S={S}, k S={Sk}")
+    if Hq % Hkv != 0:
+        raise ValueError(f"Hq={Hq} must be divisible by Hkv={Hkv}")
+
+    G   = Hq // Hkv                  # queries per KV head
+    W   = int(window_size)
+    BLK = int(block_size)
+    if BLK <= 0:
+        raise ValueError(f"block_size must be > 0, got {BLK}")
+
+    scale = 1.0 / math.sqrt(float(D))
+
+    # ── Pad sequence to a multiple of BLK ────────────────────────────────────
+    n_blocks = (S + BLK - 1) // BLK
+    S_pad    = n_blocks * BLK          # ≥ S, multiple of BLK
+
+    # ── Pad q along sequence axis ─────────────────────────────────────────────
+    # [B, Hq, S_pad, D]
+    # The extra (S_pad - S) tokens are zero-padded and trimmed from output.
+    q_pad = jnp.pad(q, ((0, 0), (0, 0), (0, S_pad - S), (0, 0)))
+
+    # ── Pad k/v on the LEFT by (W-1) for causal alignment ────────────────────
+    # [B, Hkv, S_pad + W - 1, D]
+    #
+    # After padding, for query block starting at global position b*BLK:
+    #   kv slice starts at offset b*BLK in k_pad
+    #   kv slice has static length  BLK + W - 1
+    #   It covers original positions  [b*BLK - (W-1), b*BLK + BLK - 1]
+    #   which, after clipping to ≥ 0, is exactly the causal window.
+    kv_pad_len = S_pad + W - 1        # total padded kv length (static)
+    lpad       = W - 1                # left zero-padding width
+
+    k_pad = jnp.pad(k, ((0, 0), (0, 0), (lpad, S_pad - S), (0, 0)))  # [B,Hkv,kv_pad_len,D]
+    v_pad = jnp.pad(v, ((0, 0), (0, 0), (lpad, S_pad - S), (0, 0)))
+
+    # ── Static per-block mask ─────────────────────────────────────────────────
+    # Compute a boolean mask of shape [BLK, BLK + W - 1].
+    # Entry [q_local, k_local] is True (mask=attend) when:
+    #   (a) k is not from left-padding   → k_local >= W - 1 - (something)
+    # We cannot compute the absolute positions statically because they depend
+    # on block index b.  So we record the RELATIVE offsets and apply the
+    # offset inside the map body using only static arithmetic on local indices.
+    #
+    # q_abs  = b*BLK + q_local          (q_local in 0..BLK-1)
+    # k_abs  = b*BLK + k_local - lpad   (k_local in 0..BLK+W-2)
+    #
+    # Attend iff:
+    #   k_abs >= 0                (not left padding)
+    #   k_abs <= q_abs            (causal)
+    #   q_abs - k_abs < W         (within window)
+    #
+    # q_abs - k_abs  = q_local - k_local + lpad   (b cancels out!)
+    # k_abs >= 0     ↔ k_local >= lpad - b*BLK   (depends on b → handle in body)
+    # k_abs <= q_abs ↔ k_local - q_local <= lpad  (b cancels out! → STATIC)
+    #
+    # Only the "k_abs >= 0" condition depends on b (first block only).
+    # We handle it cheaply with a dynamic mask inside the body.
+    # Everything else is b-independent and can be PRECOMPUTED once.
+
+    kv_len = BLK + W - 1   # static length of kv slice per block
+
+    q_local  = jnp.arange(BLK,    dtype=jnp.int32)   # [BLK]
+    kv_local = jnp.arange(kv_len, dtype=jnp.int32)   # [kv_len]
+
+    # Relative offset: delta[q, k] = q_local[q] - kv_local[k] + lpad
+    #   = q_abs - k_abs    (b-independent)
+    delta = q_local[:, None] - kv_local[None, :] + lpad  # [BLK, kv_len]
+
+    # Static masks (b-independent)
+    static_future_mask  = delta < 0          # k is in the future (causal)  → mask
+    static_window_mask  = delta >= W         # k is too far back (window)   → mask
+    static_base_mask    = static_future_mask | static_window_mask  # [BLK, kv_len]
+
+    # ── Map body ──────────────────────────────────────────────────────────────
+
+    def process_block(b: jnp.ndarray) -> jnp.ndarray:
+        """
+        b : scalar int32 block index in [0, n_blocks)
+
+        Tensor shapes inside this function are ALL STATIC:
+            q_blk   : [B, Hq,  BLK,    D]
+            k_blk   : [B, Hkv, kv_len, D]
+            v_blk   : [B, Hkv, kv_len, D]
+            scores  : [B, Hkv, G, BLK, kv_len]
+            probs   : same
+
+        XLA has NO gather-inside-loop here.  The dynamic_slice start index
+        `b * BLK` is a scalar multiply — XLA lowers this to a simple pointer
+        offset, not a buffer materialisation.
+        """
+        blk_start = b * BLK  # scalar traced int32
+
+        # Static-shape slices — the KEY difference from the scan approach.
+        # XLA sees shapes (B,Hq,BLK,D) and (B,Hkv,kv_len,D) as compile-time
+        # constants.  It cannot stage all blocks simultaneously because
+        # lax.map gives it one block at a time with no output accumulation.
+        q_blk = jax.lax.dynamic_slice(q_pad, (0, 0, blk_start,        0), (B, Hq,  BLK,    D))
+        k_blk = jax.lax.dynamic_slice(k_pad, (0, 0, blk_start,        0), (B, Hkv, kv_len, D))
+        v_blk = jax.lax.dynamic_slice(v_pad, (0, 0, blk_start,        0), (B, Hkv, kv_len, D))
+
+        # ── Native GQA reshape (zero-copy view) ──────────────────────────
+        # [B, Hq, BLK, D] → [B, Hkv, G, BLK, D]
+        q_g = q_blk.reshape(B, Hkv, G, BLK, D)
+
+        # ── Dot-product scores ────────────────────────────────────────────
+        # [B, Hkv, G, BLK, kv_len]   ← STATIC, FUSED by XLA
+        scores = (
+            jnp.einsum(
+                "bngqd,bnkd->bngqk",
+                q_g.astype(jnp.float32),
+                k_blk.astype(jnp.float32),
+            )
+            * scale
+        )
+
+        # ── Masking ───────────────────────────────────────────────────────
+        # (a) b-independent mask (precomputed, fused as constant)
+        mask = static_base_mask  # [BLK, kv_len]
+
+        # (b) Left-padding mask: k_abs < 0  ↔  kv_local < lpad - blk_start
+        #     Only non-trivial for b=0 (the very first block).
+        #     For all subsequent blocks, lpad - blk_start < 0, so no extra masking.
+        #     We compute it dynamically but it is a single scalar comparison
+        #     broadcast — XLA will constant-fold it for b > 0 at runtime.
+        leftpad_cutoff = lpad - blk_start   # scalar int32 (may be negative)
+        leftpad_mask   = kv_local[None, :] < leftpad_cutoff  # [1, kv_len]
+        mask = mask | leftpad_mask  # [BLK, kv_len]
+
+        scores = jnp.where(
+            mask[None, None, None, :, :],   # [1,1,1,BLK,kv_len]
+            jnp.full_like(scores, -1e30),
+            scores,
+        )
+
+        # ── Softmax + value aggregation (all FP32) ────────────────────────
+        probs  = jax.nn.softmax(scores, axis=-1)            # [B,Hkv,G,BLK,kv_len]
+        out_g  = jnp.einsum(
+            "bngqk,bnkd->bngqd",
+            probs,
+            v_blk.astype(jnp.float32),
+        )  # [B, Hkv, G, BLK, D]
+
+        # ── Reshape + cast back to input dtype ────────────────────────────
+        return out_g.reshape(B, Hq, BLK, D).astype(q.dtype)  # [B, Hq, BLK, D]
+
+    # ── Gradient checkpointing ────────────────────────────────────────────────
+    # With use_remat=True: JAX recomputes each block's forward during backward.
+    # Granularity is per-block (not per-token), so remat overhead is manageable.
+    map_fn = jax.checkpoint(process_block) if use_remat else process_block
+
+    # ── lax.map over block indices ────────────────────────────────────────────
+    # Unlike lax.scan, lax.map does NOT accumulate a carry — it writes each
+    # block's output directly.  No hidden [n_blocks, B, Hq, BLK, D] stack.
+    out_blocks = jax.lax.map(map_fn, jnp.arange(n_blocks, dtype=jnp.int32))
+    # out_blocks: [n_blocks, B, Hq, BLK, D]
+
+    # ── Assemble output ───────────────────────────────────────────────────────
+    # Transpose to [B, Hq, n_blocks, BLK, D] then reshape to [B, Hq, S_pad, D].
+    # The final slice trims padding tokens.  All zero-copy views on TPU.
+    out = out_blocks.transpose(1, 2, 0, 3, 4).reshape(B, Hq, S_pad, D)
+    return out[:, :, :S, :]
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def flash_splash_attention(
+    q: jnp.ndarray,
+    k: jnp.ndarray,
+    v: jnp.ndarray,
+    window_size: int,
+    backend: Optional[str] = None,
+    use_gqa: bool = True,
+    start_pos: int = 0,
+    block_size: int = TPU_BLOCK_SIZE,
+    use_remat: bool = True,
+) -> jnp.ndarray:
+    """
+    Block-tiled GQA-native Sliding Window Attention — main entry point.
+
+    Drop-in replacement for the previous flash_splash_attention.
+    GQA is always native (no KV expansion).  `use_gqa` is accepted for
+    API compatibility only.
+
+    Parameters
+    ----------
+    q, k, v      : [B, Hq, S, D] / [B, Hkv, S, D]
+    window_size  : causal window W
+    backend      : reserved for future Pallas/Flash kernel dispatch (no-op)
+    use_gqa      : API-compat flag; GQA is always native here
+    start_pos    : KV-cache offset for decode mode (training path: leave at 0)
+    block_size   : query block size BLK (128 recommended for TPU v5e)
+    use_remat    : gradient checkpointing in map body (recommended for training)
+
+    Returns
+    -------
+    [B, Hq, S, D] — same dtype as q
+    """
+    _ = backend
+    _ = use_gqa
+    _ = start_pos  # TODO: shift q/kv slicing for decode-mode KV cache
+
+    return _block_gqa_swa(
+        q=q,
+        k=k,
+        v=v,
+        window_size=int(window_size),
+        block_size=int(block_size),
+        use_remat=use_remat,
+    )
+def decode_swa(
+    q: jnp.ndarray,
+    k: jnp.ndarray,
+    v: jnp.ndarray,
+) -> jnp.ndarray:
+    """
+    Decode-time SWA for one-token generation.
+
+    q: [B, Hq, 1, D]
+    k: [B, Hkv, W, D]
+    v: [B, Hkv, W, D]
+
+    Returns: [B, Hq, 1, D]
+    """
+    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
+        raise ValueError(
+            f"decode_swa: expected 4-D tensors, got q={q.shape}, k={k.shape}, v={v.shape}"
+        )
+
+    B, Hq, S, D = map(int, q.shape)
+    _, Hkv, W, Dk = map(int, k.shape)
+
+    if S != 1:
+        raise ValueError(f"decode_swa expects one token, got S={S}")
+    if D != Dk:
+        raise ValueError(f"Head-dim mismatch: q D={D}, k D={Dk}")
+    if Hq % Hkv != 0:
+        raise ValueError(f"Hq={Hq} must be divisible by Hkv={Hkv}")
+
+    G = Hq // Hkv
+    scale = 1.0 / math.sqrt(float(D))
+
+    q_g = q.reshape(B, Hkv, G, 1, D)
+
+    scores = (
+        jnp.einsum(
+            "bngqd,bnkd->bngqk",
+            q_g.astype(jnp.float32),
+            k.astype(jnp.float32),
+        )
+        * scale
+    )  # [B, Hkv, G, 1, W]
+
+    probs = jax.nn.softmax(scores, axis=-1)
+
+    out_g = jnp.einsum(
+        "bngqk,bnkd->bngqd",
+        probs,
+        v.astype(jnp.float32),
+    )
+
+    return out_g.reshape(B, Hq, 1, D).astype(q.dtype)
+
+def local_causal_attention(
+    q: jnp.ndarray,
+    k: jnp.ndarray,
+    v: jnp.ndarray,
+) -> jnp.ndarray:
+    """
+    Full causal attention with native GQA — no windowing.
+
+    ⚠  Creates an O(S²) score tensor [B, Hkv, G, S, S].
+       Use only for short sequences, unit tests, or reference baselines.
+
+    q: [B, Hq, S, D]  |  k, v: [B, Hkv, S, D]
+    returns: [B, Hq, S, D]
+    """
+    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
+        raise ValueError(
+            f"local_causal_attention: expected 4-D tensors, "
+            f"got q={q.shape}  k={k.shape}  v={v.shape}"
+        )
+    B, Hq, S, D = map(int, q.shape)
+    _,  Hkv, K, _ = map(int, k.shape)
+    if Hq % Hkv != 0:
+        raise ValueError(f"Hq={Hq} must be divisible by Hkv={Hkv}")
+    G = Hq // Hkv
+    scale = 1.0 / math.sqrt(float(D))
+    q_g   = q.reshape(B, Hkv, G, S, D)
+    scores = (
+        jnp.einsum("bngsd,bnkd->bngsk", q_g.astype(jnp.float32), k.astype(jnp.float32))
+        * scale
+    )
+    qi = jnp.arange(S, dtype=jnp.int32)[:, None]
+    ki = jnp.arange(K, dtype=jnp.int32)[None, :]
+    scores = jnp.where(ki > qi, -1e30, scores)
+    probs  = jax.nn.softmax(scores, axis=-1)
+    out_g  = jnp.einsum("bngsk,bnkd->bngsd", probs, v.astype(jnp.float32))
+    return out_g.reshape(B, Hq, S, D).astype(q.dtype)
+
+
+# ---------------------------------------------------------------------------
+# Validation suite
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import sys
+
+    PASS  = "\033[92m✓\033[0m"
+    FAIL  = "\033[91m✗\033[0m"
+    HDR   = "\033[1;94m"
+    RST   = "\033[0m"
+    failures = 0
+
+    def section(t):  print(f"\n{HDR}{'─'*62}{RST}\n{HDR}  {t}{RST}\n{HDR}{'─'*62}{RST}")
+    def ok(m):       print(f"  {PASS}  {m}")
+    def fail(m):
+        global failures; failures += 1
+        print(f"  {FAIL}  {m}", file=sys.stderr)
+
+    print(f"\n{HDR}{'═'*62}{RST}")
+    print(f"{HDR}  Veylon Attention — Block-tiled GQA SWA — Validation{RST}")
+    print(f"{HDR}{'═'*62}{RST}")
+
+    # ── 1. Shape and dtype ───────────────────────────────────────────────────
+    section("1 · Shape and dtype correctness")
+
+    B, Hq, Hkv, S, D, W = 1, 8, 2, 128, 64, 32
+    ks = jax.random.split(jax.random.PRNGKey(0), 3)
+    q = jax.random.normal(ks[0], (B, Hq,  S, D), dtype=jnp.bfloat16)
+    k = jax.random.normal(ks[1], (B, Hkv, S, D), dtype=jnp.bfloat16)
+    v = jax.random.normal(ks[2], (B, Hkv, S, D), dtype=jnp.bfloat16)
+
+    out = flash_splash_attention(q, k, v, window_size=W, block_size=32)
+
+    if out.shape == (B, Hq, S, D):   ok(f"Output shape  : {out.shape}")
+    else:                             fail(f"Shape wrong — expected {(B,Hq,S,D)}, got {out.shape}")
+    if out.dtype == jnp.bfloat16:    ok(f"Output dtype  : {out.dtype}")
+    else:                             fail(f"dtype wrong — expected bfloat16, got {out.dtype}")
+    if not jnp.any(jnp.isnan(out)):  ok("No NaNs")
+    else:                             fail("Output contains NaNs")
+
+    # ── 2. Numerical agreement with reference SWA ────────────────────────────
+    section("2 · Numerical agreement: block SWA ≈ reference SWA (W=S)")
+
+    B_, Hq_, Hkv_, S_, D_ = 1, 4, 2, 48, 16
+    ks = jax.random.split(jax.random.PRNGKey(1), 3)
+    qn = jax.random.normal(ks[0], (B_, Hq_,  S_, D_), dtype=jnp.float32)
+    kn = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32)
+    vn = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32)
+
+    for W_t in [4, 8, 16, S_]:
+        out_blk  = flash_splash_attention(qn, kn, vn, window_size=W_t, block_size=16)
+        out_ref  = local_causal_attention(qn, kn, vn) if W_t == S_ else None
+
+        # Build reference SWA inline for each W_t
+        G_ = Hq_ // Hkv_
+        sc = 1.0 / math.sqrt(D_)
+        qg = qn.reshape(B_, Hkv_, G_, S_, D_)
+        sc_ref = jnp.einsum("bngsd,bnkd->bngsk", qg, kn) * sc
+        qi = jnp.arange(S_)[:, None]; ki = jnp.arange(S_)[None, :]
+        sc_ref = jnp.where((ki > qi) | (qi - ki >= W_t), -1e30, sc_ref)
+        pr_ref = jax.nn.softmax(sc_ref, axis=-1)
+        out_swa_ref = jnp.einsum("bngsk,bnkd->bngsd", pr_ref, vn).reshape(B_, Hq_, S_, D_)
+
+        err = float(jnp.max(jnp.abs(out_blk - out_swa_ref)))
+        if err < 1e-4:  ok(f"W={W_t:3d}  max|block - ref_swa| = {err:.2e}")
+        else:           fail(f"W={W_t:3d}  MISMATCH: max err = {err:.2e}")
+
+    # ── 3. Memory scaling ────────────────────────────────────────────────────
+    section("3 · Memory scaling  (2× S → ~2× memory, not 4×)")
+    print("  Shape + completion check at S = 256, 512, 1024, 2048")
+    for S_t in [256, 512, 1024, 2048]:
+        ks = jax.random.split(jax.random.PRNGKey(S_t), 3)
+        qt = jax.random.normal(ks[0], (1, 8, S_t, 64), dtype=jnp.bfloat16)
+        kt = jax.random.normal(ks[1], (1, 2, S_t, 64), dtype=jnp.bfloat16)
+        vt = jax.random.normal(ks[2], (1, 2, S_t, 64), dtype=jnp.bfloat16)
+        ot = flash_splash_attention(qt, kt, vt, window_size=128)
+        if ot.shape == (1, 8, S_t, 64): ok(f"S={S_t:5d}  →  {ot.shape}")
+        else:                            fail(f"S={S_t}  wrong shape {ot.shape}")
+
+    # ── 4. Native GQA isolation ──────────────────────────────────────────────
+    section("4 �� Native GQA head isolation (no KV duplication)")
+    B_, Hq_, Hkv_, S_, D_, W_ = 1, 4, 2, 32, 16, 16
+    G_ = Hq_ // Hkv_
+    ks  = jax.random.split(jax.random.PRNGKey(7), 3)
+    qg  = jax.random.normal(ks[0], (B_, Hq_, S_, D_), dtype=jnp.float32)
+    kgq = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32)
+    vgq = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32)
+    base  = flash_splash_attention(qg, kgq,                               vgq,                               window_size=W_, block_size=16)
+    zero0 = flash_splash_attention(qg, kgq.at[:,0].set(0.), vgq.at[:,0].set(0.), window_size=W_, block_size=16)
+    changed   = not jnp.allclose(base[:, :G_],  zero0[:, :G_],  atol=1e-4)
+    unchanged = jnp.allclose(    base[:, G_:],  zero0[:, G_:],  atol=1e-4)
+    ok(f"Q heads 0..{G_-1} changed when KV head 0 zeroed")   if changed   else fail("GQA dependency broken")
+    ok(f"Q heads {G_}..{Hq_-1} unchanged (correct isolation)") if unchanged else fail("Cross-group contamination")
+
+    # ── 5. Causality ─────────────────────────────────────────────────────────
+    section("5 · Causality")
+    B_, Hq_, Hkv_, S_, D_, W_ = 1, 4, 2, 24, 8, 8
+    pivot = S_ // 2
+    ks = jax.random.split(jax.random.PRNGKey(42), 3)
+    qc = jax.random.normal(ks[0], (B_, Hq_, S_, D_), dtype=jnp.float32)
+    kc = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32)
+    vc = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32)
+    out_base = flash_splash_attention(qc, kc, vc, window_size=W_, block_size=8)
+    out_corr = flash_splash_attention(qc,
+                                      kc.at[:,:,pivot:].set(999.),
+                                      vc.at[:,:,pivot:].set(999.),
+                                      window_size=W_, block_size=8)
+    if jnp.allclose(out_base[:,:,:pivot], out_corr[:,:,:pivot], atol=1e-5):
+        ok(f"Tokens 0..{pivot-1} unaffected by corruption of tokens {pivot}+")
+    else:
+        fail("Causality violated — future tokens leaked into past outputs")
+
+    # ── 6. Window boundary ───────────────────────────────────────────────────
+    section("6 · Window boundary  (no leakage past W tokens)")
+    B_, Hq_, Hkv_, S_, D_, W_ = 1, 2, 1, 24, 8, 4
+    ks = jax.random.split(jax.random.PRNGKey(11), 3)
+    qw = jax.random.normal(ks[0], (B_, Hq_, S_, D_), dtype=jnp.float32)
+    kw = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32)
+    vw = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32)
+    out_w  = flash_splash_attention(qw, kw, vw, window_size=W_, block_size=4)
+    out_wm = flash_splash_attention(qw,
+                                    kw.at[:,:,0].set(999.),
+                                    vw.at[:,:,0].set(999.),
+                                    window_size=W_, block_size=4)
+    pivot_w = W_  # first token where token 0 is outside the window
+    if jnp.allclose(out_w[:,:,pivot_w:], out_wm[:,:,pivot_w:], atol=1e-5):
+        ok(f"Tokens {pivot_w}+ unaffected by modifying token 0  (W={W_} boundary)")
+    else:
+        fail(f"Window boundary violated — token 0 leaked into token {pivot_w}+")
+
+    # ── 7. Non-power-of-2 sequence length ────────────────────────────────────
+    section("7 · Non-power-of-2 sequence lengths  (S=100, 200, 500)")
+    for S_t in [100, 200, 500]:
+        ks = jax.random.split(jax.random.PRNGKey(S_t+1), 3)
+        qt = jax.random.normal(ks[0], (1, 4, S_t, 16), dtype=jnp.float32)
+        kt = jax.random.normal(ks[1], (1, 2, S_t, 16), dtype=jnp.float32)
+        vt = jax.random.normal(ks[2], (1, 2, S_t, 16), dtype=jnp.float32)
+        ot = flash_splash_attention(qt, kt, vt, window_size=32, block_size=32)
+        if ot.shape == (1, 4, S_t, 16): ok(f"S={S_t}  →  {ot.shape}")
+        else:                            fail(f"S={S_t}  wrong shape {ot.shape}")
+
+    # ── Summary ──────────────────────────────────────────────────────────────
+    print(f"\n{HDR}{'═'*62}{RST}")
+    if failures == 0: print(f"  {PASS}  All tests passed.")
+    else:             print(f"  {FAIL}  {failures} test(s) failed.", file=sys.stderr)
+    print(f"{HDR}{'═'*62}{RST}\n")
+    sys.exit(failures)

veylon_model.py

@@ -0,0 +1,383 @@
+from __future__ import annotations
+
+import numpy as np
+import keras
+from keras import layers, ops
+import jax
+from veylon_attention import flash_splash_attention , decode_swa
+
+try:
+    from config import (
+        CONTEXT,
+        vocab_size as Vocab_size,
+        D_MODEL,
+        numberoflayers,
+        numberofheads,
+        d_Latent,
+        ffn_mult,
+        swa_window,
+        num_kv_heads,
+        use_moe,
+        moe_num_experts,
+        moe_top_k,
+    )
+except Exception:
+    CONTEXT = 2048
+    Vocab_size = 32000
+    D_MODEL = 512
+    numberoflayers = 8
+    numberofheads = 8
+    d_Latent = 128
+    ffn_mult = 3.5
+    swa_window = 1024
+    num_kv_heads = 2
+    use_moe = False
+    moe_num_experts = 8
+    moe_top_k = 2
+
+keras.mixed_precision.set_global_policy('mixed_bfloat16')
+
+
+@keras.saving.register_keras_serializable()
+class RMSNorm(layers.Layer):
+    def __init__(self, epsilon=1e-5, **kwargs):
+        super().__init__(**kwargs)
+        self.epsilon = epsilon
+
+    def build(self, input_shape):
+        self.weight = self.add_weight(shape=(input_shape[-1],), initializer='ones', name='gamma')
+
+    def call(self, x):
+        x_fp32 = ops.cast(x, 'float32')
+        rms = ops.sqrt(ops.mean(ops.square(x_fp32), axis=-1, keepdims=True) + self.epsilon)
+        out = x_fp32 / rms
+        return ops.cast(out, x.dtype) * self.weight
+
+    def get_config(self):
+        cfg = super().get_config()
+        cfg.update({'epsilon': self.epsilon})
+        return cfg
+
+
+@keras.saving.register_keras_serializable()
+class RotaryEmbedding(layers.Layer):
+    def __init__(self, max_seq_len, dim, theta=10000.0, **kwargs):
+        super().__init__(**kwargs)
+        if dim % 2 != 0:
+            raise ValueError('RotaryEmbedding dim must be even.')
+        self.max_seq_len = max_seq_len
+        self.dim = dim
+        self.theta = theta
+
+    def build(self, input_shape):
+        half = self.dim // 2
+        inv_freq = 1.0 / (self.theta ** (np.arange(0, self.dim, 2).astype(np.float32) / self.dim))
+        positions = np.arange(self.max_seq_len, dtype=np.float32)
+        freqs = positions[:, None] * inv_freq[None, :]
+        self.cos = self.add_weight(shape=(self.max_seq_len, half), initializer=keras.initializers.Constant(np.cos(freqs)), trainable=False, dtype='float32', name='cos_table')
+        self.sin = self.add_weight(shape=(self.max_seq_len, half), initializer=keras.initializers.Constant(np.sin(freqs)), trainable=False, dtype='float32', name='sin_table')
+
+    def call(self, x, offset=0):
+        seq_len = x.shape[1]
+        if seq_len is None:
+            seq_len = ops.shape(x)[1]
+        half = self.dim // 2
+        if offset + seq_len > self.max_seq_len:
+            raise ValueError(f'RoPE table too small: offset={offset}, seq_len={seq_len}, max={self.max_seq_len}')
+        cos = self.cos[offset:offset + seq_len, :half]
+        sin = self.sin[offset:offset + seq_len, :half]
+        cos = ops.cast(ops.reshape(cos, (1, seq_len, 1, half)), x.dtype)
+        sin = ops.cast(ops.reshape(sin, (1, seq_len, 1, half)), x.dtype)
+        x1 = x[..., :half]
+        x2 = x[..., half:]
+        return ops.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1)
+
+    def get_config(self):
+        cfg = super().get_config()
+        cfg.update({'max_seq_len': self.max_seq_len, 'dim': self.dim, 'theta': self.theta})
+        return cfg
+
+
+@keras.saving.register_keras_serializable()
+class SwiGLUFFN(layers.Layer):
+    def __init__(self, d_model, hidden_mult=3.5, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model_arg = d_model
+        self.hidden_mult = hidden_mult
+        self.hidden_dim = int(d_model * hidden_mult * 2 / 3)
+        self.hidden_dim = ((self.hidden_dim + 63) // 64) * 64
+
+    def build(self, input_shape):
+        d_model = input_shape[-1]
+        self.gate_up_proj = self.add_weight(shape=(d_model, 2 * self.hidden_dim), initializer='glorot_uniform', name='gate_up_proj')
+        self.down_proj = self.add_weight(shape=(self.hidden_dim, d_model), initializer='glorot_uniform', name='down_proj')
+
+    def call(self, x, training=False):
+        gate_up = ops.matmul(x, self.gate_up_proj)
+        gate, up = ops.split(gate_up, 2, axis=-1)
+        return ops.matmul(ops.silu(gate) * up, self.down_proj)
+
+    def get_config(self):
+        cfg = super().get_config()
+        cfg.update({'d_model': self.d_model_arg, 'hidden_mult': self.hidden_mult})
+        return cfg
+
+
+@keras.saving.register_keras_serializable()
+class MoE_FFN(layers.Layer):
+    def __init__(self, d_model, num_experts=8, top_k=2, hidden_mult=3.5, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model_arg = d_model
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.hidden_mult = hidden_mult
+        self.experts = [SwiGLUFFN(d_model, hidden_mult) for _ in range(num_experts)]
+        self.router = layers.Dense(num_experts, use_bias=False)
+
+    def _load_balancing_loss(self, router_logits, top_k_indices):
+        router_probs = ops.softmax(router_logits, axis=-1)
+        mask = ops.one_hot(top_k_indices, self.num_experts)
+        mask = ops.sum(mask, axis=2)
+        f = ops.mean(mask, axis=(0, 1))
+        p = ops.mean(router_probs, axis=(0, 1))
+        return self.num_experts * ops.sum(f * p)
+
+    def call(self, x, training=False):
+        router_logits = self.router(x)
+        top_logits, top_idx = ops.top_k(router_logits, self.top_k)
+        top_weights = ops.softmax(top_logits, axis=-1)
+        if training:
+            self.add_loss(self._load_balancing_loss(router_logits, top_idx))
+        out = ops.zeros_like(x)
+        for kk in range(self.top_k):
+            idx = top_idx[..., kk]
+            w = top_weights[..., kk]
+            for e in range(self.num_experts):
+                mask = ops.cast(idx == e, x.dtype)
+                expert_out = self.experts[e](x * mask[..., None], training=training)
+                out += expert_out * w[..., None] * mask[..., None]
+        return out
+
+    def get_config(self):
+        cfg = super().get_config()
+        cfg.update({'d_model': self.d_model_arg, 'num_experts': self.num_experts, 'top_k': self.top_k, 'hidden_mult': self.hidden_mult})
+        return cfg
+
+
+@keras.saving.register_keras_serializable()
+class MLAttention(layers.Layer):
+    def __init__(self, d_model, n_heads, d_latent, max_seq_len, num_kv_heads=2, swa_window=1024, attn_dropout=0.0, **kwargs):
+        super().__init__(**kwargs)
+        if d_model % n_heads != 0:
+            raise ValueError('d_model must be divisible by n_heads')
+        if n_heads % num_kv_heads != 0:
+            raise ValueError('n_heads must be divisible by num_kv_heads')
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.num_kv_heads = num_kv_heads
+        self.group_size = n_heads // num_kv_heads
+        self.d_head = d_model // n_heads
+        self.d_latent = d_latent
+        self.max_seq_len = max_seq_len
+        self.swa_window = swa_window
+        self.dropout = layers.Dropout(attn_dropout)
+
+    def build(self, input_shape):
+        self.W_qc = self.add_weight(shape=(self.d_model, self.d_model + self.d_latent), initializer='glorot_uniform', name='W_qc')
+        self.W_kv = self.add_weight(shape=(self.d_latent, self.num_kv_heads * 2 * self.d_head), initializer='glorot_uniform', name='W_kv')
+        self.W_o = self.add_weight(shape=(self.d_model, self.d_model), initializer='glorot_uniform', name='Wo')
+        self.rope = RotaryEmbedding(self.max_seq_len, self.d_head)
+
+    def _project_kv(self, c):
+        kv = ops.matmul(c, self.W_kv)
+        return ops.split(kv, 2, axis=-1)
+
+    def call(self, x, training=False):
+        B = ops.shape(x)[0]
+        S = ops.shape(x)[1]
+        qc = ops.matmul(x, self.W_qc)
+        q_proj, c = ops.split(qc, [self.d_model], axis=-1)
+        q = ops.reshape(q_proj, (B, S, self.n_heads, self.d_head))
+        q = self.rope(q, offset=0)
+        q = ops.transpose(q, (0, 2, 1, 3))
+        k, v = self._project_kv(c)
+        k = ops.reshape(k, (B, S, self.num_kv_heads, self.d_head))
+        v = ops.reshape(v, (B, S, self.num_kv_heads, self.d_head))
+        k = self.rope(k, offset=0)
+        k = ops.transpose(k, (0, 2, 1, 3))
+        v = ops.transpose(v, (0, 2, 1, 3))
+       
+
+        out = flash_splash_attention(
+         q,
+         k,
+         v,
+         window_size=min(self.swa_window, self.max_seq_len),
+         backend=jax.default_backend(),
+        use_gqa=True,
+        ) 
+
+        out = ops.transpose(out, (0, 2, 1, 3))
+        out = ops.reshape(out, (B, S, self.d_model))
+        out = self.dropout(out, training=training)
+        return ops.matmul(out, self.W_o)
+
+    def generate_step(self, x, cache_k=None, cache_v=None, cache_pos=0):
+     B = ops.shape(x)[0]
+     S = ops.shape(x)[1]
+
+     qc = ops.matmul(x, self.W_qc)
+     q_proj, c = ops.split(qc, [self.d_model], axis=-1)
+
+     q = ops.reshape(q_proj, (B, S, self.n_heads, self.d_head))
+     q = self.rope(q, offset=cache_pos)
+     q = ops.transpose(q, (0, 2, 1, 3))
+
+     k, v = self._project_kv(c)
+     k = ops.reshape(k, (B, S, self.num_kv_heads, self.d_head))
+     v = ops.reshape(v, (B, S, self.num_kv_heads, self.d_head))
+     k = self.rope(k, offset=cache_pos)
+     k = ops.transpose(k, (0, 2, 1, 3))
+     v = ops.transpose(v, (0, 2, 1, 3))
+
+    # Prefill path: no previous cache
+     if cache_k is None:
+        out = flash_splash_attention(
+            q,
+            k,
+            v,
+            window_size=min(self.swa_window, self.max_seq_len),
+            backend=jax.default_backend(),
+            use_gqa=True,
+        )
+        new_k = k[:, :, -self.swa_window :, :]
+        new_v = v[:, :, -self.swa_window :, :]
+     else:
+        # Decode path: one-token generation only
+        if S != 1:
+            raise ValueError(
+                f"generate_step with cache expects S=1, got S={S}"
+            )
+
+        k = ops.concatenate([cache_k, k], axis=2)
+        v = ops.concatenate([cache_v, v], axis=2)
+
+        k = k[:, :, -self.swa_window :, :]
+        v = v[:, :, -self.swa_window :, :]
+
+        out = decode_swa(q, k, v)
+        new_k = k
+        new_v = v
+
+     out = ops.transpose(out, (0, 2, 1, 3))
+     out = ops.reshape(out, (B, S, self.d_model))
+     out = ops.matmul(out, self.W_o)
+
+     return out, new_k, new_v
+
+    def get_config(self):
+        cfg = super().get_config()
+        cfg.update({'d_model': self.d_model, 'n_heads': self.n_heads, 'num_kv_heads': self.num_kv_heads, 'd_latent': self.d_latent, 'max_seq_len': self.max_seq_len, 'swa_window': self.swa_window, 'attn_dropout': self.dropout.rate})
+        return cfg
+
+
+@keras.saving.register_keras_serializable()
+class TransformerBlock(layers.Layer):
+    def __init__(self, d_model, n_heads, d_latent, ffn_layer, max_seq_len, num_kv_heads=2, swa_window=1024, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_latent = d_latent
+        self.max_seq_len = max_seq_len
+        self.num_kv_heads = num_kv_heads
+        self.swa_window = swa_window
+        self.ffn = keras.saving.deserialize_keras_object(ffn_layer) if isinstance(ffn_layer, dict) else ffn_layer
+        self.norm1 = RMSNorm()
+        self.norm2 = RMSNorm()
+        self.attn = MLAttention(d_model, n_heads, d_latent, max_seq_len, num_kv_heads=num_kv_heads, swa_window=swa_window)
+
+    def call(self, x, training=False):
+        x = x + self.attn(self.norm1(x), training=training)
+        x = x + self.ffn(self.norm2(x), training=training)
+        return x
+
+    def generate_step(self, x, cache_k=None, cache_v=None, cache_pos=0):
+     attn_out, nck, ncv = self.attn.generate_step(
+        self.norm1(x),
+        cache_k=cache_k,
+        cache_v=cache_v,
+        cache_pos=cache_pos,
+    )
+     x = x + attn_out
+     x = x + self.ffn(self.norm2(x), training=False)
+     return x, nck, ncv
+    def get_config(self):
+        cfg = super().get_config()
+        cfg.update({'d_model': self.d_model, 'n_heads': self.n_heads, 'd_latent': self.d_latent, 'ffn_layer': keras.saving.serialize_keras_object(self.ffn), 'max_seq_len': self.max_seq_len, 'num_kv_heads': self.num_kv_heads, 'swa_window': self.swa_window})
+        return cfg
+
+
+@keras.saving.register_keras_serializable()
+class VeylonModel(keras.Model):
+    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_latent, ffn_mult, max_seq_len, use_moe=False, moe_num_experts=8, moe_top_k=2, num_kv_heads=2, swa_window=1024, **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.d_latent = d_latent
+        self.ffn_mult = ffn_mult
+        self.max_seq_len = max_seq_len
+        self.use_moe = use_moe
+        self.moe_num_experts = moe_num_experts
+        self.moe_top_k = moe_top_k
+        self.num_kv_heads = num_kv_heads
+        self.swa_window = swa_window
+        self.embedding = layers.Embedding(vocab_size, d_model, name='token_embedding')
+        self.blocks = []
+        for i in range(n_layers):
+            ffn = MoE_FFN(d_model, moe_num_experts, moe_top_k, ffn_mult) if use_moe else SwiGLUFFN(d_model, ffn_mult)
+            self.blocks.append(TransformerBlock(d_model, n_heads, d_latent, ffn, max_seq_len, num_kv_heads=num_kv_heads, swa_window=swa_window, name=f'block_{i}'))
+        self.norm = RMSNorm()
+
+    def call(self, inputs, training=False):
+        x = self.embedding(inputs)
+        for block in self.blocks:
+            x = block(x, training=training)
+        x = self.norm(x)
+        embedding_weights = self.embedding.weights[0]
+        logits = ops.matmul(x, ops.transpose(embedding_weights))
+        return ops.cast(logits, 'float32')
+
+    def generate_step(self, inputs, cache_k=None, cache_v=None, cache_pos=0):
+     x = self.embedding(inputs)
+     new_cache_k = []
+     new_cache_v = []
+
+     if cache_k is None:
+        cache_k = [None] * len(self.blocks)
+        cache_v = [None] * len(self.blocks)
+
+     for i, block in enumerate(self.blocks):
+        x, nck, ncv = block.generate_step(
+            x,
+            cache_k=cache_k[i],
+            cache_v=cache_v[i],
+            cache_pos=cache_pos,
+        )
+        new_cache_k.append(nck)
+        new_cache_v.append(ncv)
+
+     x = self.norm(x)
+     embedding_weights = self.embedding.weights[0]
+     logits = ops.matmul(x, ops.transpose(embedding_weights))
+     logits = ops.cast(logits, 'float32')
+     return logits, new_cache_k, new_cache_v
+    def get_config(self):
+        cfg = super().get_config()
+        cfg.update({'vocab_size': self.vocab_size, 'd_model': self.d_model, 'n_layers': self.n_layers, 'n_heads': self.n_heads, 'd_latent': self.d_latent, 'ffn_mult': self.ffn_mult, 'max_seq_len': self.max_seq_len, 'use_moe': self.use_moe, 'moe_num_experts': self.moe_num_experts, 'moe_top_k': self.moe_top_k, 'num_kv_heads': self.num_kv_heads, 'swa_window': self.swa_window})
+        return cfg
+
+
+def create_llm(vocab_size=Vocab_size, d_model=D_MODEL, n_layers=numberoflayers, n_heads=numberofheads, d_latent=d_Latent, ffn_mult=ffn_mult, max_seq_len=CONTEXT, use_moe=use_moe, moe_num_experts=moe_num_experts, moe_top_k=moe_top_k, num_kv_heads=num_kv_heads, swa_window=swa_window):
+    return VeylonModel(vocab_size=vocab_size, d_model=d_model, n_layers=n_layers, n_heads=n_heads, d_latent=d_latent, ffn_mult=ffn_mult, max_seq_len=max_seq_len, use_moe=use_moe, moe_num_experts=moe_num_experts, moe_top_k=moe_top_k, num_kv_heads=num_kv_heads, swa_window=swa_window)