Upload folder using huggingface_hub

README.md

@@ -0,0 +1,205 @@
+---
+license: apache-2.0
+language:
+- en
+pipeline_tag: text-generation
+tags:
+- hybrid-ssm
+- mha
+- moe
+- state-space-models
+- recurrent-gated-expert
+- arg
+- hse
+- b300
+- runpod
+- instruct
+- chat
+base_model: Phase-Technologies/forge-500m-base
+---
+
+# FORGE-500M-Chat: First-Principles Optimal Recurrent-Gated Expert Architecture
+
+<div align="center">
+
+[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Parameters](https://img.shields.io/badge/Parameters-500M%20Total%20%7C%20320M%20Active-green.svg)]()
+[![Hardware](https://img.shields.io/badge/Trained%20On-4%C3%97%20NVIDIA%20B300%20SXM6-76B900.svg?logo=nvidia)]()
+[![Architecture](https://img.shields.io/badge/Architecture-Hybrid%20SSM%20%2B%20MHA%20%2B%20HSE--MoE-orange.svg)]()
+[![Training Tokens](https://img.shields.io/badge/Tokens-1.66B%20Pretrain%20%2B%20Aligned-purple.svg)]()
+
+</div>
+
+---
+
+## Executive Summary
+
+**FORGE-500M-Chat** is a state-of-the-art hybrid sub-billion parameter conversational language model designed from first principles by **Soham Pal**. It solves the compute accessibility gap in AI research by replacing quadratic self-attention with linear-time recurrent state-space recurrence, paired with hierarchical sparse expert mixture-of-experts (HSE-MoE).
+
+The model was pretrained on **1.66 Billion tokens** from scratch across 4× NVIDIA B300 SXM6 MIG compute slices, aligned via **Supervised Fine-Tuning (SFT)** on multi-turn dialogue with per-token loss masking, and polished using **Direct Preference Optimization (DPO)**.
+
+The full console training trajectory from Step 1 to Step 821 is archived in this repository under [`training_journey.log`](./training_journey.log).
+
+---
+
+## Model Architecture Specifications
+
+FORGE employs an asymmetric, multi-layer hybrid topology designed for maximum parameter utility and linear sequence scalability:
+
+| Specification | Value | Technical Description |
+| :--- | :--- | :--- |
+| **Total Parameters** | **500,019,616** | 500.0M total weights (FP16/BF16 storage: ~1.0 GB) |
+| **Active Parameters** | **~320,000,000** | Active per-token forward pass via Sparse HSE-MoE |
+| **Embedding Size** | 211,444,736 | 206,464 vocabulary × 1,024 hidden dimension |
+| **Hidden Dimension ($d_{\text{model}}$)** | 1,024 | Intermediate representation width |
+| **Total Transformer Layers** | 18 | Interleaved hybrid recurrence + global attention |
+| **SSM Recurrence Layers** | 15 (83.3%) | **Adaptive Recurrent Gating (ARG)** linear state space |
+| **Attention Layers** | 3 (16.7%) | **Global Multi-Head Attention (MHA)** with RoPE + YaRN |
+| **Attention Heads** | 8 | 8 Query heads, 8 Key/Value heads |
+| **MoE FFN Layers** | 9 | **Hierarchical Sparse Expert (HSE)** (32 experts, top-2 active) |
+| **Dense FFN Layers** | 9 | Dense SwiGLU FFN with $8/3 \times d_{\text{model}}$ expansion |
+| **Context Window** | 4,096 | Native context (scalable to 16,384 with YaRN) |
+| **Tokenizer** | CRAYON Turbo | 206,464 vocabulary (BPE with subword byte-fallback) |
+
+```
+Input Tokens (B, T)
+        │
+        ▼
+   [CRAYON Embedding (206k Vocab, 1024 d_model)]
+        │
+   ┌────┴───────────────────────────────────────┐
+   │ Layer 0..17: Interleaved Blocks           │
+   │  ├─ 15× ARG Layers: Linear SSM Recurrence │
+   │  ├─  3× Global MHA: Long-Range Attention  │
+   │  ├─  9× SwiGLU Dense FFN                  │
+   │  └─  9× Hierarchical Sparse Expert (MoE)  │
+   └────┬───────────────────────────────────────┘
+        ▼
+   [Final RMSNorm + Chunked LM Projection Head]
+        │
+        ▼
+   Target Logits / Probabilities
+```
+
+---
+
+## The Full Training Journey & Historical Narrative
+
+The pretraining and post-training of FORGE-500M represents an empirical milestone in lean, high-efficiency AI development:
+
+### 1. The Pretraining Stage
+- **Cluster**: 4× NVIDIA B300 SXM6 MIG slices (`1g.34gb`, 18 SMs each, CUDA 13.2, Driver 595.91).
+- **Global Batch**: 1,048,576 tokens per step (32 micro-batches × 32,768 tokens per batch).
+- **Physical Throughput**: 7,337 unique tokens per second (~26.4 Million tokens/hour).
+- **Step 1**: Initial Cross-Entropy Loss: **12.512**, Initial Perplexity: **> 150,000**.
+- **Step 821**: Final Loss: **4.6473**, Final Perplexity: **104.30**.
+- **Cumulative Pretrain Tokens**: **1,660,944,384 tokens** (1.66 Billion tokens).
+- **Challenges Overcome**:
+  - Survived multiple cloud host network disruptions and PCIe MIG interface resets via resilient checkpoint auto-recovery.
+  - Successfully resolved PyTorch 2.5 `nn.Module` dynamic attribute restrictions with custom parameter wrapping (`convert_to_attr_dict`).
+  - Implemented chunked linear projection (`chunk_size=1024`) on the 206,464-token vocabulary to eliminate VRAM allocation spikes.
+
+### 2. Supervised Fine-Tuning (SFT)
+- **Dataset**: Curated mixture of ShareGPT dialogues, UltraChat multi-turn instructions, CodeFeedback Python programming pairs, MetaMath mathematical reasoning, and OpenOrca synthetic data.
+- **Protocol**: Chat template formatting with assistant turn loss masking (no loss on system or user prompts).
+- **Optimization**: PyTorch DistributedDataParallel (DDP) across all 4 GPUs with cosine learning rate decay ($2\times 10^{-5} \to 2\times 10^{-6}$) and BF16 autocast.
+- **Result**: Loss dropped rapidly to **~4.0**, establishing clean turn-taking and instruction responsiveness.
+
+### 3. Direct Preference Optimization (DPO)
+- **Dataset**: Multi-domain preference pairs from UltraFeedback, Anthropic HH-RLHF, and HelpSteer2.
+- **Protocol**: Trainable policy model aligned against frozen BF16 SFT reference model with KL divergence regularization ($\beta = 0.1$).
+- **Result**: Achieved positive reward margin on chosen vs. rejected completions, eliminating repetitive degradation and improving refusal of hazardous requests.
+
+---
+
+## Complete Training Console Logs Preserved
+
+In accordance with our commitment to total scientific reproducibility, the complete console log file (`training_journey.log`) containing every training step from Step 1 to Step 821, memory footprints, step durations, and loss curves is provided directly in the repository root:
+
+- **Filename**: [`training_journey.log`](./training_journey.log)
+- **Size**: ~5.9 MB
+- **Coverage**: Steps 1 through 821, SFT, and DPO alignment logs.
+
+---
+
+## How to Run Inference
+
+### 1. Google Colab (1-Click Quickstart)
+
+In a Google Colab notebook (T4 / A100 GPU), paste and run this single cell:
+
+```bash
+# 1. Clone repository & enter directory
+!git clone https://huggingface.co/Phase-Technologies/FORGE-500M-Chat
+%cd FORGE-500M-Chat
+
+# 2. Install dependencies
+!pip install -q torch einops triton transformers huggingface_hub
+
+# 3. Generate text immediately
+!python run_inference.py --prompt "<|bos|><|usr|>Explain how quantum computers differ from classical computers.<|asst|>" --max_new_tokens 128
+```
+
+---
+
+### 2. Pure Python API (Colab / Jupyter / Local)
+
+```python
+import sys
+import torch
+from pathlib import Path
+from huggingface_hub import snapshot_download
+
+# 1. Download model repository
+repo_dir = Path(snapshot_download(repo_id="Phase-Technologies/FORGE-500M-Chat"))
+sys.path.insert(0, str(repo_dir))
+
+from tokenizer.crayon_wrapper import ForgeTokenizer
+from config import ForgeModelConfig
+from model.forge_model import build_forge_500m
+
+tokenizer = ForgeTokenizer.from_pretrained(str(repo_dir / "tokenizer"))
+config = ForgeModelConfig.from_json(str(repo_dir / "model_config.json"))
+model = build_forge_500m(config)
+
+# 3. Load BF16 weights
+weights = torch.load(repo_dir / "pytorch_model.bin", map_location="cpu")
+model.load_state_dict({k: v.to(torch.bfloat16) for k, v in weights.items()})
+model = model.to("cuda" if torch.cuda.is_available() else "cpu").to(torch.bfloat16)
+model.eval()
+
+# 4. Formatted Chat Prompt
+prompt = "<|bos|><|usr|> Explain how quantum computers differ from classical computers in simple terms.<|asst|>"
+
+input_ids = torch.tensor([tokenizer.encode(prompt, add_bos=False, add_eos=False)], device=model.device)
+
+# 5. Autoregressive Generation
+with torch.inference_mode():
+    output_ids = model.generate(
+        input_ids=input_ids,
+        max_new_tokens=128,
+        temperature=0.7,
+        top_p=0.9,
+        eos_token_id=tokenizer.special_tokens["<|eos|>"],
+    )
+
+print(tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=True))
+```
+
+### 3. Special Conversation Tokens
+
+| Token | String | Function |
+| :--- | :--- | :--- |
+| **BOS** | `<|bos|>` | Beginning of sequence |
+| **EOS** | `<|eos|>` | End of sequence / turn termination |
+| **SYS** | `<|sys|>` | System instruction block prefix |
+| **USR** | `<|usr|>` | User prompt prefix |
+| **ASST** | `<|asst|>` | Assistant completion prefix |
+| **PAD** | `<|pad|>` | Sequence padding token |
+
+---
+
+## a note from the creator.....
+
+honestly i really wanted to train this beauty for full 15b tokens..... that was the real plan from day one to see how far this custom architecture could push..... but running multiple top tier gpus burns money so fast and i just ran completely out of budget..... literally had to stop pretraining early at 821 steps..... then quickly did sft and dpo right on the pod before shutting it down..... huge shoutout and love to runpod.io for providing the insane compute and infrastructure that made training this beast possible..... every single cent went into this little 500m model and seeing it come alive from scratch with custom cuda kernels and our own tokenizer was pure magic..... maybe one day with better funding we will train the full 3b or 15b token run..... until then i hope you guys enjoy playing with this model as much as i loved building it..... thank you so much to everyone who supported this crazy journey.....
+

config.py

@@ -0,0 +1,371 @@
+"""
+FORGE-3B: Master Configuration
+All hyperparameters, paths, and GPU optimization settings in one place.
+"""
+
+from __future__ import annotations
+import os
+import json
+import math
+from dataclasses import dataclass, field, asdict
+from typing import Optional, List, Literal
+from pathlib import Path
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GPU GLOBAL OPTIMIZATION FLAGS
+# These are applied at import time across the entire process.
+# ─────────────────────────────────────────────────────────────────────────────
+
+import torch
+
+# Allow TF32 for matmuls and convolutions — massive speedup on Ampere/Hopper
+# with negligible precision loss for language model training.
+if torch.cuda.is_available():
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True          # auto-tune cuDNN kernels
+    try:
+        torch.backends.cuda.enable_flash_sdp(True)     # Flash SDP in PyTorch 2.x
+        torch.backends.cuda.enable_mem_efficient_sdp(True)
+        torch.backends.cuda.enable_math_sdp(False)     # disable slow math SDP
+    except Exception:
+        pass
+
+# NCCL tuning for 64x H100 GPU cluster (8x 8-GPU nodes via NVLink + 3.2 Tbps InfiniBand)
+os.environ.setdefault("NCCL_MIN_NCHANNELS", "32")
+os.environ.setdefault("NCCL_MAX_NCHANNELS", "64")
+os.environ.setdefault("NCCL_NSOCKS_PERTHREAD", "8")
+os.environ.setdefault("NCCL_SOCKET_NTHREADS", "8")
+os.environ.setdefault("NCCL_BUFFSIZE", "16777216")      # 16MB ring buffer for high-bandwidth AllReduce
+os.environ.setdefault("NCCL_NET_GDR_LEVEL", "5")        # GPUDirect RDMA Level 5 for NVLink + InfiniBand
+os.environ.setdefault("NCCL_CROSS_NIC", "1")            # Multi-NIC rail-optimized routing
+os.environ.setdefault("NCCL_IB_DISABLE", "0")           # InfiniBand enabled
+os.environ.setdefault("NCCL_P2P_DISABLE", "0")          # NVLink P2P enabled
+os.environ.setdefault("NCCL_NVLS_ENABLE", "1")          # NVLink SHARP collective offload
+os.environ.setdefault("CUDA_DEVICE_MAX_CONNECTIONS", "1") # Precise CUDA stream scheduling
+
+# CUDA memory allocator — reduce fragmentation for large model + optimizer states
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF",
+    "expandable_segments:True,max_split_size_mb:128,garbage_collection_threshold:0.8")
+
+# Tokenizer parallelism — disable HF tokenizer parallelism (we use CRAYON)
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# MODEL CONFIGURATION
+# ─────────────────────────────────────────────────────────────────────────────
+
+@dataclass
+class ForgeModelConfig:
+    """Complete FORGE-500M architecture configuration. 500.02M total params."""
+    
+    # ── Core Dimensions ──────────────────────────────────────────────────────
+    vocab_size: int = 206_464          # CRAYON standard profile (206,373) padded to ×128
+    d_model: int = 1024               # 500M: 1024 (down from 1280)
+    n_layers: int = 18                # 500M: 18 layers (15 ARG + 3 MHA)
+    max_seq_len: int = 4096
+    
+    # ── Layer Distribution ────────────────────────────────────────────────────
+    # Layer 0-4 = ARG; Layer 5 = MHA; repeat ×3 across 18 layers
+    mha_layer_indices: List[int] = field(default_factory=lambda: [5, 11, 17])
+    # Dense FFN on odd layers, HSE on even layers
+    dense_ffn_layer_indices: List[int] = field(default_factory=lambda: list(range(1, 18, 2)))
+    hse_ffn_layer_indices: List[int] = field(default_factory=lambda: list(range(0, 18, 2)))
+    
+    # ── ARG (Adaptive Recurrent Gating) ──────────────────────────────────────
+    arg_d_inner: int = 1024           # matches d_model for 500M
+    arg_d_state: int = 48
+    arg_d_rank: int = 48               # dt projection rank
+    arg_conv_kernel: int = 4
+    arg_local_window: int = 64         # local attention window (phase 1+2)
+    arg_local_n_heads: int = 8
+    arg_local_n_kv_heads: int = 2
+    arg_head_dim: int = 64            # 1024 / 16 heads
+    arg_scalar_gate: bool = True       # scalar gate per token (vs d_model-dim)
+    
+    # ── MHA (Multi-Head Attention) ────────────────────────────────────────────
+    mha_n_heads: int = 16
+    mha_n_kv_heads: int = 4            # GQA: 4:1 ratio
+    mha_head_dim: int = 64            # 1024 / 16 heads
+    rope_theta: float = 500_000.0      # long-context RoPE base
+    rope_scaling_type: Optional[str] = None  # None | "yarn"
+    rope_scaling_factor: float = 1.0
+    
+    # ── Dense FFN (SwiGLU) ────────────────────────────────────────────────────
+    dense_d_ff: int = 2560             # 2.5 × d_model, multiple of 128
+    
+    # ── HSE FFN (Hierarchical Sparse Expert) ──────────────────────────────────
+    hse_n_domains: int = 4
+    hse_n_experts_per_domain: int = 6  # total 24 experts (down from 32)
+    hse_top_k: int = 2                 # active per token
+    hse_d_ff_expert: int = 210        # scaled for 500M total params
+    hse_capacity_factor: float = 1.25
+    hse_expert_dropout: float = 0.1
+    hse_aux_loss_alpha: float = 0.01
+    hse_gumbel_tau_init: float = 1.0
+    hse_gumbel_tau_final: float = 0.1
+    hse_gumbel_tau_warmup_tokens: int = 10_000_000_000  # anneal over 10B tokens
+    
+    # ── Normalization ─────────────────────────────────────────────────────────
+    norm_type: str = "dgn"             # "dgn" or "rmsnorm"
+    dgn_n_groups: int = 16            # 1024 / 16 = 64 per group
+    norm_eps: float = 1e-6
+    
+    # ── Initialization ────────────────────────────────────────────────────────
+    init_std: float = 0.02
+    init_std_residual_scale: bool = True  # scale by 1/sqrt(2n_layers)
+    
+    # ── Misc ──────────────────────────────────────────────────────────────────
+    dropout: float = 0.0               # 0.0 during pretraining
+    tie_word_embeddings: bool = True
+    use_cache: bool = True             # KV cache for inference
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    
+    # ── Compile/Optimization Flags ────────────────────────────────────
+    use_flash_attention: bool = True
+    use_triton_kernels: bool = True    # Triton fused ops
+    use_torch_compile: bool = True
+    compile_mode: str = "max-autotune"  # max throughput on H100
+    use_gradient_checkpointing: bool = False  # 500M fits easily in 80GB without checkpointing
+    gradient_checkpointing_ratio: float = 0.0
+    
+    def n_params_total(self) -> int:
+        """Approximate total parameter count."""
+        embed = self.vocab_size * self.d_model
+        n_arg = len([i for i in range(self.n_layers) if i not in self.mha_layer_indices])
+        n_mha = len(self.mha_layer_indices)
+        n_dense = len(self.dense_ffn_layer_indices)
+        n_hse = len(self.hse_ffn_layer_indices)
+        
+        arg_per = (
+            self.d_model * 2 * self.arg_d_inner +      # in_proj
+            self.arg_d_inner * self.arg_conv_kernel +   # conv1d
+            self.arg_d_inner * (self.arg_d_rank + 2 * self.arg_d_state) +  # x_proj
+            self.arg_d_rank * self.arg_d_inner +        # dt_proj
+            self.arg_d_state * 2 +                      # nu, theta
+            self.arg_d_inner +                          # D
+            self.arg_d_inner * self.d_model +           # out_proj
+            self.arg_local_n_heads * self.arg_head_dim * self.d_model +  # Q
+            self.arg_local_n_kv_heads * self.arg_head_dim * self.d_model * 2 +  # KV
+            self.arg_local_n_heads * self.arg_head_dim * self.d_model +  # O
+            self.d_model * self.arg_d_state             # CPB
+        )
+        
+        mha_per = (
+            self.d_model * self.mha_n_heads * self.mha_head_dim +      # Q
+            self.d_model * self.mha_n_kv_heads * self.mha_head_dim * 2 +  # KV
+            self.d_model * self.mha_n_heads * self.mha_head_dim         # O
+        )
+        
+        dense_per = 3 * self.d_model * self.dense_d_ff
+        
+        n_experts_total = self.hse_n_domains * self.hse_n_experts_per_domain
+        hse_per = n_experts_total * 3 * self.d_model * self.hse_d_ff_expert
+        
+        total = embed + n_arg * arg_per + n_mha * mha_per + n_dense * dense_per + n_hse * hse_per
+        return total
+    
+    def to_json(self, path: str):
+        with open(path, "w") as f:
+            json.dump(asdict(self), f, indent=2)
+    
+    @classmethod
+    def from_json(cls, path: str) -> "ForgeModelConfig":
+        import inspect
+        with open(path) as f:
+            data = json.load(f)
+        sig = inspect.signature(cls)
+        filtered_data = {k: v for k, v in data.items() if k in sig.parameters}
+        return cls(**filtered_data)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# TRAINING CONFIGURATION
+# ─────────────────────────────────────────────────────────────────────────────
+
+@dataclass
+class PretrainConfig:
+    """Phase-aware pretraining configuration for FORGE-500M.
+    
+    Budget math (1x H100 SXM @ $3.29/hr):
+      - Phase 1 (vocab warmup):  2B tokens @ 80k tok/s  →  ~7h  → $23
+      - Phase 2 (core pretrain): 16B tokens @ 80k tok/s → ~56h  → $184
+      - Phase 3 (ctx extension):  2B tokens @ 60k tok/s →  ~9h  → $30
+      PRETRAIN TOTAL: ~72h → ~$237
+      SFT:  1.4B tokens → ~5h → $16
+      DPO:  0.8B tokens → ~3h → $10
+      GRAND TOTAL: ~$263 of $400 budget. $137 buffer.
+    """
+    
+    output_dir: str = "./checkpoints/forge_500m_pretrain"
+    data_dir: str = "Phase-Technologies/forge-3b-pretrain-data"  # same tokenized data
+    resume_from_checkpoint: Optional[str] = None
+    
+    # ── Token Budget (15B Tokens for 500M params: ~$245.59 total cost @ $3.50/GPU/hr on 32x H100) ──
+    # 15B tokens directly in Phase 2 at seq=2048 (Phase 1 & Phase 3 omitted)
+    phase1_tokens: int = 0                   # Omitted (start directly on Phase 2 @ seq=2048)
+    phase2_tokens: int = 15_000_000_000      # 15.0B core pretrain (seq=2048) -> ~2.08h ($233.33)
+    phase3_tokens: int = 0                   # Omitted (pure seq=2048 pretraining)
+    total_tokens: int = 15_000_000_000       # 15.0B total pretrain
+    
+    # ── Batch Config (Tuned for 32x H100 SXM — 1,048,576 tokens/step) ─────────
+    # With 32 GPUs @ micro_batch=16 (seq=2048): 32 × 16 × 2048 = 1,048,576 tokens per pass!
+    # Exactly 1 gradient accumulation step (zero idle wait, maximum NVLink saturation)
+    phase1_global_batch_tokens: int = 1_048_576   # 1M tokens/step
+    phase2_global_batch_tokens: int = 1_048_576   # 1M tokens/step (Phase 2)
+    phase3_global_batch_tokens: int = 1_048_576   # 1M tokens/step (Phase 3)
+    micro_batch_size_per_gpu: int = 16            # default fallback
+    phase1_micro_batch_per_gpu: int = 64          # 64 seqs × 512 = 32,768 tokens/GPU
+    phase2_micro_batch_per_gpu: int = 16          # 16 seqs × 2048 = 32,768 tokens/GPU (1 accum step across 32 GPUs)
+    phase3_micro_batch_per_gpu: int = 8           # 8 seqs × 4096 = 32,768 tokens/GPU
+    
+    # ── Context Lengths ───────────────────────────────────────────────────────
+    phase1_seq_len: int = 512
+    phase2_seq_len: int = 2048
+    phase3_seq_len: int = 2048
+    
+    # ── Learning Rate ─────────────────────────────────────────────────────────
+    lr_max: float = 3e-4
+    lr_min: float = 3e-5
+    lr_warmup_tokens: int = 500_000_000      # 500M tokens warmup directly in Phase 2
+    lr_schedule: str = "cosine"
+    
+    # ── AdamW ─────────────────────────────────────────────────────────────────
+    beta1: float = 0.9
+    beta2: float = 0.95
+    eps: float = 1e-8
+    weight_decay: float = 0.1
+    grad_clip: float = 1.0
+    
+    # ── Parameter Groups LR Multipliers ──────────────────────────────────────
+    embedding_lr_mult: float = 0.5
+    ssm_lr_mult: float = 0.3
+    router_lr_mult: float = 0.5
+    
+    # ── Checkpointing (Hyper-Resilient: saves every 20-30 mins OR 500M tokens) 
+    save_every_n_tokens: int = 500_000_000    # checkpoint every 500M tokens (~8 mins on 32x GPUs)
+    save_every_minutes: float = 20.0          # time-based checkpoint guard (every 20 minutes)
+    keep_last_n_checkpoints: int = 10         # keep 10 historical snapshots
+    
+    # ── Logging ───────────────────────────────────────────────────────────────
+    log_every_n_steps: int = 1           # log every step for full metrics visibility
+    eval_every_n_steps: int = 200        # eval perplexity every 200 steps
+    wandb_project: str = "forge_500m_pretrain"
+    wandb_entity: Optional[str] = None
+    
+    # ── GPU/Distributed ───────────���───────────────────────────────────────────
+    num_gpus: int = 1                          # default 1x H100; set to 16 for 16x
+    bf16: bool = True
+    fp8: bool = False                          # H100 FP8 (experimental)
+    deepspeed_config: str = "./configs/ds_zero3.json"
+    torch_compile: bool = True
+    compile_mode: str = "max-autotune"         # peak throughput
+    
+    # ── Data ──────────────────────────────────────────────────────────────────
+    num_dataloader_workers: int = 4
+    prefetch_factor: int = 4
+    seed: int = 42
+    
+    @property
+    def gradient_accumulation_steps_phase2(self) -> int:
+        """Accumulation steps for phase 2 global batch."""
+        tokens_per_gpu_step = self.micro_batch_size_per_gpu * self.phase2_seq_len
+        tokens_per_step_all_gpus = tokens_per_gpu_step * self.num_gpus
+        return self.phase2_global_batch_tokens // tokens_per_step_all_gpus
+
+
+@dataclass
+class SFTConfig:
+    """Supervised Fine-Tuning configuration for FORGE-500M."""
+    
+    base_model_path: str = "./checkpoints/forge_500m_pretrain/final"
+    output_dir: str = "./checkpoints/forge_500m_sft"
+    data_dir: str = "Phase-Technologies/forge-3b-sft-data"
+    
+    # ── Training ──────────────────────────────────────────────────────────────
+    total_tokens: int = 1_400_000_000
+    global_batch_tokens: int = 524_288  # 512K per step
+    micro_batch_size_per_gpu: int = 4
+    seq_len: int = 4096
+    
+    # ── Optimizer ─────────────────────────────────────────────────────────────
+    lr_max: float = 1e-5
+    lr_min: float = 1e-6
+    lr_schedule: str = "cosine"
+    beta1: float = 0.9
+    beta2: float = 0.95
+    eps: float = 1e-8
+    weight_decay: float = 0.0
+    grad_clip: float = 0.5
+    
+    epochs: int = 1
+    loss_on_prompt: bool = False       # only compute loss on assistant turns
+    
+    save_every_n_steps: int = 100
+    wandb_project: str = "forge_500m_sft"
+    deepspeed_config: str = "./configs/ds_zero3_sft.json"
+    num_gpus: int = 1
+    bf16: bool = True
+    torch_compile: bool = True
+    compile_mode: str = "max-autotune"
+
+
+@dataclass  
+class DPOConfig:
+    """Direct Preference Optimization configuration for FORGE-500M."""
+    
+    base_model_path: str = "./checkpoints/forge_500m_sft/final"
+    output_dir: str = "./checkpoints/forge_500m_dpo"
+    data_dir: str = "Phase-Technologies/forge-3b-dpo-data"
+    
+    # ── DPO ───────────────────────────────────────────────────────────────────
+    beta: float = 0.1                  # KL penalty coefficient
+    loss_type: str = "dpo"             # "dpo" | "ipo" | "cdpo"
+    reference_free: bool = False
+    
+    # ── Training ──────────────────────────────────────────────────────────────
+    n_preference_pairs: int = 200_000
+    batch_size_pairs: int = 8          # preference pairs per GPU step
+    gradient_accumulation_steps: int = 4
+    seq_len: int = 4096
+    
+    # ── Optimizer ─────────────────────────────────────────────────────────────
+    lr: float = 5e-7
+    beta1: float = 0.9
+    beta2: float = 0.99
+    eps: float = 1e-8
+    weight_decay: float = 0.0
+    grad_clip: float = 0.3
+    n_epochs: int = 1
+    
+    save_every_n_steps: int = 100
+    wandb_project: str = "forge_500m_dpo"
+    deepspeed_config: str = "./configs/ds_zero3_sft.json"
+    num_gpus: int = 1
+    bf16: bool = True
+    torch_compile: bool = True
+    compile_mode: str = "max-autotune"
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# TOKENIZER CONFIGURATION
+# ─────────────────────────────────────────────────────────────────────────────
+
+@dataclass
+class CrayonConfig:
+    """CRAYON tokenizer configuration."""
+    profile: str = "standard"          # "standard" (206k) | "lite" (50k)
+    device: str = "cpu"                # "cpu" is faster for tokenization (23M vs 1.2M tok/s)
+    pad_token: str = "<|pad|>"
+    bos_token: str = "<|bos|>"
+    eos_token: str = "<|eos|>"
+    sep_token: str = "<|sep|>"
+    sys_token: str = "<|sys|>"
+    usr_token: str = "<|usr|>"
+    asst_token: str = "<|asst|>"
+    tool_token: str = "<|tool|>"
+    tool_resp_token: str = "<|tool_resp|>"

model.safetensors

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

model/arg_layer.py

@@ -0,0 +1,291 @@
+"""
+ARG (Adaptive Recurrent Gating) Layer — FORGE's core sequence mixer.
+
+Combines:
+1. Complex-eigenvalue selective SSM (recurrent branch)  
+2. Local windowed GQA (local attention branch)
+3. Learned scalar gate α_t ∈ [0,1] that blends the two
+4. Compound Positional Bias (CPB) for SSM initialization
+"""
+
+from __future__ import annotations
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from .dgn_norm import build_norm
+from .rotary_embedding import RotaryEmbedding
+from .triton_kernels import fused_swiglu
+from .triton_ssm import fused_triton_ssm
+
+
+
+class ARGLayer(nn.Module):
+    """
+    Adaptive Recurrent Gating Layer.
+    
+    For each token t:
+        h_recur_t = ComplexSSM(x_{0:t})
+        h_local_t = WindowedGQA(x_{t-W:t})
+        α_t       = σ(w_g · x_t)
+        output_t  = α_t * h_local_t + (1 - α_t) * h_recur_t
+    """
+    
+    def __init__(
+        self,
+        d_model: int = 1280,
+        d_inner: int = 1280,
+        d_state: int = 48,
+        d_rank: int = 48,
+        conv_kernel: int = 4,
+        local_window: int = 64,
+        local_n_heads: int = 8,
+        local_n_kv_heads: int = 2,
+        head_dim: int = 80,
+        norm_type: str = "dgn",
+        dgn_n_groups: int = 16,
+        norm_eps: float = 1e-6,
+        use_flash_attention: bool = True,
+        rope_base: float = 500_000.0,
+        layer_idx: int = 0,
+        max_seq_len: int = 4096,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.d_inner = d_inner
+        self.d_state = d_state
+        self.d_rank = d_rank
+        self.conv_kernel = conv_kernel
+        self.local_window = local_window
+        self.local_n_heads = local_n_heads
+        self.local_n_kv_heads = local_n_kv_heads
+        self.head_dim = head_dim
+        self.use_flash_attention = use_flash_attention
+        self.layer_idx = layer_idx
+        
+        # ── Pre-norm ──────────────────────────────────────────────────────────
+        self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)
+        
+        # ── Recurrent Branch ──────────────────────────────────────────────────
+        # Input projection: x → (x_inner, z) both d_inner
+        self.in_proj = nn.Linear(d_model, 2 * d_inner, bias=False)
+        
+        # Depthwise causal conv for local mixing
+        self.conv1d = nn.Conv1d(
+            d_inner, d_inner,
+            kernel_size=conv_kernel,
+            padding=conv_kernel - 1,   # causal padding
+            groups=d_inner,
+            bias=True,
+        )
+        
+        # Input-dependent SSM parameters (low-rank for Δ)
+        self.x_proj = nn.Linear(d_inner, d_rank + 2 * d_state, bias=False)
+        self.dt_proj = nn.Linear(d_rank, d_inner, bias=True)
+        
+        # Complex SSM eigenvalue parameters
+        # Uses real-pair encoding: each complex eigenvalue λ = -exp(ν) ± jθ
+        # is represented as two real state channels, so d_state must be even.
+        # theta (oscillation) is implicitly learned through B/C projections.
+        assert d_state % 2 == 0, "d_state must be even for complex-pair encoding"
+        self.d_state_complex = d_state // 2
+
+        self.nu = nn.Parameter(torch.zeros(self.d_state_complex))    # decay (log scale)
+        
+        # Skip connection coefficient D
+        self.D = nn.Parameter(torch.ones(d_inner))
+        
+        # Output projection
+        self.out_proj = nn.Linear(d_inner, d_model, bias=False)
+        
+        # ── Local Attention Branch ────────────────────────────────────────────
+        q_dim = local_n_heads * head_dim
+        kv_dim = local_n_kv_heads * head_dim
+        
+        self.local_q = nn.Linear(d_model, q_dim, bias=False)
+        self.local_k = nn.Linear(d_model, kv_dim, bias=False)
+        self.local_v = nn.Linear(d_model, kv_dim, bias=False)
+        self.local_o = nn.Linear(q_dim, d_model, bias=False)
+        
+        # Local RoPE (initialize with safe maximum length of 8192 to avoid dynamic cache expansion during forward pass)
+        self.local_rope = RotaryEmbedding(
+            head_dim=head_dim,
+            max_seq_len=max_seq_len * 2,
+            base=rope_base,
+        )
+        
+        # ── Adaptive Gate ─────────────────────────────────────────────────────
+        self.gate_proj = nn.Linear(d_model, 1, bias=True)
+        nn.init.zeros_(self.gate_proj.weight)
+        nn.init.zeros_(self.gate_proj.bias)  # α=0.5 at init: sigmoid(0) = 0.5
+        
+        # Initialize weights
+        self._init_weights()
+    
+    def _init_weights(self):
+        """Careful initialization for stable training."""
+        # dt_proj bias: positive to encourage non-trivial step sizes at init
+        dt_init_std = self.d_rank ** -0.5
+        nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
+        dt_bias = torch.exp(
+            torch.rand(self.d_inner) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
+        )
+        with torch.no_grad():
+            self.dt_proj.bias.copy_(dt_bias + torch.log(-torch.expm1(-dt_bias)))
+        
+        # nu init — zero → moderate decay at start
+        nn.init.zeros_(self.nu)
+        
+        # Conv1d — identity-like init
+        nn.init.zeros_(self.conv1d.weight)
+        for i in range(self.d_inner):
+            self.conv1d.weight.data[i, 0, -1] = 1.0
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # COMPLEX EIGENVALUE → REAL A MATRIX
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def _build_real_A(self) -> torch.Tensor:
+        """
+        Build the (d_inner, d_state) real-valued A matrix for selective_scan_fn.
+
+        Each complex eigenvalue λ_i = -exp(ν_i) + jθ_i is encoded as two real
+        state channels with the same decay rate. The theta (oscillation) enters
+        via the B/C projections. mamba_ssm's selective_scan_fn discretizes A
+        internally via exp(A * Δ), so we pass the continuous-time log-decay
+        (negative real, per state channel).
+        """
+        decay = -F.softplus(self.nu.float())   # fp32 — mamba_ssm requires A in fp32
+        A_diag = decay.repeat_interleave(2)   # (d_state,)
+        A = A_diag.unsqueeze(0).expand(self.d_inner, -1).contiguous().float()  # (d_inner, d_state), fp32
+        return A
+    
+    # ─────────────────────────────────────────────────────────────────────────
+    # RECURRENT BRANCH
+    # ─────────────────────────────────────────────────────────────────────────
+    
+    @torch._dynamo.disable
+    def recurrent_branch(
+        self,
+        x: torch.Tensor,         # (B, T, d_model)
+        position_offset: int = 0,
+    ) -> torch.Tensor:
+        """
+        Selective SSM recurrent branch using mamba_ssm's fused CUDA kernel.
+        Returns (B, T, d_model).
+
+        Uses selective_scan_fn (same kernel as Mamba) which:
+        - Processes all channels in parallel with an efficient associative scan
+        - Applies delta_softplus internally (no manual softplus on dt)
+        - Fuses the z-gate SiLU multiplication inside the kernel
+        - Eliminates 512× Python-level kernel launches per layer
+        """
+        B, T, _ = x.shape
+
+        # ── Project inputs ────────────────────────────────────────────────────
+        xz = self.in_proj(x)                        # (B, T, 2*d_inner)
+        x_inner, z = xz.chunk(2, dim=-1)             # each (B, T, d_inner)
+
+        # ── Causal conv for local mixing ──────────────────────────────────────
+        x_conv = rearrange(x_inner, 'b t d -> b d t')
+        x_conv = self.conv1d(x_conv)[:, :, :T]       # trim causal padding, still (B, d_inner, T)
+        x_conv = F.silu(x_conv)                      # (B, d_inner, T)
+        x_conv_t = rearrange(x_conv, 'b d t -> b t d')
+
+        # ── Input-dependent SSM parameters ────────────────────────────────────
+        x_dbl = self.x_proj(x_conv_t)                # (B, T, d_rank + 2*d_state)
+        dt_raw, B_ssm, C_ssm = x_dbl.split([self.d_rank, self.d_state, self.d_state], dim=-1)
+
+        # dt: per-channel step size — (B, T, d_inner) → (B, d_inner, T)
+        dt = F.softplus(self.dt_proj(dt_raw))
+        dt = rearrange(dt, 'b t d -> b d t').contiguous()
+
+        # B, C: (B, d_state, T)
+        B_ssm = rearrange(B_ssm, 'b t n -> b n t').contiguous()
+        C_ssm = rearrange(C_ssm, 'b t n -> b n t').contiguous()
+
+        # ── Build A from complex eigenvalue params ────────────────────────────
+        A = self._build_real_A()                     # (d_inner, d_state), negative real
+
+        # ── Pure Triton Fused Selective SSM Scan ──────────────────────────────
+        z_t = rearrange(z, 'b t d -> b d t').contiguous()
+        y = fused_triton_ssm(x_conv, dt, A, B_ssm, C_ssm, self.D, z_t)
+        y = rearrange(y, 'b d t -> b t d')
+
+        return self.out_proj(y.to(dtype=x.dtype))    # (B, T, d_model)
+    
+    # ─────────────────────────────────────────────────────────────────────────
+    # LOCAL ATTENTION BRANCH
+    # ─────────────────────────────────────────────────────────────────────────
+    
+    @torch._dynamo.disable
+    def local_attention_branch(
+        self,
+        x: torch.Tensor,   # (B, T, d_model)
+    ) -> torch.Tensor:
+        """
+        Windowed GQA with window size W.
+        Uses native PyTorch SDPA with sliding window mask.
+        Returns (B, T, d_model).
+        """
+        B, T, _ = x.shape
+        W = self.local_window
+        
+        Q = self.local_q(x)  # (B, T, n_heads * head_dim)
+        K = self.local_k(x)  # (B, T, n_kv_heads * head_dim)
+        V = self.local_v(x)
+        
+        Q = rearrange(Q, 'b t (h d) -> b h t d', d=self.head_dim)  # (B, H, T, D)
+        K = rearrange(K, 'b t (h d) -> b h t d', d=self.head_dim)
+        V = rearrange(V, 'b t (h d) -> b h t d', d=self.head_dim)
+        
+        # Apply local RoPE
+        Q, K = self.local_rope.apply_rotary(Q, K)
+        
+        # GQA: expand KV to match Q heads
+        n_rep = self.local_n_heads // self.local_n_kv_heads
+        if n_rep > 1:
+            K = K.repeat_interleave(n_rep, dim=1)  # (B, H, T, D)
+            V = V.repeat_interleave(n_rep, dim=1)
+        
+        # Fast PyTorch SDPA with sliding window mask
+        indices = torch.arange(T, device=x.device)
+        mask_2d = (indices.unsqueeze(0) <= indices.unsqueeze(1)) & \
+                  (indices.unsqueeze(0) >= (indices.unsqueeze(1) - W))
+        mask_4d = mask_2d.unsqueeze(0).unsqueeze(0)
+        attn_out = F.scaled_dot_product_attention(Q, K, V, attn_mask=mask_4d)
+        out = rearrange(attn_out, 'b h t d -> b t (h d)')
+        
+        return self.local_o(out)  # (B, T, d_model)
+    
+    # ─────────────────────────────────────────────────────────────────────────
+    # FORWARD PASS
+    # ─────────────────────────────────────────────────────────────────────────
+    
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_offset: int = 0,
+    ) -> torch.Tensor:
+        """
+        x: (B, T, d_model)
+        Returns: (B, T, d_model) — residual connection applied inside
+        """
+        residual = x
+        x_norm = self.norm(x)
+        
+        # Two branches
+        h_recur = self.recurrent_branch(x_norm, position_offset)
+        h_local = self.local_attention_branch(x_norm)
+        
+        # Learned gate
+        alpha = torch.sigmoid(self.gate_proj(x_norm))  # (B, T, 1)
+        
+        # Blend and add residual
+        mixed = alpha * h_local + (1.0 - alpha) * h_recur  # (B, T, d_model)
+        
+        return residual + mixed

model/dgn_norm.py

@@ -0,0 +1,65 @@
+"""Differential Group Normalization layer."""
+
+import torch
+import torch.nn as nn
+from .triton_kernels import fused_dgn, TRITON_AVAILABLE
+
+
+class DifferentialGroupNorm(nn.Module):
+    """
+    Differential Group Normalization (DGN).
+    
+    Divides d_model features into G groups and applies independent
+    RMS normalization with learned scale+bias to each group.
+    
+    Parameters: 2 * d_model (weight + bias) vs RMSNorm's 1 * d_model.
+    The extra bias (absent in RMSNorm) provides an additive degree of freedom
+    per group, shown empirically to speed up early training convergence by ~15%.
+    """
+    
+    def __init__(self, d_model: int, n_groups: int = 16, eps: float = 1e-6):
+        super().__init__()
+        assert d_model % n_groups == 0, \
+            f"d_model={d_model} must be divisible by n_groups={n_groups}"
+        
+        self.d_model = d_model
+        self.n_groups = n_groups
+        self.group_size = d_model // n_groups
+        self.eps = eps
+        
+        # Learnable parameters — initialized to identity transform
+        self.weight = nn.Parameter(torch.ones(d_model))
+        self.bias = nn.Parameter(torch.zeros(d_model))
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (..., d_model)
+        Returns: (..., d_model), normalized within each group.
+        """
+        return fused_dgn(x, self.weight, self.bias, self.n_groups, self.eps)
+    
+    def extra_repr(self) -> str:
+        return f"d_model={self.d_model}, n_groups={self.n_groups}, eps={self.eps}"
+
+
+class RMSNorm(nn.Module):
+    """Standard RMSNorm — used as fallback or when DGN is disabled."""
+    
+    def __init__(self, d_model: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d_model))
+        self.eps = eps
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
+        return (x.float() / norm * self.weight).to(x.dtype)
+
+
+def build_norm(norm_type: str, d_model: int, n_groups: int = 16, 
+               eps: float = 1e-6) -> nn.Module:
+    if norm_type == "dgn":
+        return DifferentialGroupNorm(d_model, n_groups, eps)
+    elif norm_type == "rmsnorm":
+        return RMSNorm(d_model, eps)
+    else:
+        raise ValueError(f"Unknown norm_type: {norm_type}")

model/forge_model.py

@@ -0,0 +1,558 @@
+"""
+FORGE-3B: Complete Model Architecture
+Integrates ARG, MHA, Dense FFN, HSE FFN into the full autoregressive LM.
+"""
+
+from __future__ import annotations
+import math
+import logging
+from typing import Optional, List, Tuple, Dict, Any, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+from .arg_layer import ARGLayer
+from .mha_layer import GlobalMHALayer
+from .hse_layer import HSELayer, DenseSwiGLUFFN
+from .dgn_norm import build_norm
+
+logger = logging.getLogger(__name__)
+
+
+class ForgeBlock(nn.Module):
+    """
+    One FORGE block: SeqMixer + FFN (both with pre-norm and residual).
+    
+    Blocks are composed sequentially. The seq_mixer is either ARG or MHA;
+    the ffn is either DenseSwiGLUFFN or HSELayer — determined by layer_idx.
+    """
+    
+    def __init__(
+        self,
+        d_model: int,
+        layer_idx: int,
+        is_mha: bool,
+        is_hse: bool,
+        # ARG config
+        arg_d_inner: int = 1280,
+        arg_d_state: int = 48,
+        arg_d_rank: int = 48,
+        arg_conv_kernel: int = 4,
+        arg_local_window: int = 64,
+        arg_local_n_heads: int = 8,
+        arg_local_n_kv_heads: int = 2,
+        arg_head_dim: int = 80,
+        # MHA config
+        mha_n_heads: int = 16,
+        mha_n_kv_heads: int = 4,
+        mha_head_dim: int = 80,
+        rope_base: float = 500_000.0,
+        rope_scaling_type: Optional[str] = None,
+        rope_scaling_factor: float = 1.0,
+        max_seq_len: int = 4096,
+        # Dense FFN config
+        dense_d_ff: int = 3200,
+        # HSE config
+        hse_n_domains: int = 4,
+        hse_n_experts_per_domain: int = 8,
+        hse_top_k: int = 2,
+        hse_d_ff_expert: int = 288,
+        hse_capacity_factor: float = 1.25,
+        hse_expert_dropout: float = 0.1,
+        hse_aux_loss_alpha: float = 0.01,
+        # Norm
+        norm_type: str = "dgn",
+        dgn_n_groups: int = 16,
+        norm_eps: float = 1e-6,
+        use_flash_attention: bool = True,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.is_mha = is_mha
+        self.is_hse = is_hse
+        
+        # ── Sequence Mixer ────────────────────────────────────────────────────
+        if is_mha:
+            self.seq_mixer = GlobalMHALayer(
+                d_model=d_model,
+                n_heads=mha_n_heads,
+                n_kv_heads=mha_n_kv_heads,
+                head_dim=mha_head_dim,
+                rope_base=rope_base,
+                rope_scaling_type=rope_scaling_type,
+                rope_scaling_factor=rope_scaling_factor,
+                max_seq_len=max_seq_len,
+                dropout=dropout,
+                norm_type=norm_type,
+                dgn_n_groups=dgn_n_groups,
+                norm_eps=norm_eps,
+                use_flash_attention=use_flash_attention,
+                layer_idx=layer_idx,
+            )
+        else:
+            self.seq_mixer = ARGLayer(
+                d_model=d_model,
+                d_inner=arg_d_inner,
+                d_state=arg_d_state,
+                d_rank=arg_d_rank,
+                conv_kernel=arg_conv_kernel,
+                local_window=arg_local_window,
+                local_n_heads=arg_local_n_heads,
+                local_n_kv_heads=arg_local_n_kv_heads,
+                head_dim=arg_head_dim,
+                norm_type=norm_type,
+                dgn_n_groups=dgn_n_groups,
+                norm_eps=norm_eps,
+                use_flash_attention=use_flash_attention,
+                rope_base=rope_base,
+                layer_idx=layer_idx,
+                max_seq_len=max_seq_len,
+            )
+        
+        # ── FFN ───────────────────────────────────────────────────────────────
+        if is_hse:
+            self.ffn = HSELayer(
+                d_model=d_model,
+                n_domains=hse_n_domains,
+                n_experts_per_domain=hse_n_experts_per_domain,
+                top_k=hse_top_k,
+                d_ff_expert=hse_d_ff_expert,
+                capacity_factor=hse_capacity_factor,
+                expert_dropout=hse_expert_dropout,
+                aux_loss_alpha=hse_aux_loss_alpha,
+                norm_type=norm_type,
+                dgn_n_groups=dgn_n_groups,
+                norm_eps=norm_eps,
+                layer_idx=layer_idx,
+            )
+        else:
+            self.ffn = DenseSwiGLUFFN(
+                d_model=d_model,
+                d_ff=dense_d_ff,
+                norm_type=norm_type,
+                dgn_n_groups=dgn_n_groups,
+                norm_eps=norm_eps,
+            )
+    
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        position_offset: int = 0,
+        use_cache: bool = False,
+        return_aux_loss: bool = True,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        
+        aux_loss = None
+        
+        # ── Sequence Mixing ───────────────────────────────────────────────────
+        if self.is_mha:
+            x, _ = self.seq_mixer(
+                x, attention_mask=attention_mask,
+                position_ids=position_ids, use_cache=use_cache
+            )
+        else:
+            x = self.seq_mixer(x, position_offset=position_offset)
+        
+        # ── FFN ───────────────────────────────────────────────────────────────
+        if self.is_hse:
+            x, aux_loss = self.ffn(x, return_aux_loss=return_aux_loss)
+        else:
+            x = self.ffn(x)
+        
+        return x, aux_loss
+
+
+class ForgeModel(nn.Module):
+    """
+    FORGE-3B: Full autoregressive language model.
+    
+    Architecture:
+    - Token embedding (tied to LM head)
+    - 36 ForgeBlocks (27 ARG + 9 MHA for seq mixing; 18 Dense + 18 HSE for FFN)
+    - Final DGN normalization
+    - LM head (weight-tied to embedding)
+    """
+    
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        d = config.d_model
+        
+        # ── Token Embedding ───────────────────────────────────────────────────
+        self.embed_tokens = nn.Embedding(config.vocab_size, d, padding_idx=0)
+        
+        # ── FORGE Blocks ──────────────────────────────────────────────────────
+        mha_set = set(config.mha_layer_indices)
+        hse_set = set(config.hse_ffn_layer_indices)
+        
+        self.layers = nn.ModuleList([
+            ForgeBlock(
+                d_model=d,
+                layer_idx=i,
+                is_mha=(i in mha_set),
+                is_hse=(i in hse_set),
+                arg_d_inner=config.arg_d_inner,
+                arg_d_state=config.arg_d_state,
+                arg_d_rank=config.arg_d_rank,
+                arg_conv_kernel=config.arg_conv_kernel,
+                arg_local_window=config.arg_local_window,
+                arg_local_n_heads=config.arg_local_n_heads,
+                arg_local_n_kv_heads=config.arg_local_n_kv_heads,
+                arg_head_dim=config.arg_head_dim,
+                mha_n_heads=config.mha_n_heads,
+                mha_n_kv_heads=config.mha_n_kv_heads,
+                mha_head_dim=config.mha_head_dim,
+                rope_base=config.rope_theta,
+                rope_scaling_type=config.rope_scaling_type,
+                rope_scaling_factor=config.rope_scaling_factor,
+                max_seq_len=config.max_seq_len,
+                dense_d_ff=config.dense_d_ff,
+                hse_n_domains=config.hse_n_domains,
+                hse_n_experts_per_domain=config.hse_n_experts_per_domain,
+                hse_top_k=config.hse_top_k,
+                hse_d_ff_expert=config.hse_d_ff_expert,
+                hse_capacity_factor=config.hse_capacity_factor,
+                hse_expert_dropout=config.hse_expert_dropout,
+                hse_aux_loss_alpha=config.hse_aux_loss_alpha,
+                norm_type=config.norm_type,
+                dgn_n_groups=config.dgn_n_groups,
+                norm_eps=config.norm_eps,
+                use_flash_attention=config.use_flash_attention,
+            )
+            for i in range(config.n_layers)
+        ])
+        
+        # ── Final Norm ────────────────────────────────────────────────────────
+        self.norm_final = build_norm(config.norm_type, d, config.dgn_n_groups, config.norm_eps)
+        
+        # ── LM Head ───────────────────────────────────────────────────────────
+        self.lm_head = nn.Linear(d, config.vocab_size, bias=False)
+        
+        # Tie weights
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+        
+        # ── Gradient Checkpointing Configuration ─────────────────────────────
+        self.gradient_checkpointing_enabled = config.use_gradient_checkpointing
+        self.gradient_checkpointing_ratio = config.gradient_checkpointing_ratio
+        # Checkpoint every N-th layer
+        if config.gradient_checkpointing_ratio > 0:
+            self._checkpoint_interval = max(1, int(1.0 / config.gradient_checkpointing_ratio))
+        else:
+            self._checkpoint_interval = 1
+        
+        # Apply weight initialization
+        self._init_weights()
+        
+        # Log parameter count
+        n_params = sum(p.ds_numel if hasattr(p, "ds_numel") else p.numel() for p in self.parameters())
+        logger.info(f"ForgeModel: {n_params/1e9:.3f}B total parameters")
+        emb_p = self.embed_tokens.weight
+        emb_numel = emb_p.ds_numel if hasattr(emb_p, "ds_numel") else emb_p.numel()
+        logger.info(f"  Embedding: {emb_numel/1e6:.1f}M")
+        logger.info(f"  N layers: {config.n_layers} "
+                    f"({len(mha_set)} MHA + {config.n_layers - len(mha_set)} ARG, "
+                    f"{len(hse_set)} HSE + {config.n_layers - len(hse_set)} Dense)")
+    
+    def _init_weights(self):
+        """Initialize all model weights with scaled normal distribution."""
+        n_layers = self.config.n_layers
+        
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, std=self.config.init_std)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, std=self.config.init_std)
+                if module.padding_idx is not None:
+                    module.weight.data[module.padding_idx].zero_()
+        
+        # Residual scaling: scale output projections by 1/sqrt(2*n_layers)
+        # This prevents the residual stream from growing with depth
+        if self.config.init_std_residual_scale:
+            scale = 1.0 / math.sqrt(2 * n_layers)
+            for layer in self.layers:
+                if isinstance(layer.seq_mixer, ARGLayer):
+                    nn.init.normal_(layer.seq_mixer.out_proj.weight, std=scale)
+                    nn.init.normal_(layer.seq_mixer.local_o.weight, std=scale)
+                elif isinstance(layer.seq_mixer, GlobalMHALayer):
+                    nn.init.normal_(layer.seq_mixer.o_proj.weight, std=scale)
+                
+                if isinstance(layer.ffn, DenseSwiGLUFFN):
+                    nn.init.normal_(layer.ffn.down.weight, std=scale)
+                elif isinstance(layer.ffn, HSELayer):
+                    for expert in layer.ffn.experts:
+                        nn.init.normal_(expert.down.weight, std=scale)
+    
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+    
+    def set_input_embeddings(self, value: nn.Embedding):
+        self.embed_tokens = value
+    
+    def enable_gradient_checkpointing(self):
+        self.gradient_checkpointing_enabled = True
+        logger.info("Gradient checkpointing enabled")
+    
+    def disable_gradient_checkpointing(self):
+        self.gradient_checkpointing_enabled = False
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,                         # (B, T)
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,           # (B, T) for LM loss
+        use_cache: bool = False,
+        return_aux_loss: bool = True,
+        output_hidden_states: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Forward pass for FORGE-3B.
+        
+        Returns dict with:
+            'loss': cross-entropy + aux_loss (if labels provided)
+            'logits': (B, T, vocab_size)
+            'aux_loss': sum of MoE load-balance losses
+            'hidden_states': list of layer outputs (if output_hidden_states)
+        """
+        B, T = input_ids.shape
+        device = input_ids.device
+        
+        # ── Token Embeddings ──────────────────────────────────────────────────
+        h = self.embed_tokens(input_ids)  # (B, T, d_model)
+        
+        if position_ids is None:
+            position_ids = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
+        
+        # ── Forward Through Layers ────────────────────────────────────────────
+        all_hidden_states = [] if output_hidden_states else None
+        total_aux_loss = torch.zeros(1, device=device, dtype=h.dtype)
+        n_hse_layers = 0
+        
+        for layer_idx, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states.append(h)
+            
+            # Selective gradient checkpointing
+            if (self.gradient_checkpointing_enabled and 
+                self.training and 
+                layer_idx % self._checkpoint_interval == 0):
+                
+                def _forward_layer(h_, layer_=layer):
+                    out, aux = layer_(
+                        h_,
+                        attention_mask=attention_mask,
+                        position_ids=position_ids,
+                        position_offset=0,
+                        use_cache=False,
+                        return_aux_loss=True,
+                    )
+                    return out, aux if aux is not None else torch.zeros(1, device=h_.device)
+                
+                h, aux = checkpoint(
+                    _forward_layer, h,
+                    use_reentrant=False,  # required for PyTorch DDP to prevent double-marking ready
+                    preserve_rng_state=True,
+                )
+            else:
+                h, aux = layer(
+                    h,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    position_offset=0,
+                    use_cache=use_cache,
+                    return_aux_loss=return_aux_loss,
+                )
+            
+            if aux is not None and return_aux_loss:
+                total_aux_loss = total_aux_loss + aux
+                n_hse_layers += 1
+        
+        # ── Final Norm ────────────────────────────────────────────────────────
+        h = self.norm_final(h)  # (B, T, d_model)
+        
+        # ── LM Head & Loss Computation ────────────────────────────────────────
+        loss = None
+        if labels is not None:
+            # Shift for next-token prediction
+            shift_h = h[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous().reshape(-1)
+            
+            # Chunked LM head projection to prevent allocating massive (B, T, 206464) FP32 tensor in VRAM
+            shift_h_flat = shift_h.reshape(-1, shift_h.shape[-1])
+            chunk_size = 1024
+            n_tokens = shift_h_flat.shape[0]
+            total_ce_loss = 0.0
+            valid_tokens = 0
+            
+            for i in range(0, n_tokens, chunk_size):
+                h_chunk = shift_h_flat[i : i + chunk_size]
+                labels_chunk = shift_labels[i : i + chunk_size]
+                
+                # Check for valid tokens
+                mask = (labels_chunk != -100)
+                n_valid = mask.sum().item()
+                if n_valid == 0:
+                    continue
+                
+                logits_chunk = self.lm_head(h_chunk).float()
+                chunk_loss = F.cross_entropy(
+                    logits_chunk,
+                    labels_chunk,
+                    ignore_index=-100,
+                    reduction="sum",
+                )
+                total_ce_loss = total_ce_loss + chunk_loss
+                valid_tokens += n_valid
+            
+            ce_loss = total_ce_loss / max(1, valid_tokens) if valid_tokens > 0 else (shift_h.sum() * 0.0)
+            avg_aux = total_aux_loss / max(1, n_hse_layers)
+            loss = ce_loss + avg_aux
+            logits = None if self.training else self.lm_head(h)
+        else:
+            logits = self.lm_head(h)
+        
+        output = {
+            "loss": loss,
+            "logits": logits,
+            "aux_loss": total_aux_loss / max(1, n_hse_layers),
+        }
+        if output_hidden_states:
+            output["hidden_states"] = all_hidden_states
+        
+        return output
+    
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 256,
+        temperature: float = 1.0,
+        top_p: float = 0.9,
+        top_k: int = 50,
+        repetition_penalty: float = 1.1,
+        eos_token_id: int = 2,
+        pad_token_id: int = 0,
+    ) -> torch.Tensor:
+        """
+        Autoregressive generation with top-p/top-k sampling.
+        Clears KV caches after generation.
+        """
+        self.eval()
+        B, T_prompt = input_ids.shape
+        device = input_ids.device
+        
+        generated = input_ids.clone()
+        
+        # Clear KV caches in MHA layers
+        for layer in self.layers:
+            if isinstance(layer.seq_mixer, GlobalMHALayer):
+                layer.seq_mixer.clear_kv_cache()
+        
+        for step in range(max_new_tokens):
+            # Only process new tokens after first step
+            if step == 0:
+                curr_input = generated
+            else:
+                curr_input = generated[:, -1:]  # only last token
+            
+            outputs = self.forward(curr_input, use_cache=True, return_aux_loss=False)
+            logits = outputs["logits"][:, -1, :]  # (B, vocab_size)
+            
+            # Repetition penalty
+            if repetition_penalty != 1.0:
+                for b in range(B):
+                    for tok_id in set(generated[b].tolist()):
+                        if logits[b, tok_id] < 0:
+                            logits[b, tok_id] *= repetition_penalty
+                        else:
+                            logits[b, tok_id] /= repetition_penalty
+            
+            # Temperature
+            if temperature != 1.0:
+                logits = logits / temperature
+            
+            # Top-k filtering
+            if top_k > 0:
+                topk_vals = torch.topk(logits, min(top_k, logits.shape[-1])).values
+                logits = logits.masked_fill(logits < topk_vals[:, -1:], float('-inf'))
+            
+            # Top-p (nucleus) sampling
+            if top_p < 1.0:
+                sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+                cumprobs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                # Remove tokens beyond cumulative probability top_p
+                remove_mask = cumprobs - F.softmax(sorted_logits, dim=-1) > top_p
+                sorted_logits[remove_mask] = float('-inf')
+                logits = torch.scatter(logits, 1, sorted_idx, sorted_logits)
+            
+            # Sample
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)  # (B, 1)
+            
+            generated = torch.cat([generated, next_token], dim=1)
+            
+            # Stop if all sequences hit EOS
+            if (generated == eos_token_id).any(dim=1).all():
+                break
+        
+        # Clear caches
+        for layer in self.layers:
+            if isinstance(layer.seq_mixer, GlobalMHALayer):
+                layer.seq_mixer.clear_kv_cache()
+        
+        return generated
+    
+    def count_parameters(self) -> Dict[str, int]:
+        """Count parameters by component."""
+        counts = {}
+        counts["total"] = sum(p.numel() for p in self.parameters())
+        counts["embedding"] = self.embed_tokens.weight.numel()
+        counts["arg_layers"] = sum(
+            p.numel() for i, l in enumerate(self.layers)
+            if isinstance(l.seq_mixer, ARGLayer)
+            for p in l.seq_mixer.parameters()
+        )
+        counts["mha_layers"] = sum(
+            p.numel() for l in self.layers
+            if isinstance(l.seq_mixer, GlobalMHALayer)
+            for p in l.seq_mixer.parameters()
+        )
+        counts["dense_ffn"] = sum(
+            p.numel() for l in self.layers
+            if isinstance(l.ffn, DenseSwiGLUFFN)
+            for p in l.ffn.parameters()
+        )
+        counts["hse_ffn_total"] = sum(
+            p.numel() for l in self.layers
+            if isinstance(l.ffn, HSELayer)
+            for p in l.ffn.parameters()
+        )
+        hse_active = 0
+        for l in self.layers:
+            if isinstance(l.ffn, HSELayer):
+                hse_active += l.ffn.top_k * 3 * l.ffn.d_model * l.ffn.d_ff_expert
+        counts["hse_ffn_active"] = hse_active
+        return counts
+
+
+def build_forge_500m(config=None) -> ForgeModel:
+    """Build FORGE-500M (500M total params) from config (or default config)."""
+    from config import ForgeModelConfig
+    if config is None:
+        config = ForgeModelConfig()
+    model = ForgeModel(config)
+    return model
+
+
+def build_forge_3b(config=None) -> ForgeModel:
+    """Alias for build_forge_500m (FORGE architecture downgraded to 1B max params)."""
+    return build_forge_500m(config)

model/hse_layer.py

@@ -0,0 +1,277 @@
+"""
+HSE (Hierarchical Sparse Expert) FFN Layer.
+
+Two-tier MoE:
+  Tier 1: Select 1 semantic domain from K1=4 domains
+  Tier 2: Select top-2 specialists from K2=8 experts within that domain
+  Total: 32 experts, 2 active per token
+
+Load balanced at both tiers independently for stable training.
+"""
+
+from __future__ import annotations
+import math
+from typing import Optional, Tuple, Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .dgn_norm import build_norm
+from .triton_kernels import fused_swiglu
+
+
+class FeedForwardExpert(nn.Module):
+    """Single SwiGLU expert in the HSE layer."""
+    
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+        self.gate = nn.Linear(d_model, d_ff, bias=False)
+        self.up   = nn.Linear(d_model, d_ff, bias=False)
+        self.down = nn.Linear(d_ff, d_model, bias=False)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(fused_swiglu(self.gate(x), self.up(x)))
+
+
+class DenseSwiGLUFFN(nn.Module):
+    """Standard Dense SwiGLU FFN (used in alternating layers)."""
+    
+    def __init__(
+        self,
+        d_model: int = 1280,
+        d_ff: int = 3200,
+        norm_type: str = "dgn",
+        dgn_n_groups: int = 16,
+        norm_eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)
+        self.gate = nn.Linear(d_model, d_ff, bias=False)
+        self.up   = nn.Linear(d_model, d_ff, bias=False)
+        self.down = nn.Linear(d_ff, d_model, bias=False)
+        
+        nn.init.normal_(self.gate.weight, std=0.02)
+        nn.init.normal_(self.up.weight,   std=0.02)
+        nn.init.normal_(self.down.weight, std=0.02 / math.sqrt(2))
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        h = self.norm(x)
+        return residual + self.down(fused_swiglu(self.gate(h), self.up(h)))
+
+
+class HSELayer(nn.Module):
+    """
+    Hierarchical Sparse Expert FFN Layer.
+    
+    Architecture:
+      pre-norm → Tier1-router → select domain d*
+             → Tier2-router[d*] → select top-2 experts
+             → weighted sum of 2 expert outputs
+             → residual add
+    """
+    
+    def __init__(
+        self,
+        d_model: int = 1280,
+        n_domains: int = 4,
+        n_experts_per_domain: int = 8,
+        top_k: int = 2,
+        d_ff_expert: int = 288,
+        capacity_factor: float = 1.25,
+        expert_dropout: float = 0.1,
+        aux_loss_alpha: float = 0.01,
+        gumbel_tau_init: float = 1.0,
+        norm_type: str = "dgn",
+        dgn_n_groups: int = 16,
+        norm_eps: float = 1e-6,
+        layer_idx: int = 0,
+    ):
+        super().__init__()
+        assert top_k <= n_experts_per_domain
+        
+        self.d_model = d_model
+        self.n_domains = n_domains
+        self.n_experts_per_domain = n_experts_per_domain
+        self.n_experts_total = n_domains * n_experts_per_domain
+        self.top_k = top_k
+        self.d_ff_expert = d_ff_expert
+        self.capacity_factor = capacity_factor
+        self.expert_dropout = expert_dropout
+        self.aux_loss_alpha = aux_loss_alpha
+        self.layer_idx = layer_idx
+        
+        # Pre-norm
+        self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)
+        
+        # ── Tier-1 Router ────────────────────────────────────────────────────
+        # Selects semantic domain: d* = argmax(W_r1 · x / τ)
+        self.tier1_router = nn.Linear(d_model, n_domains, bias=False)
+        self.gumbel_tau = nn.Parameter(
+            torch.tensor(gumbel_tau_init), requires_grad=False  # annealed externally
+        )
+        
+        # ── Tier-2 Routers (one per domain) ──────────────────────────────────
+        self.tier2_routers = nn.ModuleList([
+            nn.Linear(d_model, n_experts_per_domain, bias=False)
+            for _ in range(n_domains)
+        ])
+        
+        # ── Experts ───────────────────────────────────────────────────────────
+        self.experts = nn.ModuleList([
+            FeedForwardExpert(d_model, d_ff_expert)
+            for _ in range(self.n_experts_total)
+        ])
+        
+        # ── Dropout for expert regularization ────────────────────────────────
+        self.expert_dropout_layer = nn.Dropout(expert_dropout)
+        
+        # ── Aux loss accumulator (for logging) ────────────────────────────────
+        self._last_aux_loss: Optional[torch.Tensor] = None
+        
+        self._init_weights()
+    
+    def _init_weights(self):
+        """Initialize router weights to near-uniform distribution."""
+        nn.init.normal_(self.tier1_router.weight, std=0.02)
+        for router in self.tier2_routers:
+            nn.init.normal_(router.weight, std=0.02)
+        for expert in self.experts:
+            nn.init.normal_(expert.gate.weight, std=0.02)
+            nn.init.normal_(expert.up.weight, std=0.02)
+            nn.init.normal_(expert.down.weight, std=0.02 / math.sqrt(2))
+    
+    def update_gumbel_tau(self, progress: float):
+        """
+        Anneal Gumbel-softmax temperature from init to final.
+        progress: float in [0, 1] representing training progress.
+        """
+        tau_init = 1.0
+        tau_final = 0.1
+        tau = tau_init * (tau_final / tau_init) ** progress
+        self.gumbel_tau.data.fill_(tau)
+    
+    def _compute_aux_loss(
+        self,
+        tier1_probs: torch.Tensor,  # (N, n_domains) 
+        domain_idx: torch.Tensor,   # (N,)
+        x_flat: torch.Tensor,       # (N, d_model)
+    ) -> torch.Tensor:
+        """
+        Hierarchical load-balance auxiliary loss.
+        
+        Loss = α * [Tier1_balance + Tier2_balance]
+        
+        Tier1_balance = n_domains * Σ_d [f_d * P_d]
+        Tier2_balance = Σ_d {n_exp * Σ_e [f_{d,e} * P_{d,e}]} (over assigned tokens)
+        """
+        N = x_flat.shape[0]
+        
+        # Tier-1 balance
+        # f_d: fraction of tokens routed to domain d
+        one_hot_d = F.one_hot(domain_idx, self.n_domains).float()  # (N, n_domains)
+        f1 = one_hot_d.mean(0)                                      # (n_domains,)
+        P1 = tier1_probs.mean(0)                                    # (n_domains,)
+        loss_t1 = self.n_domains * (f1 * P1).sum()
+        
+        # Tier-2 balance (over tokens in each domain)
+        loss_t2 = torch.zeros(1, device=x_flat.device)
+        
+        for d in range(self.n_domains):
+            mask = (domain_idx == d)
+            n_d = mask.sum()
+            if n_d < 2:
+                continue
+            
+            x_d = x_flat[mask]                                      # (n_d, d_model)
+            t2_logits = self.tier2_routers[d](x_d)                 # (n_d, n_exp)
+            t2_probs = F.softmax(t2_logits, dim=-1)                 # (n_d, n_exp)
+            
+            # Top-k indices for f computation
+            topk_idx = t2_logits.topk(self.top_k, dim=-1).indices  # (n_d, top_k)
+            one_hot_e = F.one_hot(topk_idx, self.n_experts_per_domain).float()  # (n_d, k, n_exp)
+            f2 = one_hot_e.sum(1).mean(0) / self.top_k             # (n_exp,)
+            P2 = t2_probs.mean(0)                                   # (n_exp,)
+            
+            loss_t2 = loss_t2 + self.n_experts_per_domain * (f2 * P2).sum()
+        
+        return self.aux_loss_alpha * (loss_t1 + loss_t2)
+    
+    def forward(
+        self,
+        x: torch.Tensor,  # (B, T, d_model)
+        return_aux_loss: bool = True,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Forward pass with hierarchical routing (vectorized for maximum GPU throughput).
+        Returns (output, aux_loss) where aux_loss is for load balancing.
+        """
+        B, T, D = x.shape
+        N = B * T
+        
+        residual = x
+        x_norm = self.norm(x)
+        x_flat = x_norm.reshape(N, D)  # (N, D)
+        
+        # ── Tier-1 Routing ────────────────────────────────────────────────────
+        tier1_logits = self.tier1_router(x_flat)              # (N, n_domains)
+        tier1_probs = F.softmax(tier1_logits / float(self.gumbel_tau), dim=-1)
+        
+        if self.training:
+            domain_idx = F.gumbel_softmax(
+                tier1_logits, 
+                tau=float(self.gumbel_tau), 
+                hard=True
+            ).argmax(dim=-1)                                   # (N,)
+        else:
+            domain_idx = tier1_probs.argmax(dim=-1)            # (N,)
+        
+        # ── Tier-2 Routing (Zero-sync batched projection + gather) ───────────
+        t2_weights = torch.cat([r.weight for r in self.tier2_routers], dim=0)  # (32, D)
+        t2_all = F.linear(x_flat, t2_weights).view(N, self.n_domains, self.n_experts_per_domain)
+        
+        d_expanded = domain_idx.view(N, 1, 1).expand(-1, 1, self.n_experts_per_domain)
+        t2_logits = torch.gather(t2_all, 1, d_expanded).squeeze(1)  # (N, 8)
+        
+        t2_probs = F.softmax(t2_logits, dim=-1)
+        w, topk_idx = t2_probs.topk(self.top_k, dim=-1)
+        w = (w / (w.sum(-1, keepdim=True) + 1e-9)).to(x.dtype)
+        
+        global_exp_idx = topk_idx + (domain_idx.unsqueeze(-1) * self.n_experts_per_domain)  # (N, 2)
+        
+        # ── Vectorized Token-Expert Dispatch ──────────────────────────────────
+        flat_exp_ids = global_exp_idx.reshape(-1)
+        flat_tokens = torch.arange(N, device=x.device).unsqueeze(1).expand(-1, self.top_k).reshape(-1)
+        flat_weights = w.reshape(-1, 1)
+        
+        sort_perm = torch.argsort(flat_exp_ids)
+        sorted_exp = flat_exp_ids[sort_perm]
+        sorted_tokens = flat_tokens[sort_perm]
+        sorted_weights = flat_weights[sort_perm]
+        
+        counts = torch.bincount(sorted_exp, minlength=self.n_experts_total).tolist()
+        out_flat = torch.zeros(N, D, device=x.device, dtype=x.dtype)
+        offset = 0
+        
+        for e in range(self.n_experts_total):
+            cnt = counts[e]
+            if cnt > 0:
+                t_idx = sorted_tokens[offset : offset + cnt]
+                w_e = sorted_weights[offset : offset + cnt]
+                x_e = x_flat[t_idx]
+                if self.training and self.expert_dropout > 0:
+                    x_e = self.expert_dropout_layer(x_e)
+                y_e = self.experts[e](x_e)
+                out_flat.index_add_(0, t_idx, w_e * y_e)
+                offset += cnt
+                
+        output = out_flat.reshape(B, T, D)
+        
+        # ── Auxiliary loss ────────────────────────────────────────────────────
+        aux_loss = None
+        if return_aux_loss and self.training:
+            aux_loss = self._compute_aux_loss(tier1_probs, domain_idx, x_flat)
+            self._last_aux_loss = aux_loss.detach()
+        
+        return residual + output, aux_loss

model/mha_layer.py

@@ -0,0 +1,170 @@
+"""
+Full Multi-Head Attention Layer for FORGE's global attention blocks.
+Uses FlashAttention-3 when available, with GQA and RoPE.
+"""
+
+from __future__ import annotations
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from .dgn_norm import build_norm
+from .rotary_embedding import RotaryEmbedding
+
+try:
+    from flash_attn import flash_attn_func
+    FLASH_ATTN = True
+except ImportError:
+    FLASH_ATTN = False
+
+
+class GlobalMHALayer(nn.Module):
+    """
+    Global Multi-Head Attention with:
+    - GQA (Grouped Query Attention) with n_kv_heads < n_heads
+    - RoPE positional encoding (long-context base 500k)
+    - FlashAttention-3 for maximum GPU utilization
+    - KV cache support for efficient inference
+    - Pre-norm with DGN
+    """
+    
+    def __init__(
+        self,
+        d_model: int = 1280,
+        n_heads: int = 16,
+        n_kv_heads: int = 4,
+        head_dim: int = 80,
+        rope_base: float = 500_000.0,
+        rope_scaling_type: Optional[str] = None,
+        rope_scaling_factor: float = 1.0,
+        max_seq_len: int = 4096,
+        dropout: float = 0.0,
+        norm_type: str = "dgn",
+        dgn_n_groups: int = 16,
+        norm_eps: float = 1e-6,
+        use_flash_attention: bool = True,
+        layer_idx: int = 0,
+    ):
+        super().__init__()
+        assert d_model == n_heads * head_dim, \
+            f"d_model={d_model} must equal n_heads*head_dim={n_heads*head_dim}"
+        assert n_heads % n_kv_heads == 0, \
+            f"n_heads={n_heads} must be divisible by n_kv_heads={n_kv_heads}"
+        
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = head_dim
+        self.n_rep = n_heads // n_kv_heads  # GQA repetition factor
+        self.dropout = dropout
+        self.use_flash = use_flash_attention and FLASH_ATTN
+        self.layer_idx = layer_idx
+        self.scale = head_dim ** -0.5
+        
+        # Pre-norm
+        self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)
+        
+        # Q/K/V projections
+        self.q_proj = nn.Linear(d_model, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(d_model, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(d_model, n_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(d_model, d_model, bias=False)
+        
+        # RoPE
+        self.rope = RotaryEmbedding(
+            head_dim=head_dim,
+            max_seq_len=max_seq_len,
+            base=rope_base,
+            scaling_type=rope_scaling_type,
+            scaling_factor=rope_scaling_factor,
+        )
+        
+        # KV cache for inference
+        self._kv_cache_k: Optional[torch.Tensor] = None
+        self._kv_cache_v: Optional[torch.Tensor] = None
+        self._kv_cache_len: int = 0
+        
+        self._init_weights()
+    
+    def _init_weights(self):
+        nn.init.normal_(self.q_proj.weight, std=0.02)
+        nn.init.normal_(self.k_proj.weight, std=0.02)
+        nn.init.normal_(self.v_proj.weight, std=0.02)
+        # Output projection: scale by 1/sqrt(2 * n_layers) for residual depth
+        nn.init.normal_(self.o_proj.weight, std=0.02)
+    
+    def clear_kv_cache(self):
+        self._kv_cache_k = None
+        self._kv_cache_v = None
+        self._kv_cache_len = 0
+    
+    def forward(
+        self,
+        x: torch.Tensor,                           # (B, T, d_model)
+        attention_mask: Optional[torch.Tensor] = None,  # (B, T) or (B, 1, T, T)
+        position_ids: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple]]:
+        
+        B, T, _ = x.shape
+        residual = x
+        x_norm = self.norm(x)
+        
+        # Project
+        Q = self.q_proj(x_norm)   # (B, T, n_heads * head_dim)
+        K = self.k_proj(x_norm)   # (B, T, n_kv_heads * head_dim)
+        V = self.v_proj(x_norm)   # (B, T, n_kv_heads * head_dim)
+        
+        # Reshape to (B, H, T, D)
+        Q = rearrange(Q, 'b t (h d) -> b h t d', d=self.head_dim)
+        K = rearrange(K, 'b t (h d) -> b h t d', d=self.head_dim)
+        V = rearrange(V, 'b t (h d) -> b h t d', d=self.head_dim)
+        
+        # RoPE
+        Q, K = self.rope.apply_rotary(Q, K, position_ids)
+        
+        # KV cache concat (inference)
+        if use_cache:
+            if self._kv_cache_k is not None:
+                K = torch.cat([self._kv_cache_k, K], dim=2)
+                V = torch.cat([self._kv_cache_v, V], dim=2)
+            self._kv_cache_k = K
+            self._kv_cache_v = V
+            self._kv_cache_len = K.shape[2]
+        
+        # GQA: expand K,V from n_kv_heads to n_heads
+        if self.n_rep > 1:
+            K = K.repeat_interleave(self.n_rep, dim=1)
+            V = V.repeat_interleave(self.n_rep, dim=1)
+        
+        # Attention
+        if self.use_flash:
+            # FlashAttention: expects (B, T, H, D)
+            Q_fa = rearrange(Q, 'b h t d -> b t h d').contiguous()
+            K_fa = rearrange(K, 'b h t d -> b t h d').contiguous()
+            V_fa = rearrange(V, 'b h t d -> b t h d').contiguous()
+            
+            out = flash_attn_func(
+                Q_fa, K_fa, V_fa,
+                dropout_p=self.dropout if self.training else 0.0,
+                causal=True,
+                softmax_scale=self.scale,
+            )  # (B, T, H, D)
+            out = rearrange(out, 'b t h d -> b t (h d)')
+        else:
+            # Standard PyTorch SDPA (fused, memory efficient)
+            out = F.scaled_dot_product_attention(
+                Q, K, V,
+                attn_mask=attention_mask,
+                dropout_p=self.dropout if self.training else 0.0,
+                is_causal=(attention_mask is None),
+                scale=self.scale,
+            )  # (B, H, T, D)
+            out = rearrange(out, 'b h t d -> b t (h d)')
+        
+        out = self.o_proj(out)  # (B, T, d_model)
+        
+        return residual + out, (K, V) if use_cache else None

model/pscan.py

@@ -0,0 +1,181 @@
+"""
+Parallel Scan (Blelloch algorithm) in pure PyTorch for ARG selective SSM recurrence.
+
+Eliminates all external dependencies (mamba-ssm, selective_scan_cuda).
+Runs in 2*log2(T) parallel steps (11 steps for T=2048) on any GPU.
+100% mathematically exact equivalence to sequential SSM recurrence.
+"""
+
+from __future__ import annotations
+import math
+import torch
+import torch.nn.functional as F
+
+
+def npo2(length: int) -> int:
+    """Returns the next power of 2 above length."""
+    return 2 ** math.ceil(math.log2(length))
+
+
+def pad_npo2(x: torch.Tensor) -> torch.Tensor:
+    """Pads input length dimension (dim 1) to the next power of 2."""
+    length = x.size(1)
+    len_p2 = npo2(length)
+    if len_p2 == length:
+        return x
+    pad_tuple = (0, 0, 0, 0, 0, len_p2 - length)
+    return F.pad(x, pad_tuple, "constant", 0)
+
+
+class PScan(torch.autograd.Function):
+    """
+    Computes the parallel associative scan:
+        H[t] = A[t] * H[t-1] + X[t]  with H[-1] = 0
+    in O(log2(T)) parallel steps via Blelloch tree reduction.
+    """
+
+    @staticmethod
+    def pscan(A: torch.Tensor, X: torch.Tensor) -> None:
+        # A: (B, D, L, N)
+        # X: (B, D, L, N)
+        B, D, L, _ = A.size()
+        num_steps = int(math.log2(L))
+
+        # ── Up sweep ──────────────────────────────────────────────────────────
+        Aa = A
+        Xa = X
+        for _ in range(num_steps - 2):
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T // 2, 2, -1)
+            Xa = Xa.view(B, D, T // 2, 2, -1)
+
+            Xa[:, :, :, 1].add_(Aa[:, :, :, 1].mul(Xa[:, :, :, 0]))
+            Aa[:, :, :, 1].mul_(Aa[:, :, :, 0])
+
+            Aa = Aa[:, :, :, 1]
+            Xa = Xa[:, :, :, 1]
+
+        if Xa.size(2) == 4:
+            Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+            Aa[:, :, 1].mul_(Aa[:, :, 0])
+            Xa[:, :, 3].add_(Aa[:, :, 3].mul(Xa[:, :, 2] + Aa[:, :, 2].mul(Xa[:, :, 1])))
+        elif Xa.size(2) == 2:
+            Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+            return
+        else:
+            return
+
+        # ── Down sweep ────────────────────────────────────────────────────────
+        Aa = A[:, :, 2 ** (num_steps - 2) - 1 : L : 2 ** (num_steps - 2)]
+        Xa = X[:, :, 2 ** (num_steps - 2) - 1 : L : 2 ** (num_steps - 2)]
+        Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 1]))
+        Aa[:, :, 2].mul_(Aa[:, :, 1])
+
+        for k in range(num_steps - 3, -1, -1):
+            stride = 2 ** k
+            Aa = A[:, :, stride - 1 : L : stride]
+            Xa = X[:, :, stride - 1 : L : stride]
+
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T // 2, 2, -1)
+            Xa = Xa.view(B, D, T // 2, 2, -1)
+
+            Xa[:, :, 1:, 0].add_(Aa[:, :, 1:, 0].mul(Xa[:, :, :-1, 1]))
+            Aa[:, :, 1:, 0].mul_(Aa[:, :, :-1, 1])
+
+    @staticmethod
+    def pscan_rev(A: torch.Tensor, X: torch.Tensor) -> None:
+        # Backward scan in reverse time order
+        B, D, L, _ = A.size()
+        num_steps = int(math.log2(L))
+
+        # Up sweep
+        Aa = A
+        Xa = X
+        for _ in range(num_steps - 2):
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T // 2, 2, -1)
+            Xa = Xa.view(B, D, T // 2, 2, -1)
+
+            Xa[:, :, :, 0].add_(Aa[:, :, :, 0].mul(Xa[:, :, :, 1]))
+            Aa[:, :, :, 0].mul_(Aa[:, :, :, 1])
+
+            Aa = Aa[:, :, :, 0]
+            Xa = Xa[:, :, :, 0]
+
+        if Xa.size(2) == 4:
+            Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 3]))
+            Aa[:, :, 2].mul_(Aa[:, :, 3])
+            Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1].add(Aa[:, :, 1].mul(Xa[:, :, 2]))))
+        elif Xa.size(2) == 2:
+            Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1]))
+            return
+        else:
+            return
+
+        # Down sweep
+        stride_top = 2 ** (num_steps - 2)
+        Aa = A[:, :, 0 : L : stride_top]
+        Xa = X[:, :, 0 : L : stride_top]
+        Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 2]))
+        Aa[:, :, 1].mul_(Aa[:, :, 2])
+
+        for k in range(num_steps - 3, -1, -1):
+            stride = 2 ** k
+            Aa = A[:, :, 0 : L : stride]
+            Xa = X[:, :, 0 : L : stride]
+
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T // 2, 2, -1)
+            Xa = Xa.view(B, D, T // 2, 2, -1)
+
+            Xa[:, :, :-1, 1].add_(Aa[:, :, :-1, 1].mul(Xa[:, :, 1:, 0]))
+            Aa[:, :, :-1, 1].mul_(Aa[:, :, 1:, 0])
+
+    @staticmethod
+    def forward(ctx, A_in: torch.Tensor, X_in: torch.Tensor) -> torch.Tensor:
+        # A_in: (B, L, D, N)
+        # X_in: (B, L, D, N)
+        L = X_in.size(1)
+
+        if L == npo2(L):
+            A = A_in.clone()
+            X = X_in.clone()
+        else:
+            A = pad_npo2(A_in).clone()
+            X = pad_npo2(X_in).clone()
+
+        A = A.transpose(2, 1).contiguous()  # (B, D, L_p2, N)
+        X = X.transpose(2, 1).contiguous()  # (B, D, L_p2, N)
+
+        PScan.pscan(A, X)
+
+        ctx.save_for_backward(A_in, X)
+        return X.transpose(2, 1)[:, :L].contiguous()
+
+    @staticmethod
+    def backward(ctx, grad_output_in: torch.Tensor):
+        A_in, X = ctx.saved_tensors
+        L = grad_output_in.size(1)
+
+        if L == npo2(L):
+            grad_output = grad_output_in.clone()
+        else:
+            grad_output = pad_npo2(grad_output_in).clone()
+            A_in = pad_npo2(A_in)
+
+        grad_output = grad_output.transpose(2, 1).contiguous()  # (B, D, L_p2, N)
+        A_in_t = A_in.transpose(2, 1).contiguous()             # (B, D, L_p2, N)
+
+        # Shift A 1 step to the left for adjoint recurrence
+        A_shifted = F.pad(A_in_t[:, :, 1:], (0, 0, 0, 1))
+
+        PScan.pscan_rev(A_shifted, grad_output)
+
+        grad_A = torch.zeros_like(X)
+        grad_A[:, :, 1:].add_(X[:, :, :-1] * grad_output[:, :, 1:])
+
+        return grad_A.transpose(2, 1)[:, :L].contiguous(), grad_output.transpose(2, 1)[:, :L].contiguous()
+
+
+pscan = PScan.apply

model/rotary_embedding.py

@@ -0,0 +1,168 @@
+"""
+RoPE (Rotary Positional Embedding) with YaRN support.
+GPU-optimized via cached cos/sin buffers and optional Triton kernel.
+"""
+
+from __future__ import annotations
+import math
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple
+
+
+class RotaryEmbedding(nn.Module):
+    """
+    Rotary Positional Embedding (RoPE) with:
+    - Long-context theta (500,000 base)
+    - YaRN scaling for context extension
+    - GPU-resident cos/sin cache (no recomputation per step)
+    - BF16-native computation
+    """
+    
+    def __init__(
+        self,
+        head_dim: int,
+        max_seq_len: int = 4096,
+        base: float = 500_000.0,
+        scaling_type: Optional[str] = None,
+        scaling_factor: float = 1.0,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.max_seq_len = max_seq_len
+        self.base = base
+        self.scaling_type = scaling_type
+        self.scaling_factor = scaling_factor
+        
+        # Compute inverse frequencies
+        inv_freq = self._build_inv_freq(head_dim, base, scaling_type, scaling_factor)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        
+        # Pre-compute and cache cos/sin for all positions
+        self._cos_cache: Optional[torch.Tensor] = None
+        self._sin_cache: Optional[torch.Tensor] = None
+        self._cache_seq_len: int = 0
+        
+        # Build initial cache
+        self._build_cache(max_seq_len, device)
+    
+    def _build_inv_freq(
+        self,
+        head_dim: int,
+        base: float,
+        scaling_type: Optional[str],
+        scaling_factor: float,
+    ) -> torch.Tensor:
+        """Compute inverse frequencies with optional YaRN scaling."""
+        dim = head_dim
+        # Standard RoPE frequencies
+        freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        
+        if scaling_type == "yarn":
+            # YaRN: scale only the low-frequency dimensions
+            # High-frequency dims (short-range) are left unchanged
+            # Low-frequency dims (long-range) are scaled down by 1/factor
+            mscale = 0.1 * math.log(scaling_factor) + 1.0  # YaRN mscale
+            
+            # Threshold: dims where wavelength > original_max_len
+            low_freq_factor = 1.0
+            high_freq_factor = 4.0
+            original_max_len = self.max_seq_len
+            new_base = base * (
+                (scaling_factor * original_max_len / (2 * math.pi)) -
+                (original_max_len / (2 * math.pi))
+            ) ** (dim / (dim - 2))
+            
+            # Apply NTK-by-parts
+            d = torch.arange(0, dim, 2).float()
+            wavelength = 2 * math.pi / freqs
+            
+            # Smooth interpolation between no-scale and full-scale
+            scale = torch.where(
+                wavelength < original_max_len / high_freq_factor,
+                torch.ones_like(freqs),
+                torch.where(
+                    wavelength > original_max_len / low_freq_factor,
+                    torch.full_like(freqs, 1.0 / scaling_factor),
+                    (original_max_len / wavelength / low_freq_factor - 1.0) / 
+                    (high_freq_factor / low_freq_factor - 1.0) *
+                    (1.0 - 1.0 / scaling_factor) + 1.0 / scaling_factor
+                )
+            )
+            freqs = freqs * scale
+        
+        elif scaling_type == "linear":
+            freqs = freqs / scaling_factor
+        
+        return freqs
+    
+    @torch.no_grad()
+    def _build_cache(self, seq_len: int, device=None):
+        """Build/extend the cos/sin cache up to seq_len."""
+        if seq_len <= self._cache_seq_len:
+            return
+        
+        device = device or (self.inv_freq.device if hasattr(self, 'inv_freq') 
+                            else torch.device('cpu'))
+        
+        t = torch.arange(seq_len, device=device, dtype=torch.float32)
+        freqs = torch.outer(t, self.inv_freq.to(device))
+        emb = torch.cat([freqs, freqs], dim=-1)  # (seq_len, head_dim)
+        
+        # Store in BF16 for fast loading
+        self._cos_cache = emb.cos().to(torch.bfloat16)
+        self._sin_cache = emb.sin().to(torch.bfloat16)
+        self._cache_seq_len = seq_len
+    
+    def get_cos_sin(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return cached cos/sin for positions 0..seq_len-1."""
+        if seq_len > self._cache_seq_len:
+            self._build_cache(seq_len * 2, device)  # extend with buffer
+        
+        cos = self._cos_cache[:seq_len].to(device)
+        sin = self._sin_cache[:seq_len].to(device)
+        return cos, sin  # each (seq_len, head_dim)
+    
+    def apply_rotary(
+        self,
+        q: torch.Tensor,   # (B, n_heads, T, head_dim)
+        k: torch.Tensor,   # (B, n_kv_heads, T, head_dim)
+        position_ids: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Apply RoPE to Q and K."""
+        T = q.shape[2]
+        device = q.device
+        
+        cos, sin = self.get_cos_sin(T, device)  # (T, head_dim)
+        
+        if position_ids is not None:
+            # Custom positions (for packed sequences or position bias)
+            cos = cos[position_ids]   # (B, T, head_dim) or (T, head_dim)
+            sin = sin[position_ids]
+            if position_ids.dim() == 2:
+                cos = cos.unsqueeze(1)  # (B, 1, T, head_dim)
+                sin = sin.unsqueeze(1)
+            else:
+                cos = cos.unsqueeze(0).unsqueeze(0)  # (1, 1, T, head_dim)
+                sin = sin.unsqueeze(0).unsqueeze(0)
+        else:
+            # Reshape for broadcasting: (1, 1, T, head_dim)
+            cos = cos.unsqueeze(0).unsqueeze(0)
+            sin = sin.unsqueeze(0).unsqueeze(0)
+        
+        q_out = _rotate_half_apply(q.float(), cos, sin).to(q.dtype)
+        k_out = _rotate_half_apply(k.float(), cos, sin).to(k.dtype)
+        
+        return q_out, k_out
+    
+    def forward(self, q, k, position_ids=None):
+        return self.apply_rotary(q, k, position_ids)
+
+
+def _rotate_half_apply(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Apply rotary: x * cos + rotate_half(x) * sin"""
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    rotated = torch.cat([-x2, x1], dim=-1)
+    return x * cos + rotated * sin

model/triton_kernels.py

@@ -0,0 +1,457 @@
+"""
+FORGE-3B Triton GPU Kernels
+Hand-written Triton kernels for maximum GPU throughput:
+- Fused RMSNorm / DGN (Differential Group Normalization)
+- Fused SwiGLU activation
+- Complex SSM scan kernel
+- Fused MoE dispatch + expert compute + gather
+- Fused RoPE injection
+
+All kernels target H100/A100 with BF16 inputs.
+"""
+
+from __future__ import annotations
+import math
+import torch
+import torch.nn.functional as F
+
+import os
+try:
+    import triton
+    import triton.language as tl
+    TRITON_AVAILABLE = os.environ.get("FORGE_NO_TRITON") != "1"
+except ImportError:
+    TRITON_AVAILABLE = False
+    import warnings
+    warnings.warn("Triton not available — falling back to PyTorch ops. "
+                  "Install triton for 2-4× kernel speedup.")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# FUSED RMSNorm KERNEL
+# ─────────────────────────────────────────────────────────────────────────────
+
+if TRITON_AVAILABLE:
+    @triton.jit
+    def _rms_norm_fwd_kernel(
+        X,          # input ptr (N, D)
+        W,          # weight ptr (D,)
+        B,          # bias ptr (D,) or None
+        Y,          # output ptr (N, D)
+        Rstd,       # 1/rms ptr (N,) — saved for backward
+        stride_x,   # stride along N dim
+        N,          # number of rows
+        D,          # row dimension
+        eps,
+        HAS_BIAS: tl.constexpr,
+        BLOCK_D: tl.constexpr,
+    ):
+        row = tl.program_id(0)
+        X_ptr = X + row * stride_x
+        Y_ptr = Y + row * stride_x
+        
+        # Compute mean square
+        ms = tl.zeros([BLOCK_D], dtype=tl.float32)
+        for off in range(0, D, BLOCK_D):
+            cols = off + tl.arange(0, BLOCK_D)
+            mask = cols < D
+            x = tl.load(X_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+            ms += x * x
+        ms = tl.sum(ms, axis=0) / D
+        rstd = 1.0 / tl.sqrt(ms + eps)
+        tl.store(Rstd + row, rstd)
+        
+        # Normalize and scale
+        for off in range(0, D, BLOCK_D):
+            cols = off + tl.arange(0, BLOCK_D)
+            mask = cols < D
+            x = tl.load(X_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+            w = tl.load(W + cols, mask=mask, other=1.0).to(tl.float32)
+            y = x * rstd * w
+            if HAS_BIAS:
+                b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32)
+                y = y + b
+            tl.store(Y_ptr + cols, y.to(tl.bfloat16), mask=mask)
+    
+    
+    @triton.jit
+    def _dgn_fwd_kernel(
+        X,              # (N, D) input
+        W,              # (G, Dg) per-group weight
+        B,              # (G, Dg) per-group bias
+        Y,              # (N, D) output
+        Rstd,           # (N, G) per-group rstd
+        stride_x,
+        N, D, G,        # n_rows, d_model, n_groups
+        Dg,             # group size = D / G
+        eps,
+        BLOCK_Dg: tl.constexpr,
+    ):
+        """Differential Group Normalization: separate RMS per feature group."""
+        row = tl.program_id(0)
+        grp = tl.program_id(1)
+        
+        base_x = row * stride_x + grp * Dg
+        
+        # Compute group RMS
+        ms = tl.zeros([BLOCK_Dg], dtype=tl.float32)
+        for off in range(0, Dg, BLOCK_Dg):
+            cols = off + tl.arange(0, BLOCK_Dg)
+            mask = cols < Dg
+            x = tl.load(X + base_x + cols, mask=mask, other=0.0).to(tl.float32)
+            ms += x * x
+        ms = tl.sum(ms, axis=0) / Dg
+        rstd = 1.0 / tl.sqrt(ms + eps)
+        tl.store(Rstd + row * G + grp, rstd)
+        
+        # Normalize, scale, bias
+        for off in range(0, Dg, BLOCK_Dg):
+            cols = off + tl.arange(0, BLOCK_Dg)
+            mask = cols < Dg
+            x = tl.load(X + base_x + cols, mask=mask, other=0.0).to(tl.float32)
+            w = tl.load(W + grp * Dg + cols, mask=mask, other=1.0).to(tl.float32)
+            b = tl.load(B + grp * Dg + cols, mask=mask, other=0.0).to(tl.float32)
+            y = x * rstd * w + b
+            tl.store(Y + base_x + cols, y.to(tl.bfloat16), mask=mask)
+
+
+    @triton.jit
+    def _swiglu_fwd_kernel(
+        Gate,   # (N, D) gate branch
+        Up,     # (N, D) up branch  
+        Out,    # (N, D) output = silu(Gate) * Up
+        N, D,
+        BLOCK_D: tl.constexpr,
+    ):
+        """Fused SwiGLU: Out = SiLU(Gate) * Up"""
+        row = tl.program_id(0)
+        for off in range(0, D, BLOCK_D):
+            cols = off + tl.arange(0, BLOCK_D)
+            mask = cols < D
+            g = tl.load(Gate + row * D + cols, mask=mask).to(tl.float32)
+            u = tl.load(Up   + row * D + cols, mask=mask).to(tl.float32)
+            # SiLU(x) = x * sigmoid(x)
+            silu_g = g * tl.sigmoid(g)
+            out = silu_g * u
+            tl.store(Out + row * D + cols, out.to(tl.bfloat16), mask=mask)
+
+    
+    @triton.jit
+    def _rope_fwd_kernel(
+        Q,          # (B, H, T, D) queries
+        K,          # (B, H_kv, T, D) keys
+        Cos,        # (T, D/2) cosine
+        Sin,        # (T, D/2) sine
+        Q_out,
+        K_out,
+        B, H_q, H_kv, T, D,
+        BLOCK_T: tl.constexpr,
+        BLOCK_D: tl.constexpr,
+    ):
+        """Fused RoPE for Q and K."""
+        b_idx = tl.program_id(0)
+        h_idx = tl.program_id(1)
+        is_kv_head = h_idx >= H_q
+        
+        for t in range(0, T, BLOCK_T):
+            t_off = t + tl.arange(0, BLOCK_T)
+            t_mask = t_off < T
+            
+            for d in range(0, D // 2, BLOCK_D):
+                d_off = d + tl.arange(0, BLOCK_D)
+                d_mask = d_off < (D // 2)
+                
+                # Load cos/sin
+                cos = tl.load(Cos + t_off[:, None] * (D // 2) + d_off[None, :],
+                              mask=t_mask[:, None] & d_mask[None, :], other=1.0).to(tl.float32)
+                sin = tl.load(Sin + t_off[:, None] * (D // 2) + d_off[None, :],
+                              mask=t_mask[:, None] & d_mask[None, :], other=0.0).to(tl.float32)
+                
+                if not is_kv_head:
+                    ptr_base = Q + b_idx * H_q * T * D + h_idx * T * D
+                    out_base = Q_out + b_idx * H_q * T * D + h_idx * T * D
+                    # x1 = x[..., :D/2], x2 = x[..., D/2:]
+                    x1 = tl.load(ptr_base + t_off[:, None] * D + d_off[None, :],
+                                  mask=t_mask[:, None] & d_mask[None, :]).to(tl.float32)
+                    x2 = tl.load(ptr_base + t_off[:, None] * D + (D // 2) + d_off[None, :],
+                                  mask=t_mask[:, None] & d_mask[None, :]).to(tl.float32)
+                    y1 = x1 * cos - x2 * sin
+                    y2 = x1 * sin + x2 * cos
+                    tl.store(out_base + t_off[:, None] * D + d_off[None, :],
+                             y1.to(tl.bfloat16), mask=t_mask[:, None] & d_mask[None, :])
+                    tl.store(out_base + t_off[:, None] * D + (D // 2) + d_off[None, :],
+                             y2.to(tl.bfloat16), mask=t_mask[:, None] & d_mask[None, :])
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# PYTHON WRAPPERS AROUND TRITON KERNELS
+# ─────────────────────────────────────────────────────────────────────────────
+
+class FusedDGN(torch.autograd.Function):
+    """
+    Fused Differential Group Normalization using Triton kernel.
+    Fallback to PyTorch if Triton unavailable.
+    """
+    
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, weight: torch.Tensor, 
+                bias: torch.Tensor, n_groups: int, eps: float):
+        """
+        x: (batch, seq, d_model) or (N, d_model)
+        weight, bias: (d_model,)
+        """
+        orig_shape = x.shape
+        ctx.orig_shape = orig_shape
+        x_2d = x.reshape(-1, x.shape[-1]).contiguous()
+        N, D = x_2d.shape
+        G = n_groups
+        Dg = D // G
+        
+        if not TRITON_AVAILABLE or not x.is_cuda:
+            # Pure PyTorch fallback
+            out = _dgn_pytorch(x_2d, weight, bias, G, Dg, eps)
+            ctx.save_for_backward(x_2d, weight, bias, 
+                                  _compute_rstd_pytorch(x_2d, G, Dg, eps))
+            ctx.n_groups = G
+            ctx.Dg = Dg
+            ctx.eps = eps
+            return out.reshape(orig_shape)
+        
+        y = torch.empty_like(x_2d)
+        rstd = torch.empty(N, G, device=x.device, dtype=torch.float32)
+        
+        BLOCK_Dg = min(triton.next_power_of_2(Dg), 128)
+        grid = (N, G)
+        
+        _dgn_fwd_kernel[grid](
+            x_2d, weight.reshape(G, Dg), bias.reshape(G, Dg), 
+            y, rstd,
+            x_2d.stride(0), N, D, G, Dg, eps,
+            BLOCK_Dg=BLOCK_Dg,
+        )
+        
+        ctx.save_for_backward(x_2d, weight, bias, rstd)
+        ctx.n_groups = G
+        ctx.Dg = Dg
+        ctx.eps = eps
+        return y.reshape(orig_shape)
+    
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, weight, bias, rstd = ctx.saved_tensors
+        G, Dg, eps = ctx.n_groups, ctx.Dg, ctx.eps
+        
+        # Backward via PyTorch autograd (exact, Triton bwd TBD)
+        with torch.enable_grad():
+            x_r = x.detach().requires_grad_(True)
+            w_r = weight.detach().requires_grad_(True)
+            b_r = bias.detach().requires_grad_(True)
+            y = _dgn_pytorch(x_r, w_r, b_r, G, Dg, eps)
+            y.backward(grad_output.reshape_as(y))
+        
+        return x_r.grad.reshape(ctx.orig_shape), w_r.grad, b_r.grad, None, None
+
+
+def _dgn_pytorch(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor,
+                  G: int, Dg: int, eps: float) -> torch.Tensor:
+    """PyTorch reference for DGN — used in backward and as fallback."""
+    N, D = x.shape
+    x_g = x.reshape(N, G, Dg)
+    rms = x_g.float().pow(2).mean(dim=-1, keepdim=True).add(eps).sqrt()
+    x_norm = (x_g / rms).reshape(N, D)
+    return x_norm.to(x.dtype) * weight + bias
+
+
+def _compute_rstd_pytorch(x: torch.Tensor, G: int, Dg: int, eps: float) -> torch.Tensor:
+    N, D = x.shape
+    x_g = x.reshape(N, G, Dg).float()
+    rms = x_g.pow(2).mean(dim=-1).add(eps).sqrt()
+    return 1.0 / rms
+
+
+def fused_dgn(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor,
+               n_groups: int = 16, eps: float = 1e-6) -> torch.Tensor:
+    """Entry point for fused DGN — uses Triton if available, else PyTorch."""
+    return FusedDGN.apply(x, weight, bias, n_groups, eps)
+
+
+def fused_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    """Fused SwiGLU: output = SiLU(gate) * up"""
+    if TRITON_AVAILABLE and gate.is_cuda and gate.dtype == torch.bfloat16:
+        orig_shape = gate.shape
+        gate_2d = gate.reshape(-1, gate.shape[-1]).contiguous()
+        up_2d = up.reshape(-1, up.shape[-1]).contiguous()
+        N, D = gate_2d.shape
+        out = torch.empty_like(gate_2d)
+        
+        BLOCK_D = min(triton.next_power_of_2(D), 512)
+        grid = (N,)
+        _swiglu_fwd_kernel[grid](gate_2d, up_2d, out, N, D, BLOCK_D=BLOCK_D)
+        return out.reshape(orig_shape)
+    else:
+        return F.silu(gate) * up
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# COMPLEX SSM PARALLEL SCAN (Triton)
+# Implements O(T) parallel prefix scan for the ARG recurrent branch.
+# For complex-valued states stored as (real, imag) interleaved.
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@torch.jit.script
+def _ssm_forward_scan(
+    ar: torch.Tensor,     # (B, T, d_state) float32
+    ai: torch.Tensor,     # (B, T, d_state)
+    bx: torch.Tensor,     # (B, T, d_state) = B_bar * x_inner[:,:,:d_state]
+    h0_real: torch.Tensor,  # (B, d_state)
+    h0_imag: torch.Tensor,  # (B, d_state)
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """TorchScript-compiled sequential SSM scan — zero Python overhead per step."""
+    B = ar.size(0)
+    T = ar.size(1)
+    d_state = ar.size(2)
+    h_real = torch.empty(B, T + 1, d_state, dtype=ar.dtype, device=ar.device)
+    h_imag = torch.empty(B, T + 1, d_state, dtype=ar.dtype, device=ar.device)
+    h_real[:, 0] = h0_real
+    h_imag[:, 0] = h0_imag
+    for t in range(T):
+        hr = h_real[:, t]
+        hi = h_imag[:, t]
+        art = ar[:, t]
+        ait = ai[:, t]
+        h_real[:, t + 1] = art * hr - ait * hi + bx[:, t]
+        h_imag[:, t + 1] = art * hi + ait * hr
+    return h_real, h_imag
+
+
+@torch.jit.script
+def _ssm_backward_scan(
+    ar: torch.Tensor,              # (B, T, d_state)
+    ai: torch.Tensor,              # (B, T, d_state)
+    bb: torch.Tensor,              # (B, T, d_state)
+    xi_slice: torch.Tensor,        # (B, T, d_state) — x_inner[:,:,:d_state]
+    cc: torch.Tensor,              # (B, T, d_state)
+    h_real: torch.Tensor,          # (B, T+1, d_state)
+    h_imag: torch.Tensor,          # (B, T+1, d_state)
+    go_scalar: torch.Tensor,       # (B, T)
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """TorchScript-compiled reverse accumulation backward — zero Python overhead per step."""
+    B = ar.size(0)
+    T = ar.size(1)
+    d_state = ar.size(2)
+
+    grad_hreal_from_y = go_scalar.unsqueeze(-1) * cc   # (B, T, d_state)
+
+    grad_ar      = torch.zeros_like(ar)
+    grad_ai      = torch.zeros_like(ai)
+    grad_bb      = torch.zeros_like(bb)
+    grad_xi_ssm  = torch.zeros_like(xi_slice)
+    delta_hr     = torch.zeros(B, d_state, dtype=ar.dtype, device=ar.device)
+    delta_hi     = torch.zeros(B, d_state, dtype=ar.dtype, device=ar.device)
+
+    for t in range(T - 1, -1, -1):
+        hr_t = h_real[:, t]
+        hi_t = h_imag[:, t]
+        art  = ar[:, t]
+        ait  = ai[:, t]
+
+        total_hr = grad_hreal_from_y[:, t] + delta_hr
+        total_hi = delta_hi
+
+        grad_ar[:, t] = total_hr * hr_t + total_hi * hi_t
+        grad_ai[:, t] = -total_hr * hi_t + total_hi * hr_t
+
+        grad_bx_t = total_hr
+        grad_bb[:, t]     = grad_bx_t * xi_slice[:, t]
+        grad_xi_ssm[:, t] = grad_bx_t * bb[:, t]
+
+        delta_hr = art * total_hr + ait * total_hi
+        delta_hi = -ait * total_hr + art * total_hi
+
+    grad_h0_real = delta_hr
+    grad_h0_imag = delta_hi
+    return grad_ar, grad_ai, grad_bb, grad_xi_ssm, grad_h0_real, grad_h0_imag
+
+
+class ComplexSSMScan(torch.autograd.Function):
+    """
+    Fused complex SSM scan with TorchScript forward + analytic backward.
+    
+    Both the forward scan and backward accumulation loops run inside
+    @torch.jit.script functions — no Python interpreter overhead per timestep.
+    """
+
+    @staticmethod
+    def forward(ctx,
+                x_inner:   torch.Tensor,  # (B, T, d_inner)
+                A_bar_real: torch.Tensor, # (B, T, d_state)
+                A_bar_imag: torch.Tensor, # (B, T, d_state)
+                B_bar:     torch.Tensor,  # (B, T, d_state)
+                C:         torch.Tensor,  # (B, T, d_state)
+                D:         torch.Tensor,  # (d_inner,)
+                h0_real:   torch.Tensor,  # (B, d_state)
+                h0_imag:   torch.Tensor,  # (B, d_state)
+               ):
+        B, T, d_inner = x_inner.shape
+        d_state = A_bar_real.shape[2]
+        dtype = x_inner.dtype
+
+        ar = A_bar_real.float()
+        ai = A_bar_imag.float()
+        bb = B_bar.float()
+        cc = C.float()
+        xi = x_inner.float()
+        D_f = D.float()
+
+        bx = bb * xi[:, :, :d_state]   # (B, T, d_state)
+
+        # TorchScript-compiled forward scan — no Python loop overhead
+        h_real, h_imag = _ssm_forward_scan(ar, ai, bx, h0_real.float(), h0_imag.float())
+
+        y_ssm = (cc * h_real[:, 1:]).sum(-1, keepdim=True)  # (B, T, 1)
+        y = y_ssm.expand(B, T, d_inner) + D_f.unsqueeze(0).unsqueeze(0) * xi
+
+        ctx.save_for_backward(ar, ai, bb, cc, xi, D_f, h_real, h_imag)
+        ctx.input_dtype = dtype
+        ctx.d_state = d_state
+        return y.to(dtype)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        ar, ai, bb, cc, xi, D_f, h_real, h_imag = ctx.saved_tensors
+        B, T, d_inner = grad_output.shape
+        d_state = ctx.d_state
+        go = grad_output.float()
+
+        grad_D  = (go * xi).sum([0, 1])
+        grad_xi = go * D_f.unsqueeze(0).unsqueeze(0)
+
+        go_scalar = go.sum(-1)          # (B, T)
+        grad_C    = go_scalar.unsqueeze(-1) * h_real[:, 1:]  # (B, T, d_state)
+
+        xi_slice = xi[:, :, :d_state]
+
+        # TorchScript-compiled backward scan — no Python loop overhead
+        grad_ar, grad_ai, grad_bb, grad_xi_ssm, grad_h0_real, grad_h0_imag = \
+            _ssm_backward_scan(ar, ai, bb, xi_slice, cc, h_real, h_imag, go_scalar)
+
+        grad_xi[:, :, :d_state] = grad_xi[:, :, :d_state] + grad_xi_ssm
+
+        dtype = ctx.input_dtype
+        return (
+            grad_xi.to(dtype),
+            grad_ar.to(dtype),
+            grad_ai.to(dtype),
+            grad_bb.to(dtype),
+            grad_C.to(dtype),
+            grad_D.to(dtype),
+            grad_h0_real.to(dtype),
+            grad_h0_imag.to(dtype),
+        )
+
+
+# Export
+def complex_ssm_scan(x_inner, A_bar_real, A_bar_imag, B_bar, C, D, h0_real, h0_imag):
+    """Main entry point for complex SSM scan."""
+    return ComplexSSMScan.apply(x_inner, A_bar_real, A_bar_imag, B_bar, C, D, h0_real, h0_imag)

model/triton_ssm.py

@@ -0,0 +1,254 @@
+"""
+Pure Triton Fused Selective State Space Model (SSM) kernel.
+Eliminates mamba_ssm / selective_scan_cuda dependencies entirely.
+Compiles natively on NVIDIA B300 (sm_100), H100 (sm_90), A100 (sm_80).
+
+Fuses:
+1. Discretization: dA = exp(dt * A)
+2. State input update: u = dt * x * B
+3. Recurrent state accumulation: h_t = dA * h_{t-1} + u_t
+4. State output contraction: y_t = sum(h_t * C_t)
+5. Skip connection: y_t + x_t * D
+6. Output gating: out_t = (y_t + x_t * D) * silu(z_t)
+"""
+
+from __future__ import annotations
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def ssm_fwd_kernel(
+    # Pointers
+    x_ptr, dt_ptr, A_ptr, B_ptr, C_ptr, D_ptr, z_ptr, out_ptr, h_all_ptr,
+    # Strides
+    stride_xb, stride_xd, stride_xt,
+    stride_dtb, stride_dtd, stride_dtt,
+    stride_Ad, stride_An,
+    stride_Bb, stride_Bn, stride_Bt,
+    stride_Cb, stride_Cn, stride_Ct,
+    stride_outb, stride_outd, stride_outt,
+    stride_hb, stride_hd, stride_ht, stride_hn,
+    # Dimensions
+    T: tl.constexpr,
+    N: tl.constexpr,
+    N_BLOCK: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_d = tl.program_id(1)
+
+    offs_n = tl.arange(0, N_BLOCK)
+    mask_n = offs_n < N
+
+    # Load continuous-time decay A and skip coefficient D
+    A = tl.load(A_ptr + pid_d * stride_Ad + offs_n * stride_An, mask=mask_n, other=0.0)
+    D = tl.load(D_ptr + pid_d)
+
+    # State h accumulator in registers (FP32 precision)
+    h = tl.zeros([N_BLOCK], dtype=tl.float32)
+
+    for t in range(0, T):
+        dt_val = tl.load(dt_ptr + pid_b * stride_dtb + pid_d * stride_dtd + t * stride_dtt).to(tl.float32)
+        x_val = tl.load(x_ptr + pid_b * stride_xb + pid_d * stride_xd + t * stride_xt).to(tl.float32)
+        z_val = tl.load(z_ptr + pid_b * stride_xb + pid_d * stride_xd + t * stride_xt).to(tl.float32)
+
+        B_vec = tl.load(B_ptr + pid_b * stride_Bb + offs_n * stride_Bn + t * stride_Bt, mask=mask_n, other=0.0).to(tl.float32)
+        C_vec = tl.load(C_ptr + pid_b * stride_Cb + offs_n * stride_Cn + t * stride_Ct, mask=mask_n, other=0.0).to(tl.float32)
+
+        # Discretize: dA = exp(dt * A)
+        dA = tl.exp(dt_val * A)
+
+        # u = dt * x * B
+        u = (dt_val * x_val) * B_vec
+
+        # State transition: h = dA * h + u
+        h = dA * h + u
+
+        # Save state snapshot for exact backward gradient flow
+        tl.store(h_all_ptr + pid_b * stride_hb + pid_d * stride_hd + t * stride_ht + offs_n * stride_hn, h, mask=mask_n)
+
+        # Output contraction: y = sum(h * C)
+        y = tl.sum(tl.where(mask_n, h * C_vec, 0.0), axis=0)
+
+        # Skip connection + SiLU gate: out = (y + x * D) * silu(z)
+        sig_z = 1.0 / (1.0 + tl.exp(-z_val))
+        silu_z = z_val * sig_z
+        out_val = (y + x_val * D) * silu_z
+
+        tl.store(out_ptr + pid_b * stride_outb + pid_d * stride_outd + t * stride_outt, out_val)
+
+
+@triton.jit
+def ssm_bwd_fast_kernel(
+    dout_ptr, x_ptr, dt_ptr, A_ptr, B_ptr, C_ptr, D_ptr, z_ptr, h_all_ptr,
+    dx_ptr, ddt_ptr, dA_ptr, dD_ptr, dz_ptr, dy_ptr, dh_lambda_ptr,
+    stride_xb, stride_xd, stride_xt,
+    stride_dtb, stride_dtd, stride_dtt,
+    stride_Ad, stride_An,
+    stride_Bb, stride_Bn, stride_Bt,
+    stride_Cb, stride_Cn, stride_Ct,
+    stride_doutb, stride_doutd, stride_doutt,
+    stride_hb, stride_hd, stride_ht, stride_hn,
+    T: tl.constexpr, N: tl.constexpr, N_BLOCK: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_d = tl.program_id(1)
+
+    offs_n = tl.arange(0, N_BLOCK)
+    mask_n = offs_n < N
+
+    A = tl.load(A_ptr + pid_d * stride_Ad + offs_n * stride_An, mask=mask_n, other=0.0)
+    D = tl.load(D_ptr + pid_d)
+
+    lambda_h = tl.zeros([N_BLOCK], dtype=tl.float32)
+    accum_dD = 0.0
+    accum_dA = tl.zeros([N_BLOCK], dtype=tl.float32)
+
+    for t in range(T - 1, -1, -1):
+        dout_val = tl.load(dout_ptr + pid_b * stride_doutb + pid_d * stride_doutd + t * stride_doutt).to(tl.float32)
+        x_val = tl.load(x_ptr + pid_b * stride_xb + pid_d * stride_xd + t * stride_xt).to(tl.float32)
+        z_val = tl.load(z_ptr + pid_b * stride_xb + pid_d * stride_xd + t * stride_xt).to(tl.float32)
+        dt_val = tl.load(dt_ptr + pid_b * stride_dtb + pid_d * stride_dtd + t * stride_dtt).to(tl.float32)
+
+        B_vec = tl.load(B_ptr + pid_b * stride_Bb + offs_n * stride_Bn + t * stride_Bt, mask=mask_n, other=0.0).to(tl.float32)
+        C_vec = tl.load(C_ptr + pid_b * stride_Cb + offs_n * stride_Cn + t * stride_Ct, mask=mask_n, other=0.0).to(tl.float32)
+        h_t = tl.load(h_all_ptr + pid_b * stride_hb + pid_d * stride_hd + t * stride_ht + offs_n * stride_hn, mask=mask_n, other=0.0).to(tl.float32)
+
+        if t > 0:
+            h_prev = tl.load(h_all_ptr + pid_b * stride_hb + pid_d * stride_hd + (t - 1) * stride_ht + offs_n * stride_hn, mask=mask_n, other=0.0).to(tl.float32)
+        else:
+            h_prev = tl.zeros([N_BLOCK], dtype=tl.float32)
+
+        sig_z = 1.0 / (1.0 + tl.exp(-z_val))
+        silu_z = z_val * sig_z
+        dsilu_z = sig_z * (1.0 + z_val * (1.0 - sig_z))
+
+        y = tl.sum(tl.where(mask_n, h_t * C_vec, 0.0), axis=0)
+        core = y + x_val * D
+
+        dz_val = dout_val * core * dsilu_z
+        dy_val = dout_val * silu_z
+
+        tl.store(dz_ptr + pid_b * stride_xb + pid_d * stride_xd + t * stride_xt, dz_val)
+        tl.store(dy_ptr + pid_b * stride_xb + pid_d * stride_xd + t * stride_xt, dy_val)
+
+        accum_dD += dy_val * x_val
+
+        dh_t = dy_val * C_vec + lambda_h
+        tl.store(dh_lambda_ptr + pid_b * stride_hb + pid_d * stride_hd + t * stride_ht + offs_n * stride_hn, dh_t, mask=mask_n)
+
+        du_x_dt = tl.sum(tl.where(mask_n, dh_t * B_vec, 0.0), axis=0)
+        dx_val = dy_val * D + du_x_dt * dt_val
+        tl.store(dx_ptr + pid_b * stride_xb + pid_d * stride_xd + t * stride_xt, dx_val)
+
+        dA_step = tl.exp(dt_val * A)
+        term_A = tl.sum(tl.where(mask_n, dh_t * h_prev * A * dA_step, 0.0), axis=0)
+        ddt_val = du_x_dt * x_val + term_A
+        tl.store(ddt_ptr + pid_b * stride_dtb + pid_d * stride_dtd + t * stride_dtt, ddt_val)
+
+        accum_dA += tl.where(mask_n, dh_t * h_prev * dt_val * dA_step, 0.0)
+        lambda_h = dh_t * dA_step
+
+    tl.atomic_add(dA_ptr + pid_d * stride_Ad + offs_n * stride_An, accum_dA, mask=mask_n)
+    tl.atomic_add(dD_ptr + pid_d, accum_dD)
+
+
+class TritonSSMFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, dt, A, B_ssm, C_ssm, D_skip, z):
+        B_sz, d_in, T_sz = x.shape
+        d_st = A.shape[1]
+        out = torch.empty_like(x)
+        h_all = torch.empty(B_sz, d_in, T_sz, d_st, device=x.device, dtype=torch.bfloat16)
+
+        grid = (B_sz, d_in)
+        ssm_fwd_kernel[grid](
+            x, dt, A, B_ssm, C_ssm, D_skip, z, out, h_all,
+            x.stride(0), x.stride(1), x.stride(2),
+            dt.stride(0), dt.stride(1), dt.stride(2),
+            A.stride(0), A.stride(1),
+            B_ssm.stride(0), B_ssm.stride(1), B_ssm.stride(2),
+            C_ssm.stride(0), C_ssm.stride(1), C_ssm.stride(2),
+            out.stride(0), out.stride(1), out.stride(2),
+            h_all.stride(0), h_all.stride(1), h_all.stride(2), h_all.stride(3),
+            T=T_sz, N=d_st, N_BLOCK=64,
+        )
+        ctx.save_for_backward(x, dt, A, B_ssm, C_ssm, D_skip, z, h_all)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        x, dt, A, B_ssm, C_ssm, D_skip, z, h_all = ctx.saved_tensors
+        B_sz, d_in, T_sz = x.shape
+        d_st = A.shape[1]
+
+        dx = torch.empty_like(x)
+        ddt = torch.empty_like(dt)
+        dz = torch.empty_like(z)
+        dy = torch.empty_like(x)
+        dh_lambda = torch.empty_like(h_all)
+        dA = torch.zeros_like(A, dtype=torch.float32)
+        dD = torch.zeros_like(D_skip, dtype=torch.float32)
+
+        grid = (B_sz, d_in)
+        ssm_bwd_fast_kernel[grid](
+            dout, x, dt, A, B_ssm, C_ssm, D_skip, z, h_all,
+            dx, ddt, dA, dD, dz, dy, dh_lambda,
+            x.stride(0), x.stride(1), x.stride(2),
+            dt.stride(0), dt.stride(1), dt.stride(2),
+            A.stride(0), A.stride(1),
+            B_ssm.stride(0), B_ssm.stride(1), B_ssm.stride(2),
+            C_ssm.stride(0), C_ssm.stride(1), C_ssm.stride(2),
+            dout.stride(0), dout.stride(1), dout.stride(2),
+            h_all.stride(0), h_all.stride(1), h_all.stride(2), h_all.stride(3),
+            T=T_sz, N=d_st, N_BLOCK=64,
+        )
+
+        # Zero atomic collisions: fast vectorized contractions
+        dC = torch.einsum('bdt, bdtn -> bnt', dy.float(), h_all.float()).to(C_ssm.dtype)
+        dt_x = (dt * x).float()
+        dB = torch.einsum('bdt, bdtn -> bnt', dt_x, dh_lambda.float()).to(B_ssm.dtype)
+
+        return (
+            dx,
+            ddt,
+            dA.to(A.dtype),
+            dB,
+            dC,
+            dD.to(D_skip.dtype),
+            dz,
+        )
+
+
+def fused_triton_ssm(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B_ssm: torch.Tensor,
+    C_ssm: torch.Tensor,
+    D_skip: torch.Tensor,
+    z: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Fused selective scan forward and backward in pure Triton.
+    Inputs:
+      x:      (B, d_inner, T)
+      dt:     (B, d_inner, T)
+      A:      (d_inner, d_state)
+      B_ssm:  (B, d_state, T)
+      C_ssm:  (B, d_state, T)
+      D_skip: (d_inner,)
+      z:      (B, d_inner, T)
+    Returns:
+      out:    (B, d_inner, T)
+    """
+    return TritonSSMFunction.apply(
+        x.contiguous(),
+        dt.contiguous(),
+        A.contiguous(),
+        B_ssm.contiguous(),
+        C_ssm.contiguous(),
+        D_skip.contiguous(),
+        z.contiguous(),
+    )

model_bf16.pt

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

model_config.json

@@ -0,0 +1,76 @@
+{
+  "vocab_size": 206464,
+  "d_model": 1024,
+  "n_layers": 18,
+  "max_seq_len": 4096,
+  "mha_layer_indices": [
+    5,
+    11,
+    17
+  ],
+  "dense_ffn_layer_indices": [
+    1,
+    3,
+    5,
+    7,
+    9,
+    11,
+    13,
+    15,
+    17
+  ],
+  "hse_ffn_layer_indices": [
+    0,
+    2,
+    4,
+    6,
+    8,
+    10,
+    12,
+    14,
+    16
+  ],
+  "arg_d_inner": 1024,
+  "arg_d_state": 48,
+  "arg_d_rank": 48,
+  "arg_conv_kernel": 4,
+  "arg_local_window": 64,
+  "arg_local_n_heads": 8,
+  "arg_local_n_kv_heads": 2,
+  "arg_head_dim": 64,
+  "arg_scalar_gate": true,
+  "mha_n_heads": 16,
+  "mha_n_kv_heads": 4,
+  "mha_head_dim": 64,
+  "rope_theta": 500000.0,
+  "rope_scaling_type": null,
+  "rope_scaling_factor": 1.0,
+  "dense_d_ff": 2560,
+  "hse_n_domains": 4,
+  "hse_n_experts_per_domain": 6,
+  "hse_top_k": 2,
+  "hse_d_ff_expert": 210,
+  "hse_capacity_factor": 1.25,
+  "hse_expert_dropout": 0.1,
+  "hse_aux_loss_alpha": 0.01,
+  "hse_gumbel_tau_init": 1.0,
+  "hse_gumbel_tau_final": 0.1,
+  "hse_gumbel_tau_warmup_tokens": 10000000000,
+  "norm_type": "dgn",
+  "dgn_n_groups": 16,
+  "norm_eps": 1e-06,
+  "init_std": 0.02,
+  "init_std_residual_scale": true,
+  "dropout": 0.0,
+  "tie_word_embeddings": true,
+  "use_cache": true,
+  "pad_token_id": 0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "use_flash_attention": true,
+  "use_triton_kernels": true,
+  "use_torch_compile": true,
+  "compile_mode": "max-autotune",
+  "use_gradient_checkpointing": true,
+  "gradient_checkpointing_ratio": 1.0
+}

requirements.txt

@@ -0,0 +1,16 @@
+# FORGE-3B dependencies
+torch>=2.3.0
+transformers>=4.41.0
+datasets>=2.19.0
+tokenizers>=0.19.1
+wandb==0.17.0
+deepspeed>=0.14.4
+flash-attn>=2.5.8
+einops==0.8.0
+accelerate==0.30.1
+bitsandbytes==0.43.1
+mamba-ssm>=1.2.0
+causal-conv1d>=1.3.0
+numpy>=1.24.0
+triton>=2.1.0
+xerv-crayon

run_inference.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+import os
+import argparse
+import logging
+import torch
+from pathlib import Path
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format="[%(asctime)s][%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Run inference on FORGE-500M using a trained model checkpoint")
+    parser.add_argument("--model_path", type=str, default=".",
+                        help="Path to the trained model directory containing config, tokenizer, and weights (defaults to current directory)")
+    parser.add_argument("--prompt", type=str, default="The future of artificial intelligence is",
+                        help="Prompt to generate text from")
+    parser.add_argument("--max_new_tokens", type=int, default=128,
+                        help="Max number of tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.7,
+                        help="Sampling temperature for text generation")
+    parser.add_argument("--top_k", type=int, default=50,
+                        help="Top-k sampling parameter")
+    parser.add_argument("--top_p", type=float, default=0.9,
+                        help="Top-p sampling parameter")
+    parser.add_argument("--repetition_penalty", type=float, default=1.1,
+                        help="Repetition penalty parameter")
+    return parser.parse_args()
+
+def main():
+    args = get_args()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Using device: {device}")
+    
+    model_dir = Path(args.model_path)
+    if not model_dir.exists():
+        raise FileNotFoundError(f"Model path does not exist: {model_dir}")
+        
+    # 1. Load Tokenizer
+    from tokenizer.crayon_wrapper import ForgeTokenizer
+    tok_dir = model_dir / "tokenizer"
+    if tok_dir.exists():
+        logger.info(f"Loading tokenizer from {tok_dir}...")
+        tokenizer = ForgeTokenizer.from_pretrained(str(tok_dir))
+    else:
+        logger.info(f"Loading tokenizer from model directory root {model_dir}...")
+        tokenizer = ForgeTokenizer.from_pretrained(str(model_dir))
+        
+    # 2. Load Model Config
+    from config import ForgeModelConfig
+    config_path = model_dir / "model_config.json"
+    if not config_path.exists():
+        config_path = model_dir / "config.json"
+    if not config_path.exists():
+        raise FileNotFoundError(f"Could not find model_config.json or config.json in {model_dir}")
+        
+    logger.info(f"Loading config from {config_path}")
+    model_config = ForgeModelConfig.from_json(str(config_path))
+    model_config.vocab_size = tokenizer.vocab_size
+    
+    # 3. Build Model
+    from model.forge_model import build_forge_500m
+    logger.info("Building FORGE model...")
+    model = build_forge_500m(model_config)
+    
+    # 4. Load Model Weights
+    loaded_weights = False
+    for weight_filename in ("model_bf16.pt", "model.pt", "pytorch_model.bin"):
+        weight_path = model_dir / weight_filename
+        if weight_path.exists():
+            logger.info(f"Loading model weights from {weight_path}...")
+            state_dict = torch.load(str(weight_path), map_location="cpu")
+            # Convert values to bfloat16
+            state_dict = {k: v.to(torch.bfloat16) for k, v in state_dict.items()}
+            model.load_state_dict(state_dict, strict=False)
+            del state_dict
+            torch.cuda.empty_cache()
+            loaded_weights = True
+            break
+            
+    if not loaded_weights:
+        raise FileNotFoundError(
+            f"No model weights found in {model_dir}. "
+            "Expected one of: model_bf16.pt, model.pt, pytorch_model.bin"
+        )
+        
+    model = model.to(device).to(torch.bfloat16)
+    model.eval()
+    
+    # 5. Encode prompt
+    logger.info(f"Prompt: {args.prompt!r}")
+    input_ids = torch.tensor([tokenizer.encode(args.prompt, add_bos=True, add_eos=False)], dtype=torch.long, device=device)
+    
+    # 6. Generate
+    logger.info("Generating tokens...")
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids=input_ids,
+            max_new_tokens=args.max_new_tokens,
+            temperature=args.temperature,
+            top_k=args.top_k,
+            top_p=args.top_p,
+            repetition_penalty=args.repetition_penalty,
+            eos_token_id=tokenizer.special_tokens["<|eos|>"],
+            pad_token_id=tokenizer.special_tokens["<|pad|>"]
+        )
+    
+    # 7. Decode output
+    generated_text = tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=False)
+    logger.info(f"Generated text:\n{generated_text}")
+
+if __name__ == "__main__":
+    main()

tokenizer/crayon_wrapper.py

@@ -0,0 +1,653 @@
+"""
+CRAYON Tokenizer Wrapper
+Wraps CrayonVocab into a full-featured tokenizer compatible with FORGE training.
+Exploits CRAYON's 24M+ tok/sec CPU throughput via parallel multi-process pipeline.
+"""
+
+from __future__ import annotations
+
+import os
+import gc
+import math
+import time
+import logging
+import threading
+import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
+from typing import List, Union, Optional, Dict, Any, Iterator, Tuple
+from pathlib import Path
+import queue
+
+# Fix Windows console UTF-8 printing issues for external packages (e.g., Crayon print statements with emoji)
+import sys
+import io
+if sys.platform.startswith("win"):
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+        sys.stderr.reconfigure(encoding='utf-8')
+    except AttributeError:
+        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
+        sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
+
+import torch
+import numpy as np
+
+from crayon import CrayonVocab
+
+logger = logging.getLogger(__name__)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# THREAD-LOCAL TOKENIZER POOL
+# Each worker thread gets its own CrayonVocab instance to avoid any GIL
+# contention on the underlying C++ DAT engine.
+# ─────────────────────────────────────────────────────────────────────────────
+
+_thread_local = threading.local()
+_init_lock = threading.Lock()
+
+def _get_thread_local_tokenizer(profile: str = "standard") -> CrayonVocab:
+    """Return a thread-local CrayonVocab instance, creating one if needed."""
+    if not hasattr(_thread_local, "tokenizer"):
+        with _init_lock:
+            # CPU is 20× faster than CUDA for CRAYON (23.84M vs 1.22M tok/s)
+            _thread_local.tokenizer = CrayonVocab(device="cpu")
+            _thread_local.tokenizer.load_profile(profile)
+            logger.debug(f"Thread {threading.current_thread().name}: "
+                         f"CrayonVocab loaded with profile='{profile}'")
+    return _thread_local.tokenizer
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# SPECIAL TOKEN REGISTRY
+# We manage special tokens separately since CRAYON's vocabulary may not have
+# them natively — we append them to the vocab ID space.
+# ─────────────────────────────────────────────────────────────────────────────
+
+SPECIAL_TOKENS = {
+    "<|pad|>":       0,
+    "<|bos|>":       1,
+    "<|eos|>":       2,
+    "<|sep|>":       3,
+    "<|sys|>":       4,
+    "<|/sys|>":      5,
+    "<|usr|>":       6,
+    "<|/usr|>":      7,
+    "<|asst|>":      8,
+    "<|/asst|>":     9,
+    "<|tool|>":      10,
+    "<|tool_resp|>": 11,
+    "<|unk|>":       12,
+}
+
+# For the standard profile: actual vocab starts at ID 13
+# Effective vocab_size = 206,373 (crayon) + 13 (specials) = 206,386 → pad to 206,464
+SPECIAL_TOKEN_OFFSET = len(SPECIAL_TOKENS)   # 13
+CRAYON_STANDARD_VOCAB = 206_373
+CRAYON_LITE_VOCAB = 50_000
+
+# Reverse mapping for decode
+ID_TO_SPECIAL = {v: k for k, v in SPECIAL_TOKENS.items()}
+
+
+class ForgeTokenizer:
+    """
+    Production-grade FORGE tokenizer built on top of CRAYON.
+    
+    Features:
+    - Thread-safe, zero-copy multi-threaded tokenization
+    - Batched encoding with automatic chunking for throughput
+    - Full special token support with chat template
+    - Streaming decode
+    - Vocabulary size adapter between CRAYON profile and model vocab
+    - Offline preprocessing at 24M+ tokens/sec on CPU
+    
+    Usage:
+        tokenizer = ForgeTokenizer(profile="standard")
+        ids = tokenizer.encode("Hello world")
+        text = tokenizer.decode(ids)
+        
+        # Chat format
+        ids = tokenizer.encode_chat([
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "What is 2+2?"},
+        ])
+    """
+    
+    def __init__(
+        self,
+        profile: str = "standard",
+        device: str = "cpu",
+        n_workers: int = None,
+        max_length: Optional[int] = None,
+        add_bos: bool = True,
+        add_eos: bool = True,
+    ):
+        self.profile = profile
+        self.device_str = device   # "cpu" recommended for CRAYON throughput
+        self.n_workers = n_workers or max(1, mp.cpu_count() - 2)
+        self.max_length = max_length
+        self.add_bos = add_bos
+        self.add_eos = add_eos
+        
+        # Primary tokenizer (used for single-threaded ops)
+        self._tokenizer = CrayonVocab(device=device)
+        self._tokenizer.load_profile(profile)
+        
+        # Special tokens
+        self.special_tokens = SPECIAL_TOKENS.copy()
+        self.id_to_special = ID_TO_SPECIAL.copy()
+        self.special_token_offset = SPECIAL_TOKEN_OFFSET
+        
+        # Derive vocab size
+        if profile == "standard":
+            self._crayon_vocab_size = CRAYON_STANDARD_VOCAB
+        else:
+            self._crayon_vocab_size = CRAYON_LITE_VOCAB
+        
+        # Effective vocab size = specials + crayon tokens
+        # Padded to nearest multiple of 128 for efficient matmuls
+        raw_vocab = self._crayon_vocab_size + SPECIAL_TOKEN_OFFSET
+        self.vocab_size = math.ceil(raw_vocab / 128) * 128
+        
+        # Lock for thread safety on single-instance operations
+        self._lock = threading.Lock()
+        
+        # Thread pool for parallel tokenization
+        self._thread_pool = ThreadPoolExecutor(
+            max_workers=self.n_workers,
+            thread_name_prefix="forge_tok"
+        )
+        
+        logger.info(
+            f"ForgeTokenizer initialized: profile={profile}, "
+            f"vocab_size={self.vocab_size} (crayon={self._crayon_vocab_size} "
+            f"+ specials={SPECIAL_TOKEN_OFFSET} + padding={self.vocab_size - raw_vocab}), "
+            f"n_workers={self.n_workers}, throughput ~24M tok/s (CPU)"
+        )
+    
+    # ─────────────────────────────────────────────────────────────────────────
+    # TOKEN ID CONVERSION HELPERS
+    # CRAYON returns IDs in [0, crayon_vocab_size).
+    # We shift them by SPECIAL_TOKEN_OFFSET to [SPECIAL_TOKEN_OFFSET, vocab_size).
+    # Special tokens occupy IDs [0, SPECIAL_TOKEN_OFFSET).
+    # ─────────────────────────────────────────────────────────────────────────
+    
+    def _shift_crayon_ids(self, ids: List[int]) -> List[int]:
+        """Shift CRAYON token IDs to make room for special tokens at the front."""
+        return [i + self.special_token_offset for i in ids]
+    
+    def _unshift_crayon_ids(self, ids: List[int]) -> List[int]:
+        """Inverse of _shift_crayon_ids for decode path."""
+        out = []
+        for i in ids:
+            if i < self.special_token_offset:
+                out.append(None)   # special token — handled separately
+            else:
+                out.append(i - self.special_token_offset)
+        return out
+    
+    # ─────────────────────────────────────────────────────────────────────────
+    # CORE ENCODE / DECODE
+    # ─────────────────────────────────────────────────────────────────────────
+    
+    def encode(
+        self,
+        text: str,
+        add_bos: Optional[bool] = None,
+        add_eos: Optional[bool] = None,
+        truncate: bool = True,
+    ) -> List[int]:
+        """
+        Encode a single string to token IDs.
+        Thread-safe via CRAYON's internal lock or thread-local tokenizers.
+        """
+        add_bos_ = self.add_bos if add_bos is None else add_bos
+        add_eos_ = self.add_eos if add_eos is None else add_eos
+        
+        # CRAYON tokenization
+        with self._lock:
+            raw_ids = self._tokenizer.tokenize(text)
+        
+        # Handle both list and tensor returns from CRAYON
+        if isinstance(raw_ids, torch.Tensor):
+            raw_ids = raw_ids.tolist()
+        
+        ids = self._shift_crayon_ids(raw_ids)
+        
+        # Prepend/append special tokens
+        if add_bos_:
+            ids = [self.special_tokens["<|bos|>"]] + ids
+        if add_eos_:
+            ids = ids + [self.special_tokens["<|eos|>"]]
+        
+        # Truncate if needed
+        if truncate and self.max_length is not None:
+            ids = ids[:self.max_length]
+        
+        return ids
+    
+    def encode_batch(
+        self,
+        texts: List[str],
+        add_bos: Optional[bool] = None,
+        add_eos: Optional[bool] = None,
+        pad: bool = True,
+        pad_to_multiple_of: int = 8,
+        return_tensors: str = "pt",
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Encode a batch of texts in parallel using thread pool.
+        Leverages CRAYON's 24M tok/s CPU throughput across all threads.
+        
+        Returns dict with 'input_ids' and 'attention_mask'.
+        """
+        add_bos_ = self.add_bos if add_bos is None else add_bos
+        add_eos_ = self.add_eos if add_eos is None else add_eos
+        
+        # ── Parallel tokenization ──────────────────────────────────────────
+        def _tokenize_one(text: str) -> List[int]:
+            tok = _get_thread_local_tokenizer(self.profile)
+            raw = tok.tokenize(text)
+            if isinstance(raw, torch.Tensor):
+                raw = raw.tolist()
+            ids = self._shift_crayon_ids(raw)
+            if add_bos_:
+                ids = [self.special_tokens["<|bos|>"]] + ids
+            if add_eos_:
+                ids = ids + [self.special_tokens["<|eos|>"]]
+            if self.max_length is not None:
+                ids = ids[:self.max_length]
+            return ids
+        
+        if self.n_workers <= 1:
+            # Run sequentially on the main thread using the primary tokenizer instance
+            # to completely avoid any C++ multi-threading deadlocks.
+            all_ids = []
+            for t in texts:
+                raw = self._tokenizer.tokenize(t)
+                if isinstance(raw, torch.Tensor):
+                    raw = raw.tolist()
+                ids = self._shift_crayon_ids(raw)
+                if add_bos_:
+                    ids = [self.special_tokens["<|bos|>"]] + ids
+                if add_eos_:
+                    ids = ids + [self.special_tokens["<|eos|>"]]
+                if self.max_length is not None:
+                    ids = ids[:self.max_length]
+                all_ids.append(ids)
+        else:
+            # Submit all texts to thread pool
+            futures = [self._thread_pool.submit(_tokenize_one, t) for t in texts]
+            all_ids = [f.result() for f in futures]
+        
+        if not pad or return_tensors is None:
+            return {"input_ids": all_ids}
+        
+        # ── Padding ────────────────────────────────────────────────────────
+        max_len = max(len(ids) for ids in all_ids)
+        # Pad to multiple for efficiency
+        if pad_to_multiple_of > 1:
+            max_len = math.ceil(max_len / pad_to_multiple_of) * pad_to_multiple_of
+        
+        pad_id = self.special_tokens["<|pad|>"]
+        padded_ids = []
+        attention_masks = []
+        
+        for ids in all_ids:
+            n_pad = max_len - len(ids)
+            padded_ids.append(ids + [pad_id] * n_pad)
+            attention_masks.append([1] * len(ids) + [0] * n_pad)
+        
+        if return_tensors == "pt":
+            return {
+                "input_ids": torch.tensor(padded_ids, dtype=torch.long),
+                "attention_mask": torch.tensor(attention_masks, dtype=torch.long),
+            }
+        return {"input_ids": padded_ids, "attention_mask": attention_masks}
+    
+    def decode(
+        self,
+        ids: Union[List[int], torch.Tensor],
+        skip_special_tokens: bool = True,
+    ) -> str:
+        """Decode token IDs back to text."""
+        if isinstance(ids, torch.Tensor):
+            ids = ids.tolist()
+        
+        # Separate special tokens from CRAYON tokens
+        crayon_ids = []
+        special_strs = []
+        
+        for tok_id in ids:
+            if tok_id < self.special_token_offset:
+                if not skip_special_tokens and tok_id in self.id_to_special:
+                    special_strs.append(self.id_to_special[tok_id])
+            else:
+                crayon_ids.append(tok_id - self.special_token_offset)
+        
+        # CRAYON decode
+        if crayon_ids:
+            with self._lock:
+                text = self._tokenizer.decode(crayon_ids)
+        else:
+            text = ""
+        
+        # Re-insert special tokens if not skipping
+        if special_strs and not skip_special_tokens:
+            text = " ".join(special_strs) + " " + text
+        
+        return text
+    
+    def decode_batch(
+        self,
+        batch_ids: Union[List[List[int]], torch.Tensor],
+        skip_special_tokens: bool = True,
+    ) -> List[str]:
+        """Decode a batch in parallel."""
+        if isinstance(batch_ids, torch.Tensor):
+            batch_ids = batch_ids.tolist()
+        
+        if self.n_workers <= 1:
+            return [self.decode(ids, skip_special_tokens) for ids in batch_ids]
+            
+        futures = [
+            self._thread_pool.submit(self.decode, ids, skip_special_tokens)
+            for ids in batch_ids
+        ]
+        return [f.result() for f in futures]
+    
+    # ─────────────────────────────────────────────────────────────────────────
+    # CHAT TEMPLATE
+    # ──────────────────────────────────────────────────────────────────���──────
+    
+    def encode_chat(
+        self,
+        messages: List[Dict[str, str]],
+        add_generation_prompt: bool = False,
+        return_as_tensor: bool = True,
+    ) -> Union[List[int], torch.Tensor]:
+        """
+        Encode a conversation using FORGE chat template.
+        
+        Template:
+            <|bos|><|sys|>{system}<|/sys|>
+            <|usr|>{user}<|/usr|>
+            <|asst|>{assistant}<|/asst|>
+            ...
+        
+        Args:
+            messages: List of {"role": "system|user|assistant", "content": str}
+            add_generation_prompt: If True, append <|asst|> at end (inference)
+        """
+        token_ids = [self.special_tokens["<|bos|>"]]
+        
+        role_to_open = {
+            "system":    "<|sys|>",
+            "user":      "<|usr|>",
+            "assistant": "<|asst|>",
+        }
+        role_to_close = {
+            "system":    "<|/sys|>",
+            "user":      "<|/usr|>",
+            "assistant": "<|/asst|>",
+        }
+        
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+            
+            # Open tag
+            open_tag = role_to_open.get(role, "<|usr|>")
+            close_tag = role_to_close.get(role, "<|/usr|>")
+            
+            token_ids.append(self.special_tokens.get(open_tag, self.special_tokens["<|unk|>"]))
+            
+            # Encode content (no bos/eos, no truncation)
+            content_ids = self.encode(content, add_bos=False, add_eos=False, truncate=False)
+            token_ids.extend(content_ids)
+            
+            token_ids.append(self.special_tokens.get(close_tag, self.special_tokens["<|unk|>"]))
+        
+        if add_generation_prompt:
+            token_ids.append(self.special_tokens["<|asst|>"])
+        
+        if self.max_length is not None:
+            token_ids = token_ids[:self.max_length]
+        
+        if return_as_tensor:
+            return torch.tensor(token_ids, dtype=torch.long)
+        return token_ids
+    
+    def encode_chat_with_labels(
+        self,
+        messages: List[Dict[str, str]],
+        loss_on_prompt: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Encode a conversation and produce labels for SFT.
+        Labels are -100 (ignored) for non-assistant turns unless loss_on_prompt=True.
+        
+        Returns dict with 'input_ids' and 'labels'.
+        """
+        input_ids = []
+        labels = []
+        
+        input_ids.append(self.special_tokens["<|bos|>"])
+        labels.append(-100 if not loss_on_prompt else self.special_tokens["<|bos|>"])
+        
+        role_to_open = {
+            "system":    "<|sys|>",
+            "user":      "<|usr|>",
+            "assistant": "<|asst|>",
+        }
+        role_to_close = {
+            "system":    "<|/sys|>",
+            "user":      "<|/usr|>",
+            "assistant": "<|/asst|>",
+        }
+        
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+            is_assistant = (role == "assistant")
+            
+            open_tag_id = self.special_tokens.get(role_to_open.get(role, "<|usr|>"), 
+                                                    self.special_tokens["<|unk|>"])
+            close_tag_id = self.special_tokens.get(role_to_close.get(role, "<|/usr|>"),
+                                                     self.special_tokens["<|unk|>"])
+            
+            content_ids = self.encode(content, add_bos=False, add_eos=False, truncate=False)
+            
+            turn_ids = [open_tag_id] + content_ids + [close_tag_id]
+            input_ids.extend(turn_ids)
+            
+            if is_assistant or loss_on_prompt:
+                # Compute loss on assistant content + close tag, NOT on open tag
+                turn_labels = [-100] + content_ids + [close_tag_id]
+                labels.extend(turn_labels)
+            else:
+                labels.extend([-100] * len(turn_ids))
+        
+        if self.max_length is not None:
+            input_ids = input_ids[:self.max_length]
+            labels = labels[:self.max_length]
+        
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+        }
+    
+    # ─────────────────────────────────────────────────────────────────────────
+    # HIGH-THROUGHPUT OFFLINE PREPROCESSING
+    # ─────────────────────────────────────────────────────────────────────────
+    
+    def tokenize_file_streaming(
+        self,
+        input_path: str,
+        output_path: str,
+        chunk_size: int = 100_000,       # documents per chunk
+        seq_len: int = 2048,
+        append_eos: bool = True,
+        n_processes: int = None,
+    ) -> int:
+        """
+        High-throughput offline tokenization of a large text file.
+        
+        Architecture:
+        - Producer process reads text lines
+        - N worker processes each run their own CrayonVocab instance
+        - Consumer aggregates and packs sequences
+        
+        Throughput: ~5-10GB/min depending on hardware.
+        Returns: total number of tokens written.
+        """
+        n_processes = n_processes or max(1, mp.cpu_count() - 2)
+        total_tokens = 0
+        
+        logger.info(f"Offline tokenization: {input_path} → {output_path}, "
+                    f"n_processes={n_processes}, seq_len={seq_len}")
+        
+        # We use process-based parallelism to bypass GIL fully
+        with ProcessPoolExecutor(max_workers=n_processes) as pool:
+            with open(input_path, "r", encoding="utf-8", errors="replace") as fin:
+                chunk_buffer = []
+                futures = []
+                
+                for line in fin:
+                    chunk_buffer.append(line.strip())
+                    if len(chunk_buffer) >= chunk_size:
+                        futures.append(
+                            pool.submit(
+                                _tokenize_chunk_worker,
+                                chunk_buffer.copy(),
+                                self.profile,
+                                self.special_token_offset,
+                                SPECIAL_TOKENS["<|eos|>"] if append_eos else None,
+                            )
+                        )
+                        chunk_buffer.clear()
+                
+                # Submit remaining
+                if chunk_buffer:
+                    futures.append(
+                        pool.submit(
+                            _tokenize_chunk_worker,
+                            chunk_buffer,
+                            self.profile,
+                            self.special_token_offset,
+                            SPECIAL_TOKENS["<|eos|>"] if append_eos else None,
+                        )
+                    )
+                
+                # Collect and pack
+                token_stream = []
+                packed_sequences = []
+                
+                for future in as_completed(futures):
+                    chunk_tokens = future.result()  # List[int]
+                    token_stream.extend(chunk_tokens)
+                    
+                    # Pack into fixed-length sequences
+                    while len(token_stream) >= seq_len:
+                        packed_sequences.append(token_stream[:seq_len])
+                        token_stream = token_stream[seq_len:]
+                        total_tokens += seq_len
+                    
+                    # Flush to disk in chunks
+                    if len(packed_sequences) >= 10_000:
+                        _flush_to_disk(packed_sequences, output_path, total_tokens)
+                        packed_sequences.clear()
+                
+                # Flush remaining
+                if packed_sequences:
+                    _flush_to_disk(packed_sequences, output_path, total_tokens)
+        
+        logger.info(f"Tokenization complete: {total_tokens:,} tokens written to {output_path}")
+        return total_tokens
+    
+    def __repr__(self) -> str:
+        return (f"ForgeTokenizer(profile={self.profile!r}, "
+                f"vocab_size={self.vocab_size:,}, "
+                f"n_workers={self.n_workers})")
+    
+    def save_pretrained(self, path: str):
+        """Save tokenizer config for later loading."""
+        import json
+        os.makedirs(path, exist_ok=True)
+        config = {
+            "profile": self.profile,
+            "device": self.device_str,
+            "vocab_size": self.vocab_size,
+            "special_tokens": self.special_tokens,
+            "special_token_offset": self.special_token_offset,
+            "add_bos": self.add_bos,
+            "add_eos": self.add_eos,
+            "max_length": self.max_length,
+        }
+        with open(os.path.join(path, "tokenizer_config.json"), "w") as f:
+            json.dump(config, f, indent=2)
+        logger.info(f"Tokenizer config saved to {path}")
+    
+    @classmethod
+    def from_pretrained(cls, path: str) -> "ForgeTokenizer":
+        import json
+        with open(os.path.join(path, "tokenizer_config.json")) as f:
+            config = json.load(f)
+        return cls(
+            profile=config["profile"],
+            device=config.get("device", "cpu"),
+            max_length=config.get("max_length"),
+            add_bos=config.get("add_bos", True),
+            add_eos=config.get("add_eos", True),
+        )
+    
+    def __del__(self):
+        """Graceful cleanup of thread pool."""
+        try:
+            self._thread_pool.shutdown(wait=False)
+        except Exception:
+            pass
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# WORKER FUNCTIONS (must be module-level for pickling with ProcessPoolExecutor)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _tokenize_chunk_worker(
+    texts: List[str],
+    profile: str,
+    offset: int,
+    eos_id: Optional[int],
+) -> List[int]:
+    """
+    Worker function for parallel tokenization.
+    Runs in a separate process with its own CrayonVocab instance.
+    """
+    tok = CrayonVocab(device="cpu")
+    tok.load_profile(profile)
+    
+    all_ids = []
+    for text in texts:
+        if not text:
+            continue
+        raw = tok.tokenize(text)
+        if isinstance(raw, torch.Tensor):
+            raw = raw.tolist()
+        ids = [i + offset for i in raw]
+        if eos_id is not None:
+            ids.append(eos_id)
+        all_ids.extend(ids)
+    
+    return all_ids
+
+
+def _flush_to_disk(sequences: List[List[int]], output_path: str, offset: int):
+    """Append packed sequences to a memory-mapped numpy file."""
+    arr = np.array(sequences, dtype=np.uint32)
+    mode = "ab"  # append binary
+    with open(output_path + ".npy", mode) as f:
+        np.save(f, arr)
+    # Also write offset index for random access
+    with open(output_path + ".idx", "ab") as f:
+        np.save(f, np.array([offset], dtype=np.int64))

tokenizer/tokenizer_config.json

@@ -0,0 +1,24 @@
+{
+  "profile": "standard",
+  "device": "cpu",
+  "vocab_size": 206464,
+  "special_tokens": {
+    "<|pad|>": 0,
+    "<|bos|>": 1,
+    "<|eos|>": 2,
+    "<|sep|>": 3,
+    "<|sys|>": 4,
+    "<|/sys|>": 5,
+    "<|usr|>": 6,
+    "<|/usr|>": 7,
+    "<|asst|>": 8,
+    "<|/asst|>": 9,
+    "<|tool|>": 10,
+    "<|tool_resp|>": 11,
+    "<|unk|>": 12
+  },
+  "special_token_offset": 13,
+  "add_bos": true,
+  "add_eos": true,
+  "max_length": 4096
+}

tokenizer_config.json

@@ -0,0 +1,24 @@
+{
+  "profile": "standard",
+  "device": "cpu",
+  "vocab_size": 206464,
+  "special_tokens": {
+    "<|pad|>": 0,
+    "<|bos|>": 1,
+    "<|eos|>": 2,
+    "<|sep|>": 3,
+    "<|sys|>": 4,
+    "<|/sys|>": 5,
+    "<|usr|>": 6,
+    "<|/usr|>": 7,
+    "<|asst|>": 8,
+    "<|/asst|>": 9,
+    "<|tool|>": 10,
+    "<|tool_resp|>": 11,
+    "<|unk|>": 12
+  },
+  "special_token_offset": 13,
+  "add_bos": true,
+  "add_eos": true,
+  "max_length": 4096
+}